zhtaoxiang commented on code in PR #9825: URL: https://github.com/apache/pinot/pull/9825#discussion_r1036364399
########## pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java: ########## @@ -276,6 +299,87 @@ public List<SegmentConversionResult> executeTask(PinotTaskConfig pinotTaskConfig } } + private void pushSegment(String tableName, Map<String, String> taskConfigs, URI outputSegmentTarURI, + List<Header> headers, List<NameValuePair> parameters) throws Exception { + String pushMode = taskConfigs.get(BatchConfigProperties.PUSH_MODE); + LOGGER.info("Trying to push Pinot segment with push mode {} from {}", pushMode, outputSegmentTarURI); + + PushJobSpec pushJobSpec = new PushJobSpec(); + pushJobSpec.setPushAttempts(DEFUALT_PUSH_ATTEMPTS); + pushJobSpec.setPushParallelism(DEFAULT_PUSH_PARALLELISM); + pushJobSpec.setPushRetryIntervalMillis(DEFAULT_PUSH_RETRY_INTERVAL_MILLIS); + pushJobSpec.setSegmentUriPrefix(taskConfigs.get(BatchConfigProperties.PUSH_SEGMENT_URI_PREFIX)); + pushJobSpec.setSegmentUriSuffix(taskConfigs.get(BatchConfigProperties.PUSH_SEGMENT_URI_SUFFIX)); + + SegmentGenerationJobSpec spec = generatePushJobSpec(tableName, taskConfigs, pushJobSpec); + + URI outputSegmentDirURI = null; Review Comment: Should we fail early is outputSegmentDirURI is null? ########## pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java: ########## @@ -276,6 +299,87 @@ public List<SegmentConversionResult> executeTask(PinotTaskConfig pinotTaskConfig } } + private void pushSegment(String tableName, Map<String, String> taskConfigs, URI outputSegmentTarURI, + List<Header> headers, List<NameValuePair> parameters) throws Exception { + String pushMode = taskConfigs.get(BatchConfigProperties.PUSH_MODE); + LOGGER.info("Trying to push Pinot segment with push mode {} from {}", pushMode, outputSegmentTarURI); + + PushJobSpec pushJobSpec = new PushJobSpec(); + pushJobSpec.setPushAttempts(DEFUALT_PUSH_ATTEMPTS); + pushJobSpec.setPushParallelism(DEFAULT_PUSH_PARALLELISM); + pushJobSpec.setPushRetryIntervalMillis(DEFAULT_PUSH_RETRY_INTERVAL_MILLIS); + pushJobSpec.setSegmentUriPrefix(taskConfigs.get(BatchConfigProperties.PUSH_SEGMENT_URI_PREFIX)); + pushJobSpec.setSegmentUriSuffix(taskConfigs.get(BatchConfigProperties.PUSH_SEGMENT_URI_SUFFIX)); + + SegmentGenerationJobSpec spec = generatePushJobSpec(tableName, taskConfigs, pushJobSpec); + + URI outputSegmentDirURI = null; + if (taskConfigs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) { + outputSegmentDirURI = URI.create(taskConfigs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)); + } + try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(taskConfigs, outputSegmentDirURI)) { + switch (BatchConfigProperties.SegmentPushType.valueOf(pushMode.toUpperCase())) { + case TAR: + try (PinotFS pinotFS = MinionTaskUtils.getLocalPinotFs()) { + SegmentPushUtils.pushSegments( + spec, pinotFS, Arrays.asList(outputSegmentTarURI.toString()), headers, parameters); + } catch (RetriableOperationException | AttemptsExceededException e) { + throw new RuntimeException(e); + } + break; + case METADATA: + try { + Map<String, String> segmentUriToTarPathMap = + SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI, pushJobSpec, + new String[]{outputSegmentTarURI.toString()}); + SegmentPushUtils.sendSegmentUriAndMetadata(spec, outputFileFS, segmentUriToTarPathMap, headers, parameters); + } catch (RetriableOperationException | AttemptsExceededException e) { + throw new RuntimeException(e); + } + break; + default: + throw new UnsupportedOperationException("Unrecognized push mode - " + pushMode); + } + } + } + + private SegmentGenerationJobSpec generatePushJobSpec(String tableName, Map<String, String> taskConfigs, + PushJobSpec pushJobSpec) { + + TableSpec tableSpec = new TableSpec(); + tableSpec.setTableName(tableName); + + PinotClusterSpec pinotClusterSpec = new PinotClusterSpec(); + pinotClusterSpec.setControllerURI(taskConfigs.get(BatchConfigProperties.PUSH_CONTROLLER_URI)); + PinotClusterSpec[] pinotClusterSpecs = new PinotClusterSpec[]{pinotClusterSpec}; + + SegmentGenerationJobSpec spec = new SegmentGenerationJobSpec(); + spec.setPushJobSpec(pushJobSpec); + spec.setTableSpec(tableSpec); + spec.setPinotClusterSpecs(pinotClusterSpecs); + spec.setAuthToken(taskConfigs.get(BatchConfigProperties.AUTH_TOKEN)); + + return spec; + } + + private URI moveSegmentToOutputPinotFS(Map<String, String> taskConfigs, File localSegmentTarFile) + throws Exception { + if (!taskConfigs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) { Review Comment: 1/ I feel that the task generator should always make sure the output segment dir is provided? 2/ using the local segment tar file does not help here, because most of the time, servers cannot download segment from it. ########## pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/MinionTaskUtils.java: ########## @@ -67,7 +72,35 @@ static PinotFS getOutputPinotFS(Map<String, String> taskConfigs, URI fileURI) return PinotFSFactory.create(fileURIScheme); } - static PinotFS getLocalPinotFs() { + public static Map<String, String> getPushTaskConfig(String tableName, Map<String, String> taskConfigs, + ClusterInfoAccessor clusterInfoAccessor) { + try { + URI outputDirURI = URI.create(clusterInfoAccessor.getDataDir() + "/" + tableName); Review Comment: I agree that it is not really used in the case when push mode is TAR. However, I feel the current logic is a little bit hard to follow if we don't have the full context. If we can check push mode first then decide what to do, it's easier to follow. WDYT? -- 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: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org