kfaraz commented on code in PR #17653:
URL: https://github.com/apache/druid/pull/17653#discussion_r1947547224


##########
server/src/main/java/org/apache/druid/metadata/SqlSegmentsMetadataQuery.java:
##########
@@ -264,51 +348,64 @@ public Set<SegmentId> retrieveUsedSegmentIds(
       );
     }
 
-    return connector.inReadOnlyTransaction(
-        (handle, status) -> {
-          final Query<Map<String, Object>> sql = handle
-              .createQuery(StringUtils.format(sb.toString(), 
dbTables.getSegmentsTable()))
-              .setFetchSize(connector.getStreamingFetchSize())
-              .bind("used", true)
-              .bind("dataSource", dataSource);
+    final Query<Map<String, Object>> sql = handle
+        .createQuery(StringUtils.format(sb.toString(), 
dbTables.getSegmentsTable()))
+        .setFetchSize(connector.getStreamingFetchSize())
+        .bind("used", true)
+        .bind("dataSource", dataSource);
 
-          if (compareAsString) {
-            bindIntervalsToQuery(sql, Collections.singletonList(interval));
-          }
+    if (compareAsString) {
+      bindIntervalsToQuery(sql, Collections.singletonList(interval));
+    }
+
+    final Set<SegmentId> segmentIds = new HashSet<>();
+    try (final ResultIterator<String> iterator = 
sql.map(StringMapper.FIRST).iterator()) {
+      while (iterator.hasNext()) {
+        final String id = iterator.next();
+        final SegmentId segmentId = SegmentId.tryParse(dataSource, id);
+        if (segmentId == null) {
+          throw DruidException.defensive(
+              "Failed to parse SegmentId for id[%s] and dataSource[%s].",
+              id, dataSource
+          );
+        }
+        if (IntervalMode.OVERLAPS.apply(interval, segmentId.getInterval())) {
+          segmentIds.add(segmentId);
+        }
+      }
+    }
+    return segmentIds;
 
-          final Set<SegmentId> segmentIds = new HashSet<>();
-          try (final ResultIterator<String> iterator = sql.map((index, r, ctx) 
-> r.getString(1)).iterator()) {
-            while (iterator.hasNext()) {
-              final String id = iterator.next();
-              final SegmentId segmentId = SegmentId.tryParse(dataSource, id);
-              if (segmentId == null) {
-                throw DruidException.defensive(
-                    "Failed to parse SegmentId for id[%s] and dataSource[%s].",
-                    id, dataSource
-                );
-              }
-              if (IntervalMode.OVERLAPS.apply(interval, 
segmentId.getInterval())) {
-                segmentIds.add(segmentId);
-              }
-            }
-          }
-          return segmentIds;
-        });
   }
 
   public List<DataSegmentPlus> retrieveSegmentsById(
       String datasource,
       Set<String> segmentIds
   )
+  {
+    try (CloseableIterator<DataSegmentPlus> iterator
+             = retrieveSegmentsByIdIterator(datasource, segmentIds)) {
+      return ImmutableList.copyOf(iterator);
+    }
+    catch (IOException e) {
+      throw DruidException.defensive(e, "Error while retrieving segments from 
metadata store");
+    }
+  }
+
+  public CloseableIterator<DataSegmentPlus> retrieveSegmentsByIdIterator(
+      String datasource,
+      Set<String> segmentIds
+  )
   {
     final List<List<String>> partitionedSegmentIds
         = Lists.partition(new ArrayList<>(segmentIds), 100);
 
-    final List<DataSegmentPlus> fetchedSegments = new 
ArrayList<>(segmentIds.size());
+    final List<CloseableIterator<DataSegmentPlus>> fetchedSegments
+        = new ArrayList<>(partitionedSegmentIds.size());
     for (List<String> partition : partitionedSegmentIds) {
-      fetchedSegments.addAll(retrieveSegmentBatchById(datasource, partition, 
false));
+      fetchedSegments.add(retrieveSegmentBatchById(datasource, partition, 
false));

Review Comment:
   Thanks for the suggestion! Added a CloseableIterator which keeps only one 
result set open at a time.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to