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 781140a978b..35df577c868 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 @@ -400,7 +400,7 @@ public ValidDocIdsBitmapResponse getValidDocIdsBitmapFromServer(String tableName private String generateAggregateSegmentMetadataServerURL(String tableNameWithType, List columns, String endpoint) { tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8); - String paramsStr = generateColumnsParam(columns); + String paramsStr = generateColumnsParam("columns", columns); return String.format("%s/tables/%s/metadata?%s", endpoint, tableNameWithType, paramsStr); } @@ -408,10 +408,18 @@ private String generateSegmentMetadataServerURL(String tableNameWithType, String String endpoint) { tableNameWithType = URLEncoder.encode(tableNameWithType, StandardCharsets.UTF_8); segmentName = URLEncoder.encode(segmentName, StandardCharsets.UTF_8); - String paramsStr = generateColumnsParam(columns); + String paramsStr = generateColumnsParam("columns", columns); 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", columns) + + generateColumnsParam("segmentsToInclude", 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); @@ -458,14 +466,14 @@ private Pair generateValidDocIdsMetadataURL(String tableNameWith return Pair.of(url, jsonTableSegments); } - private String generateColumnsParam(List columns) { + private String generateColumnsParam(String param, List values) { String paramsStr = ""; - if (columns == null || columns.isEmpty()) { + if (values == null || values.isEmpty()) { return paramsStr; } - List params = new ArrayList<>(columns.size()); - for (String column : columns) { - params.add(String.format("columns=%s", column)); + List params = new ArrayList<>(values.size()); + for (String value : values) { + params.add(String.format("%s=%s", param, 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 c9e87b396b5..ccec2c454a4 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 @@ -100,21 +100,40 @@ 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()); + Set segmentNames = new HashSet<>(); + responseJson.fieldNames().forEachRemaining(segmentNames::add); + if (segmentsToInclude != null && !segmentsToInclude.isEmpty()) { + segmentNames.retainAll(segmentsToInclude); + } + for (String segmentName : segmentNames) { + JsonNode segmentJson = responseJson.get(segmentName); + response.put(segmentName, segmentJson); + } + +// Iterator> fields = responseJson.fields(); +// while (fields.hasNext()) { +// Map.Entry field = fields.next(); +// String segmentName = field.getKey(); +// JsonNode segmentJson = field.getValue(); +// if (segmentsToInclude == null || segmentsToInclude.isEmpty() || segmentsToInclude.contains(segmentName)) { +// response.put(segmentName, segmentJson); +// } +// } } return JsonUtils.objectToJsonNode(response); } diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentRestletResourceTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentRestletResourceTest.java index 78a9923d999..80b86b82105 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentRestletResourceTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentRestletResourceTest.java @@ -31,6 +31,8 @@ import org.apache.pinot.segment.spi.SegmentMetadata; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.spi.utils.builder.ControllerRequestURLBuilder; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; @@ -169,6 +171,37 @@ public void testSegmentCrcApi() checkCrcRequest(rawTableName, segmentMetadataTable, 9); } + @Test + public void testGetServerMetadataAPI() throws Exception { + // Adding table and schema + String rawTableName = "serverMetadataTestTable"; + Schema schema = new Schema.SchemaBuilder() + .setSchemaName(rawTableName) + .addSingleValueDimension("dimension1", FieldSpec.DataType.STRING) + .addMetric("metric1", FieldSpec.DataType.INT) + .build(); + PinotHelixResourceManager resourceManager = TEST_INSTANCE.getHelixResourceManager(); + resourceManager.addSchema(schema, true, false); + String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(rawTableName); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(rawTableName).setNumReplicas(1) + .setDeletedSegmentsRetentionPeriod("0d").build(); + resourceManager.addTable(tableConfig); + // Upload Segments + Map segmentMetadataTable = new HashMap<>(); + for (int i = 0; i < 2; i++) { + SegmentMetadata segmentMetadata = SegmentMetadataMockUtils.mockSegmentMetadata(rawTableName); + resourceManager.addNewSegment(offlineTableName, segmentMetadata, "downloadUrl"); + segmentMetadataTable.put(segmentMetadata.getName(), segmentMetadata); + } + // check that metadata matches + String sampleSegment = segmentMetadataTable.keySet().iterator().next(); + List columns = List.of("dimension1"); + String resp = ControllerTest.sendGetRequest( + TEST_INSTANCE.getControllerRequestURLBuilder().forSegmentsMetadataFromServer(rawTableName)); + Map fetchedMetadata = JsonUtils.stringToObject(resp, Map.class); + assertEquals(fetchedMetadata.size(), 2); + } + @Test public void testDeleteSegmentsWithTimeWindow() throws Exception { 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 ce85ec3f312..e1bc2c9f844 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<>();