Modify segment metadata call

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java

@@ -911,7 +911,8 @@ public String getServerMetadata(
       @ApiParam(value = "Name of the table", required = true) @PathParam("tableName") String tableName,
       @ApiParam(value = "OFFLINE|REALTIME") @QueryParam("type") String tableTypeStr,
       @ApiParam(value = "Columns name", allowMultiple = true) @QueryParam("columns") @DefaultValue("")
-          List<String> columns, @Context HttpHeaders headers) {
+      List<String> columns, @ApiParam(value = "Segments name", allowMultiple = true) @QueryParam("segmentsToInclude")
+      @DefaultValue("") Set<String> segmentsToInclude, @Context HttpHeaders headers) {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
     LOGGER.info("Received a request to fetch metadata for all segments for table {}", tableName);
     TableType tableType = Constants.validateTableType(tableTypeStr);
@@ -920,7 +921,7 @@ public String getServerMetadata(
         ResourceUtils.getExistingTableNamesWithType(_pinotHelixResourceManager, tableName, tableType, LOGGER).get(0);
     String segmentsMetadata;
     try {
-      JsonNode segmentsMetadataJson = getSegmentsMetadataFromServer(tableNameWithType, columns);
+      JsonNode segmentsMetadataJson = getSegmentsMetadataFromServer(tableNameWithType, columns, segmentsToInclude);
       segmentsMetadata = JsonUtils.objectToPrettyString(segmentsMetadataJson);
     } catch (InvalidConfigException e) {
       throw new ControllerApplicationException(LOGGER, e.getMessage(), Status.BAD_REQUEST);
@@ -1142,12 +1143,14 @@ public List<Map<TableType, List<String>>> getSelectedSegments(
    * @param columns name of the columns
    * @return Map<String, String>  metadata of the table segments -> map of segment name to its metadata
    */
-  private JsonNode getSegmentsMetadataFromServer(String tableNameWithType, List<String> columns)
+  private JsonNode getSegmentsMetadataFromServer(String tableNameWithType, List<String> columns,
+      Set<String> segmentsToInclude)
       throws InvalidConfigException, IOException {
     TableMetadataReader tableMetadataReader =
         new TableMetadataReader(_executor, _connectionManager, _pinotHelixResourceManager);
     return tableMetadataReader
-        .getSegmentsMetadata(tableNameWithType, columns, _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
+        .getSegmentsMetadata(tableNameWithType, columns, segmentsToInclude,
+            _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
   }
 
   @POST

pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerSegmentMetadataReader.java

@@ -442,6 +442,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(segmentsToInclude);
+    return String.format("%s/tables/%s/segments/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);
@@ -505,4 +513,17 @@ private String generateStaleSegmentsServerURL(String tableNameWithType, String e
     tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
     return String.format("%s/tables/%s/segments/isStale", endpoint, tableNameWithType);
   }
+
+  private String generateSegmentsParam(Set<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;
+  }
 }

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;
@@ -101,21 +102,26 @@ 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();
+        response.put(field.getKey(), field.getValue());
+      }
     }
     return JsonUtils.objectToJsonNode(response);
   }

pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pinot.server.api.resources;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.base.Preconditions;
 import io.swagger.annotations.Api;
 import io.swagger.annotations.ApiKeyAuthDefinition;
@@ -195,7 +196,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", allowMultiple = true) @QueryParam("segmentsToInclude")
+      @DefaultValue("") List<String> segmentsToInclude, @Context HttpHeaders headers)
       throws WebApplicationException {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
     InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager();
@@ -214,6 +217,14 @@ public String getSegmentMetadata(
       decodedColumns.add(URIUtils.decode(column));
     }
 
+    List<String> decodedSegments = new ArrayList<>();
+    if (segmentsToInclude != null && !segmentsToInclude.isEmpty()) {
+      for (String segment : segmentsToInclude) {
+        if (!segment.isEmpty()) {
+          decodedSegments.add(URIUtils.decode(segment));
+        }
+      }
+    }
     boolean allColumns = false;
     // For robustness, loop over all columns, if any of the columns is "*", return metadata for all columns.
     for (String column : decodedColumns) {
@@ -223,8 +234,12 @@ public String getSegmentMetadata(
       }
     }
     Set<String> columnSet = allColumns ? null : new HashSet<>(decodedColumns);
-
-    List<SegmentDataManager> segmentDataManagers = tableDataManager.acquireAllSegments();
+    List<SegmentDataManager> segmentDataManagers;
+    if (!decodedSegments.isEmpty()) {
+      segmentDataManagers = tableDataManager.acquireSegments(decodedSegments, new ArrayList<>());
+    } else {
+      segmentDataManagers = tableDataManager.acquireAllSegments();
+    }
     long totalSegmentSizeBytes = 0;
     long totalNumRows = 0;
     Map<String, Double> columnLengthMap = new HashMap<>();
@@ -402,6 +417,62 @@ public String getSegmentMetadata(
     }
   }
 
+  @GET
+  @Encoded
+  @Path("/tables/{tableName}/segments/metadata")
+  @Produces(MediaType.APPLICATION_JSON)
+  @ApiOperation(value = "Provide segments metadata", notes = "Provide segments metadata for the segments on server")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Success"),
+      @ApiResponse(code = 500, message = "Internal server error", response = ErrorInfo.class),
+      @ApiResponse(code = 404, message = "Table or segment not found", response = ErrorInfo.class)
+  })
+  public String getSegmentsMetadata(
+      @ApiParam(value = "Table name including type", required = true, example = "myTable_OFFLINE")
+      @PathParam("tableName") String tableName,
+      @ApiParam(value = "Segment names to include", allowMultiple = true) @QueryParam("segmentsToInclude")
+      @DefaultValue("") List<String> segmentsToInclude,
+      @ApiParam(value = "Column name", allowMultiple = true) @QueryParam("columns") @DefaultValue("")
+      List<String> columns, @Context HttpHeaders headers) {
+    tableName = DatabaseUtils.translateTableName(tableName, headers);
+    TableDataManager tableDataManager = ServerResourceUtils.checkGetTableDataManager(_serverInstance, tableName);
+    // decode columns and segments
+    List<String> decodedSegments = new ArrayList<>();
+    if (segmentsToInclude != null && !segmentsToInclude.isEmpty()) {
+      for (String segment : segmentsToInclude) {
+        if (!segment.isEmpty()) {
+          decodedSegments.add(URIUtils.decode(segment));
+        }
+      }
+    }
+    List<SegmentDataManager> segmentDataManagers;
+    if (!decodedSegments.isEmpty()) {
+      segmentDataManagers = tableDataManager.acquireSegments(decodedSegments, new ArrayList<>());
+    } else {
+      segmentDataManagers = tableDataManager.acquireAllSegments();
+    }
+    for (int i = 0; i < columns.size(); i++) {
+      columns.set(i, URIUtils.decode(columns.get(i)));
+    }
+    // get metadata for every segment in the list
+    Map<String, JsonNode> response = new HashMap<>();
+    for (SegmentDataManager segmentDataManager: segmentDataManagers) {
+      String segmentName = segmentDataManager.getSegmentName();
+      try {
+        String segmentMetadata = SegmentMetadataFetcher.getSegmentMetadata(segmentDataManager, columns);
+        JsonNode segmentMetadataJson = JsonUtils.stringToJsonNode(segmentMetadata);
+        response.put(segmentName, segmentMetadataJson);
+      } catch (Exception e) {
+        LOGGER.error("Failed to convert table {} segment {} to json", tableName, segmentName);
+        throw new WebApplicationException("Failed to convert segment metadata to json",
+            Response.Status.INTERNAL_SERVER_ERROR);
+      } finally {
+        tableDataManager.releaseSegment(segmentDataManager);
+      }
+    }
+    return ResourceUtils.convertToJsonString(response);
+  }
+
   @GET
   @Path("/tables/{tableName}/segments/crc")
   @Produces(MediaType.APPLICATION_JSON)