Modify segment metadata call (#14250)

* Modify segment metadata call

* fix tests

* add per server metadata api

* remove v2 api

* Add v2 API and server test, separate functions

* fix test

* merge getSegmentsMetadataInternalV2 into getSegmentsMetadataInternal

* Address comment and add test

* final fixes

* address comments

* address comments

* update log

Author: cypherean

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

@@ -926,17 +926,20 @@ private void deleteSegmentsInternal(String tableNameWithType, List<String> segme
   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) {
+      @Encoded @ApiParam(value = "Segments to include (all if not specified)", allowMultiple = true)
+      @QueryParam("segments") @Nullable List<String> segments,
+      @Encoded @ApiParam(value = "Columns name", allowMultiple = true) @QueryParam("columns")
+      @Nullable List<String> columns, @Context HttpHeaders headers) {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
-    LOGGER.info("Received a request to fetch metadata for all segments for table {}", tableName);
+    String segmentCount = (segments == null) ? "all" : String.valueOf(segments.size());
+    LOGGER.info("Received a request to fetch metadata for {} segments for table {}", segmentCount, tableName);
     TableType tableType = Constants.validateTableType(tableTypeStr);
 
     String tableNameWithType =
         ResourceUtils.getExistingTableNamesWithType(_pinotHelixResourceManager, tableName, tableType, LOGGER).get(0);
     String segmentsMetadata;
     try {
-      JsonNode segmentsMetadataJson = getSegmentsMetadataFromServer(tableNameWithType, columns);
+      JsonNode segmentsMetadataJson = getSegmentsMetadataFromServer(tableNameWithType, columns, segments);
       segmentsMetadata = JsonUtils.objectToPrettyString(segmentsMetadataJson);
     } catch (InvalidConfigException e) {
       throw new ControllerApplicationException(LOGGER, e.getMessage(), Status.BAD_REQUEST);
@@ -1156,14 +1159,17 @@ public List<Map<TableType, List<String>>> getSelectedSegments(
    * This is a helper method to get the metadata for all segments for a given table name.
    * @param tableNameWithType name of the table along with its type
    * @param columns name of the columns
+   * @param segments name of the segments to include in metadata
    * @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, @Nullable List<String> columns,
+      @Nullable List<String> segments)
       throws InvalidConfigException, IOException {
     TableMetadataReader tableMetadataReader =
         new TableMetadataReader(_executor, _connectionManager, _pinotHelixResourceManager);
     return tableMetadataReader
-        .getSegmentsMetadata(tableNameWithType, columns, _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
+        .getSegmentsMetadata(tableNameWithType, columns, segments,
+            _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
   }
 
   @POST

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

@@ -39,6 +39,7 @@
 import javax.ws.rs.client.ClientBuilder;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
+import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManager;
 import org.apache.hc.client5.http.io.HttpClientConnectionManager;
@@ -65,6 +66,8 @@
  */
 public class ServerSegmentMetadataReader {
   private static final Logger LOGGER = LoggerFactory.getLogger(ServerSegmentMetadataReader.class);
+  private static final String COLUMNS_KEY = "columns";
+  private static final String SEGMENTS_KEY = "segments";
 
   private final Executor _executor;
   private final HttpClientConnectionManager _connectionManager;
@@ -430,18 +433,25 @@ public Map<String, TableStaleSegmentResponse> getStaleSegmentsFromServer(
   private String generateAggregateSegmentMetadataServerURL(String tableNameWithType, List<String> columns,
       String endpoint) {
     tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
-    String paramsStr = generateColumnsParam(columns);
+    String paramsStr = generateParam(COLUMNS_KEY, columns);
     return String.format("%s/tables/%s/metadata?%s", endpoint, tableNameWithType, paramsStr);
   }
 
-  private String generateSegmentMetadataServerURL(String tableNameWithType, String segmentName, List<String> columns,
-      String endpoint) {
+  public String generateSegmentMetadataServerURL(String tableNameWithType, String segmentName,
+      @Nullable List<String> columns, String endpoint) {
     tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
     segmentName = URLEncoder.encode(segmentName, StandardCharsets.UTF_8);
-    String paramsStr = generateColumnsParam(columns);
+    String paramsStr = generateParam(COLUMNS_KEY, columns);
     return String.format("%s/tables/%s/segments/%s/metadata?%s", endpoint, tableNameWithType, segmentName, paramsStr);
   }
 
+  public String generateTableMetadataServerURL(String tableNameWithType, @Nullable List<String> columns,
+      @Nullable List<String> segmentsToInclude, String endpoint) {
+    tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
+    String paramsStr = generateParam(COLUMNS_KEY, columns) + "&" + generateParam(SEGMENTS_KEY, 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);
@@ -488,24 +498,24 @@ private Pair<String, String> generateValidDocIdsMetadataURL(String tableNameWith
     return Pair.of(url, jsonTableSegments);
   }
 
-  private String generateColumnsParam(List<String> columns) {
+  private String generateStaleSegmentsServerURL(String tableNameWithType, String endpoint) {
+    tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
+    return String.format("%s/tables/%s/segments/isStale", endpoint, tableNameWithType);
+  }
+
+  private String generateParam(String key, List<String> values) {
     String paramsStr = "";
-    if (columns == null || columns.isEmpty()) {
+    if (CollectionUtils.isEmpty(values)) {
       return paramsStr;
     }
-    List<String> params = new ArrayList<>(columns.size());
-    for (String column : columns) {
-      params.add(String.format("columns=%s", column));
+    List<String> params = new ArrayList<>(values.size());
+    for (String value : values) {
+      params.add(key + "=" + value);
     }
     paramsStr = String.join("&", params);
     return paramsStr;
   }
 
-  private String generateStaleSegmentsServerURL(String tableNameWithType, String endpoint) {
-    tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8);
-    return String.format("%s/tables/%s/segments/isStale", endpoint, tableNameWithType);
-  }
-
   public class TableReloadResponse {
     private int _numFailedResponses;
     private List<String> _serverReloadResponses;

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

@@ -19,8 +19,11 @@
 package org.apache.pinot.controller.util;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.common.collect.BiMap;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -29,6 +32,7 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import org.apache.hc.client5.http.io.HttpClientConnectionManager;
 import org.apache.helix.model.ExternalView;
 import org.apache.pinot.common.exception.InvalidConfigException;
@@ -40,6 +44,8 @@
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.JsonUtils;
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 
 /**
@@ -50,6 +56,7 @@
  * the column indexes available.
  */
 public class TableMetadataReader {
+  private static final Logger log = LoggerFactory.getLogger(TableMetadataReader.class);
   private final Executor _executor;
   private final HttpClientConnectionManager _connectionManager;
   private final PinotHelixResourceManager _pinotHelixResourceManager;
@@ -127,50 +134,95 @@ private TableReloadJsonResponse processSegmentMetadataReloadResponse(
 
   /**
    * This api takes in list of segments for which we need the metadata.
+   * This calls the server to get the metadata for all segments instead of making a call per segment.
    */
-  public JsonNode getSegmentsMetadata(String tableNameWithType, List<String> columns, Set<String> segmentsToInclude,
-      int timeoutMs)
+  public JsonNode getSegmentsMetadata(String tableNameWithType, @Nullable List<String> columns,
+      @Nullable List<String> segments, int timeoutMs)
       throws InvalidConfigException, IOException {
-    return getSegmentsMetadataInternal(tableNameWithType, columns, segmentsToInclude, timeoutMs);
+    return getSegmentsMetadataInternal(tableNameWithType, columns, segments, timeoutMs);
   }
 
-  private JsonNode getSegmentsMetadataInternal(String tableNameWithType, List<String> columns,
-      Set<String> segmentsToInclude, int timeoutMs)
+  /**
+   * Common helper used by both the new (server-level) and legacy (segment-level) endpoints.
+   */
+  private JsonNode fetchAndAggregateMetadata(List<String> urls, BiMap<String, String> endpoints, boolean perSegmentJson,
+      String tableNameWithType, int timeoutMs)
       throws InvalidConfigException, IOException {
-    final Map<String, List<String>> serverToSegmentsMap =
-        _pinotHelixResourceManager.getServerToSegmentsMap(tableNameWithType);
-    BiMap<String, String> endpoints =
-        _pinotHelixResourceManager.getDataInstanceAdminEndpoints(serverToSegmentsMap.keySet());
-    ServerSegmentMetadataReader serverSegmentMetadataReader =
-        new ServerSegmentMetadataReader(_executor, _connectionManager);
+    CompletionServiceHelper cs = new CompletionServiceHelper(_executor, _connectionManager, endpoints);
+    CompletionServiceHelper.CompletionServiceResponse resp =
+        cs.doMultiGetRequest(urls, tableNameWithType, perSegmentJson, timeoutMs);
+    // all requests will fail if new server endpoint is not available
+    if (resp._failedResponseCount > 0) {
+      throw new RuntimeException(String.format("Got %d failed responses from total %d server instances. "
+          + "Falling back to legacy segment metadata api", resp._failedResponseCount, urls.size()));
+    }
 
-    // Filter segments that we need
-    for (Map.Entry<String, List<String>> serverToSegment : serverToSegmentsMap.entrySet()) {
-      List<String> segments = serverToSegment.getValue();
-      if (segmentsToInclude != null && !segmentsToInclude.isEmpty()) {
-        segments.retainAll(segmentsToInclude);
+    ObjectMapper mapper = new ObjectMapper();
+    ObjectNode aggregatedNode = mapper.createObjectNode();
+    for (String body : resp._httpResponses.values()) {
+      JsonNode node = JsonUtils.stringToJsonNode(body);
+      // legacy returns one JSON per segment; new returns one JSON with many fields
+      if (perSegmentJson) {
+        String segmentName = node.get("segmentName").asText();
+        aggregatedNode.set(segmentName, node);
+      } else {
+        node.fields().forEachRemaining(entry -> aggregatedNode.set(entry.getKey(), entry.getValue()));
       }
     }
+    return aggregatedNode;
+  }
 
-    List<String> segmentsMetadata =
-        serverSegmentMetadataReader.getSegmentMetadataFromServer(tableNameWithType, serverToSegmentsMap, endpoints,
-            columns, timeoutMs);
-    Map<String, JsonNode> response = new HashMap<>();
-    for (String segmentMetadata : segmentsMetadata) {
-      JsonNode responseJson = JsonUtils.stringToJsonNode(segmentMetadata);
-      response.put(responseJson.get("segmentName").asText(), responseJson);
+  private List<String> buildTableLevelUrls(Map<String, List<String>> serverToSegs, BiMap<String, String> endpoints,
+      String tableNameWithType, List<String> columns, List<String> segmentsFilter, ServerSegmentMetadataReader reader) {
+    List<String> urls = new ArrayList<>(serverToSegs.size());
+    for (String server : serverToSegs.keySet()) {
+      urls.add(reader.generateTableMetadataServerURL(
+          tableNameWithType, columns, segmentsFilter, endpoints.get(server)));
     }
-    return JsonUtils.objectToJsonNode(response);
+    return urls;
   }
 
-  /**
-   * This method retrieves the full segment metadata for a given table.
-   * Currently supports only OFFLINE tables.
-   * @return a map of segmentName to its metadata
-   */
-  public JsonNode getSegmentsMetadata(String tableNameWithType, List<String> columns, int timeoutMs)
+  private List<String> buildSegmentLevelUrls(Map<String, List<String>> serverToSegs, BiMap<String, String> endpoints,
+      String tableNameWithType, List<String> columns, List<String> segmentsFilter, ServerSegmentMetadataReader reader) {
+    List<String> urls = new ArrayList<>();
+    for (Map.Entry<String, List<String>> e : serverToSegs.entrySet()) {
+      for (String segment : e.getValue()) {
+        if (segmentsFilter == null || segmentsFilter.isEmpty()
+            || segmentsFilter.contains(segment)) {
+          urls.add(reader.generateSegmentMetadataServerURL(
+              tableNameWithType, segment, columns, endpoints.get(e.getKey())));
+        }
+      }
+    }
+    return urls;
+  }
+
+  private JsonNode getSegmentsMetadataInternal(String tableNameWithType, @Nullable List<String> columns,
+      @Nullable List<String> segments, int timeoutMs)
       throws InvalidConfigException, IOException {
-    return getSegmentsMetadataInternal(tableNameWithType, columns, null, timeoutMs);
+    Map<String, List<String>> serverToSegs =
+        _pinotHelixResourceManager.getServerToSegmentsMap(tableNameWithType);
+    BiMap<String, String> endpoints =
+        _pinotHelixResourceManager.getDataInstanceAdminEndpoints(serverToSegs.keySet());
+    ServerSegmentMetadataReader reader =
+        new ServerSegmentMetadataReader(_executor, _connectionManager);
+
+    // try table level endpoint first
+    try {
+      List<String> tableUrls = buildTableLevelUrls(serverToSegs, endpoints,
+          tableNameWithType, columns, segments, reader);
+      return fetchAndAggregateMetadata(tableUrls, endpoints, /*perSegmentJson=*/false,
+          tableNameWithType, timeoutMs);
+    } catch (RuntimeException e) {
+      log.warn("Failed to fetch table metadata for table {} using new server endpoint, falling back to legacy "
+              + "per-segment endpoint", tableNameWithType, e);
+    }
+
+    // legacy per segment endpoint
+    List<String> segmentUrls = buildSegmentLevelUrls(serverToSegs, endpoints,
+        tableNameWithType, columns, segments, reader);
+    return fetchAndAggregateMetadata(segmentUrls, endpoints.inverse(), /*perSegmentJson=*/true,
+        tableNameWithType, timeoutMs);
   }
 
   /**

pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java

@@ -873,6 +873,14 @@ protected JsonNode getColumnIndexSize(String column)
         .get("columnIndexSizeMap").get(column);
   }
 
+  /**
+   * Get all segment names for a given tableName and tableType.
+   */
+  protected List<String> getSegmentNames(String tableName, @Nullable String tableType)
+      throws Exception {
+    return getControllerRequestClient().listSegments(tableName, tableType, true);
+  }
+
   protected List<ValidDocIdsMetadataInfo> getValidDocIdsMetadata(String tableNameWithType,
       ValidDocIdsType validDocIdsType)
       throws Exception {