xiangfu0 commented on code in PR #19895:
URL: https://github.com/apache/pinot/pull/19895#discussion_r4288752207
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -507,36 +541,225 @@ public static void
sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, Pin
public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
throws Exception {
+ sendSegmentsUriAndMetadata(spec, segmentUriToMetadataFileMap, headers,
parameters, 1);
+ }
+
+ @VisibleForTesting
+ static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
+ Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters,
+ int parallelism)
+ throws Exception {
Map<String, File> segmentMetadataFileMap = new HashMap<>();
List<String> segmentURIs = new ArrayList<>();
for (Map.Entry<String, File> entry :
segmentUriToMetadataFileMap.entrySet()) {
String segmentName = getSegmentNameFromMetadataFile(entry.getValue());
- segmentMetadataFileMap.put(segmentName, entry.getValue());
+
Preconditions.checkArgument(segmentMetadataFileMap.putIfAbsent(segmentName,
entry.getValue()) == null,
+ "Multiple metadata files for segment: %s", segmentName);
segmentURIs.add(segmentName);
segmentURIs.add(entry.getKey());
}
- pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters);
+ pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters, parallelism);
}
private static void pushSegmentsMetadata(SegmentGenerationJobSpec spec,
Collection<String> segmentURIs,
- Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
+ Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters, int parallelism)
+ throws Exception {
+ List<MetadataBatch> batches = partitionMetadataBatches(segmentURIs,
segmentMetadataFileMap, parallelism);
+ validateMetadataControllers(spec);
+ if (batches.isEmpty()) {
+ return;
+ }
+ if (batches.size() == 1) {
+ FileUploadDownloadClient client =
getOrCreateFileUploadDownloadClient(spec);
+ try {
+ pushMetadataBatch(spec, batches.get(0), headers, parameters, client,
() -> false);
+ } finally {
+ closeFileUploadDownloadClient(spec, client);
+ }
+ return;
+ }
+ pushMetadataBatches(spec, batches, headers, parameters);
+ }
+
+ /// Each batch owns its immutable name/URI pairs and metadata-file mapping;
source files stay caller-owned.
+ private static final class MetadataBatch {
+ private final List<String> _segmentUris;
+ private final Map<String, File> _metadataFiles;
+
+ private MetadataBatch(List<String> segmentUris, Map<String, File>
metadataFiles) {
+ _segmentUris = List.copyOf(segmentUris);
+ _metadataFiles = Map.copyOf(metadataFiles);
+ }
+ }
+
+ private static List<MetadataBatch>
partitionMetadataBatches(Collection<String> segmentURIs,
+ Map<String, File> metadataFiles, int parallelism) {
+ List<String> pairs = new ArrayList<>(segmentURIs);
+ Preconditions.checkArgument(pairs.size() == 2 * metadataFiles.size(),
"Incomplete metadata name/URI pairs");
+ Map<String, String> urisByName = new LinkedHashMap<>();
+ for (int i = 0; i < pairs.size(); i += 2) {
+ String name = pairs.get(i);
+ SegmentNameUtils.validatePartialOrFullSegmentName(name);
+ String uri = pairs.get(i + 1);
+ validateMetadataUris(List.of(uri));
Review Comment:
The download URIs get validated a few times over now: `validateMetadataUris`
on the whole key set in the entry method, then again here one URI at a time
inside the partition loop. Correct, just redundant. Since partitioning already
walks every pair, you could drop the earlier set-level pass and rely on this
one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
