Modify segment metadata call

apache · Oct 23, 2024 · f1488d8 · f1488d8
1 parent b556b37
commit f1488d8
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 14 deletions.
diff --git a/...ontroller/src/main/java/org/apache/pinot/controller/util/ServerSegmentMetadataReader.java b/...ontroller/src/main/java/org/apache/pinot/controller/util/ServerSegmentMetadataReader.java
@@ -412,6 +412,14 @@ private String generateSegmentMetadataServerURL(String tableNameWithType, String
     return String.format("%s/tables/%s/segments/%s/metadata?%s", endpoint, tableNameWithType, segmentName, paramsStr);
   }
 
+  public String generateTableMetadataServerURL(String tableNameWithType, List<String> columns,
+                                               Set<String> segmentsToInclude, String endpoint) {
+    tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
+    String paramsStr = generateColumnsParam(columns)
+            + generateSegmentsParam(new ArrayList<>(segmentsToInclude));
+    return String.format("%s/tables/%s/metadata?%s", endpoint, tableNameWithType, paramsStr);
+  }
+
   private String generateCheckReloadSegmentsServerURL(String tableNameWithType, String endpoint) {
     tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
     return String.format("%s/tables/%s/segments/needReload", endpoint, tableNameWithType);
@@ -470,4 +478,17 @@ private String generateColumnsParam(List<String> columns) {
     paramsStr = String.join("&", params);
     return paramsStr;
   }
+
+  private String generateSegmentsParam(List<String> values) {
+    String paramsStr = "";
+    if (values == null || values.isEmpty()) {
+      return paramsStr;
+    }
+    List<String> params = new ArrayList<>(values.size());
+    for (String value : values) {
+      params.add(String.format("segmentsToInclude=%s", value));
+    }
+    paramsStr = String.join("&", params);
+    return paramsStr;
+  }
 }
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/util/TableMetadataReader.java b/pinot-controller/src/main/java/org/apache/pinot/controller/util/TableMetadataReader.java
@@ -24,6 +24,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -100,21 +101,28 @@ private JsonNode getSegmentsMetadataInternal(String tableNameWithType, List<Stri
     ServerSegmentMetadataReader serverSegmentMetadataReader =
         new ServerSegmentMetadataReader(_executor, _connectionManager);
 
-    // Filter segments that we need
+    List<String> serverURL = new java.util.ArrayList<>(List.of());
     for (Map.Entry<String, List<String>> serverToSegment : serverToSegmentsMap.entrySet()) {
-      List<String> segments = serverToSegment.getValue();
-      if (segmentsToInclude != null && !segmentsToInclude.isEmpty()) {
-        segments.retainAll(segmentsToInclude);
-      }
+      String serverInstance = serverToSegment.getKey();
+      serverURL.add(serverSegmentMetadataReader.generateTableMetadataServerURL(tableNameWithType,
+              columns, segmentsToInclude, endpoints.get(serverInstance)));
     }
 
-    List<String> segmentsMetadata =
-        serverSegmentMetadataReader.getSegmentMetadataFromServer(tableNameWithType, serverToSegmentsMap, endpoints,
-            columns, timeoutMs);
+    CompletionServiceHelper completionServiceHelper =
+            new CompletionServiceHelper(_executor, _connectionManager, endpoints);
+    CompletionServiceHelper.CompletionServiceResponse serviceResponse =
+            completionServiceHelper.doMultiGetRequest(serverURL, tableNameWithType, false, timeoutMs);
+
     Map<String, JsonNode> response = new HashMap<>();
-    for (String segmentMetadata : segmentsMetadata) {
-      JsonNode responseJson = JsonUtils.stringToJsonNode(segmentMetadata);
-      response.put(responseJson.get("segmentName").asText(), responseJson);
+    for (Map.Entry<String, String> serverToSegmentsMetadata : serviceResponse._httpResponses.entrySet()) {
+      JsonNode responseJson = JsonUtils.stringToJsonNode(serverToSegmentsMetadata.getValue());
+      Iterator<Map.Entry<String, JsonNode>> fields = responseJson.fields();
+      while (fields.hasNext()) {
+        Map.Entry<String, JsonNode> field = fields.next();
+        String segmentName = field.getKey();
+        JsonNode segmentJson = field.getValue();
+        response.put(segmentName, segmentJson);
+      }
     }
     return JsonUtils.objectToJsonNode(response);
   }

diff --git a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java
@@ -191,7 +191,9 @@ public String listTableSegments(
   public String getSegmentMetadata(
       @ApiParam(value = "Table Name with type", required = true) @PathParam("tableName") String tableName,
       @ApiParam(value = "Column name", allowMultiple = true) @QueryParam("columns") @DefaultValue("")
-      List<String> columns, @Context HttpHeaders headers)
+      List<String> columns,
+      @ApiParam(value = "List of segments to fetch metadata for") @QueryParam("segmentsToInclude") @DefaultValue("")
+      List<String> segmentsToInclude, @Context HttpHeaders headers)
       throws WebApplicationException {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
     InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager();
@@ -219,8 +221,12 @@ public String getSegmentMetadata(
       }
     }
     Set<String> columnSet = allColumns ? null : new HashSet<>(decodedColumns);
-
-    List<SegmentDataManager> segmentDataManagers = tableDataManager.acquireAllSegments();
+    List<SegmentDataManager> segmentDataManagers;
+    if (segmentsToInclude != null && !segmentsToInclude.isEmpty()) {
+      segmentDataManagers = tableDataManager.acquireSegments(segmentsToInclude, new ArrayList<>());
+    } else {
+      segmentDataManagers = tableDataManager.acquireAllSegments();
+    }
     long totalSegmentSizeBytes = 0;
     long totalNumRows = 0;
     Map<String, Double> columnLengthMap = new HashMap<>();