diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerSegmentMetadataReader.java b/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerSegmentMetadataReader.java index 781140a978ba..9583a04e32fb 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerSegmentMetadataReader.java +++ b/pinot-controller/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 columns, + Set 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 columns) { paramsStr = String.join("&", params); return paramsStr; } + + private String generateSegmentsParam(List values) { + String paramsStr = ""; + if (values == null || values.isEmpty()) { + return paramsStr; + } + List 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 index c9e87b396b53..045db6cbefed 100644 --- 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 serverURL = new java.util.ArrayList<>(List.of()); for (Map.Entry> serverToSegment : serverToSegmentsMap.entrySet()) { - List 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 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 response = new HashMap<>(); - for (String segmentMetadata : segmentsMetadata) { - JsonNode responseJson = JsonUtils.stringToJsonNode(segmentMetadata); - response.put(responseJson.get("segmentName").asText(), responseJson); + for (Map.Entry serverToSegmentsMetadata : serviceResponse._httpResponses.entrySet()) { + JsonNode responseJson = JsonUtils.stringToJsonNode(serverToSegmentsMetadata.getValue()); + Iterator> fields = responseJson.fields(); + while (fields.hasNext()) { + Map.Entry 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 index ce85ec3f3123..e1bc2c9f844f 100644 --- 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 columns, @Context HttpHeaders headers) + List columns, + @ApiParam(value = "List of segments to fetch metadata for") @QueryParam("segmentsToInclude") @DefaultValue("") + List segmentsToInclude, @Context HttpHeaders headers) throws WebApplicationException { tableName = DatabaseUtils.translateTableName(tableName, headers); InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager(); @@ -219,8 +221,12 @@ public String getSegmentMetadata( } } Set columnSet = allColumns ? null : new HashSet<>(decodedColumns); - - List segmentDataManagers = tableDataManager.acquireAllSegments(); + List segmentDataManagers; + if (segmentsToInclude != null && !segmentsToInclude.isEmpty()) { + segmentDataManagers = tableDataManager.acquireSegments(segmentsToInclude, new ArrayList<>()); + } else { + segmentDataManagers = tableDataManager.acquireAllSegments(); + } long totalSegmentSizeBytes = 0; long totalNumRows = 0; Map columnLengthMap = new HashMap<>();