From b8ce40239d087ef7be4edc0914d8828f85647923 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 25 Jul 2026 06:15:35 +0000 Subject: [PATCH] Parallelize segment staging copy to deep store Batch Spark/Hadoop generation jobs spent wall-clock time on a single-threaded PinotFS move of hundreds of segment tars from staging to deep store. Parallel copy with bounded concurrency cuts job latency without changing push semantics. --- .../common/SegmentGenerationJobUtils.java | 169 ++++++++++++++++-- .../common/SegmentGenerationJobUtilsTest.java | 157 ++++++++++++++++ .../HadoopSegmentGenerationJobRunner.java | 8 +- .../SparkSegmentGenerationJobRunner.java | 8 +- 4 files changed, 321 insertions(+), 21 deletions(-) diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtils.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtils.java index 1129bcce128a..8ed5128b8571 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtils.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtils.java @@ -29,6 +29,10 @@ import java.nio.file.SimpleFileVisitor; import java.util.ArrayList; import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.apache.commons.io.FileUtils; import org.apache.pinot.common.segment.generation.SegmentGenerationUtils; import org.apache.pinot.common.utils.TarCompressionUtils; @@ -53,6 +57,13 @@ private SegmentGenerationJobUtils() { // Field names in the executionFrameworkSpec/extraConfigs section shared across ingestion frameworks public static final String DEPENDENCY_JAR_DIR = "dependencyJarDir"; public static final String STAGING_DIR = "stagingDir"; + /** + * Optional extraConfigs key controlling parallel staging → deep-store copy threads. + * Default is {@link #DEFAULT_STAGING_COPY_PARALLELISM}. Cap is {@link #MAX_STAGING_COPY_PARALLELISM}. + */ + public static final String STAGING_COPY_PARALLELISM = "stagingCopyParallelism"; + public static final int DEFAULT_STAGING_COPY_PARALLELISM = 10; + public static final int MAX_STAGING_COPY_PARALLELISM = 64; /** * Always use local directory sequence id unless explicitly config: "use.global.directory.sequence.id". @@ -108,27 +119,155 @@ public static void moveLocalTarFileToRemote(File localMetadataTarFile, URI outpu * If is true, and the source file exists in the destination directory, then replace it, otherwise * log a warning and continue. We assume that source and destination directories are on the same filesystem, * so that move() can be used. + * Uses {@link #DEFAULT_STAGING_COPY_PARALLELISM} worker threads. + *

+ * The shared {@link PinotFS} instance must support concurrent {@code move} of distinct paths (true for + * {@code LocalPinotFS} and typical remote implementations). * - * @param fs - * @param sourceDir - * @param destDir - * @param overwrite - * @throws IOException - * @throws URISyntaxException + * @param fs filesystem used for both source and destination + * @param sourceDir source directory URI + * @param destDir destination directory URI + * @param overwrite whether to overwrite existing destination files + * @throws IOException on listing or move failure + * @throws URISyntaxException on URI construction failure */ public static void moveFiles(PinotFS fs, URI sourceDir, URI destDir, boolean overwrite) - throws IOException, URISyntaxException { + throws IOException, URISyntaxException { + moveFiles(fs, sourceDir, destDir, overwrite, DEFAULT_STAGING_COPY_PARALLELISM); + } + + /** + * Move all files from the to the using up to {@code parallelism} threads. + * Directories in the source listing are skipped; parent directories on the destination are created by + * {@link PinotFS#move}. Relative path layout under {@code sourceDir} is preserved. + *

+ * On partial failure, remaining in-flight moves are allowed to finish, then an {@link IOException} is thrown + * with any additional failures attached as suppressed exceptions. + * + * @param fs filesystem used for both source and destination (must be safe for concurrent move of distinct paths) + * @param sourceDir source directory URI + * @param destDir destination directory URI + * @param overwrite whether to overwrite existing destination files + * @param parallelism number of concurrent move workers; values <= 1 run serially + * @throws IOException on listing or move failure + * @throws URISyntaxException on URI construction failure + */ + public static void moveFiles(PinotFS fs, URI sourceDir, URI destDir, boolean overwrite, int parallelism) + throws IOException, URISyntaxException { + List sourceFileUris = listSourceFiles(fs, sourceDir); + if (sourceFileUris.isEmpty()) { + return; + } + int effectiveParallelism = Math.max(1, Math.min(parallelism, sourceFileUris.size())); + LOGGER.info("Moving {} files from [{}] to [{}] with parallelism {}", sourceFileUris.size(), sourceDir, destDir, + effectiveParallelism); + if (effectiveParallelism == 1) { + for (URI sourceFileUri : sourceFileUris) { + moveOneFile(fs, sourceDir, sourceFileUri, destDir, overwrite); + } + return; + } + + ExecutorService executor = Executors.newFixedThreadPool(effectiveParallelism, r -> { + Thread t = new Thread(r, "pinot-staging-copy"); + t.setDaemon(true); + return t; + }); + try { + List> futures = new ArrayList<>(sourceFileUris.size()); + for (URI sourceFileUri : sourceFileUris) { + futures.add(executor.submit(() -> { + moveOneFile(fs, sourceDir, sourceFileUri, destDir, overwrite); + return null; + })); + } + IOException firstFailure = null; + boolean interrupted = false; + for (Future future : futures) { + try { + future.get(); + } catch (InterruptedException e) { + interrupted = true; + future.cancel(true); + if (firstFailure == null) { + firstFailure = new IOException("Interrupted while moving files from " + sourceDir + " to " + destDir, e); + } else { + firstFailure.addSuppressed(e); + } + } catch (Exception e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + if (firstFailure == null) { + firstFailure = cause instanceof IOException ? (IOException) cause + : new IOException("Failed to move files from " + sourceDir + " to " + destDir, cause); + } else { + firstFailure.addSuppressed(cause); + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + if (firstFailure != null) { + throw firstFailure; + } + } finally { + executor.shutdownNow(); + } + } + + /** + * Resolve staging-copy parallelism from job {@code executionFrameworkSpec.extraConfigs}. + * Missing/invalid/non-positive values fall back to {@link #DEFAULT_STAGING_COPY_PARALLELISM}. + * Values above {@link #MAX_STAGING_COPY_PARALLELISM} are capped. + */ + public static int getStagingCopyParallelism(Map extraConfigs) { + if (extraConfigs == null) { + return DEFAULT_STAGING_COPY_PARALLELISM; + } + String value = extraConfigs.get(STAGING_COPY_PARALLELISM); + if (value == null || value.isEmpty()) { + return DEFAULT_STAGING_COPY_PARALLELISM; + } + try { + int parallelism = Integer.parseInt(value.trim()); + if (parallelism < 1) { + LOGGER.warn("Invalid {}={}, using default {}", STAGING_COPY_PARALLELISM, value, + DEFAULT_STAGING_COPY_PARALLELISM); + return DEFAULT_STAGING_COPY_PARALLELISM; + } + if (parallelism > MAX_STAGING_COPY_PARALLELISM) { + LOGGER.warn("Capping {}={} to max {}", STAGING_COPY_PARALLELISM, parallelism, MAX_STAGING_COPY_PARALLELISM); + return MAX_STAGING_COPY_PARALLELISM; + } + return parallelism; + } catch (NumberFormatException e) { + LOGGER.warn("Invalid {}={}, using default {}", STAGING_COPY_PARALLELISM, value, DEFAULT_STAGING_COPY_PARALLELISM); + return DEFAULT_STAGING_COPY_PARALLELISM; + } + } + + private static List listSourceFiles(PinotFS fs, URI sourceDir) + throws IOException, URISyntaxException { + List sourceFileUris = new ArrayList<>(); for (String sourcePath : fs.listFiles(sourceDir, true)) { URI sourceFileUri = SegmentGenerationUtils.getFileURI(sourcePath, sourceDir); - String sourceFilename = SegmentGenerationUtils.getFileName(sourceFileUri); - URI destFileUri = - SegmentGenerationUtils.getRelativeOutputPath(sourceDir, sourceFileUri, destDir).resolve(sourceFilename); - - if (!overwrite && fs.exists(destFileUri)) { - LOGGER.warn("Can't overwrite existing output segment tar file: {}", destFileUri); - } else { - fs.move(sourceFileUri, destFileUri, true); + if (fs.isDirectory(sourceFileUri)) { + continue; } + sourceFileUris.add(sourceFileUri); + } + return sourceFileUris; + } + + private static void moveOneFile(PinotFS fs, URI sourceDir, URI sourceFileUri, URI destDir, boolean overwrite) + throws IOException, URISyntaxException { + String sourceFilename = SegmentGenerationUtils.getFileName(sourceFileUri); + URI destFileUri = + SegmentGenerationUtils.getRelativeOutputPath(sourceDir, sourceFileUri, destDir).resolve(sourceFilename); + if (!overwrite && fs.exists(destFileUri)) { + LOGGER.warn("Can't overwrite existing output segment tar file: {}", destFileUri); + } else { + fs.move(sourceFileUri, destFileUri, true); } } } diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/test/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtilsTest.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/test/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtilsTest.java index 93aba0c5a59a..1d25f345662a 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/test/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtilsTest.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/test/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtilsTest.java @@ -19,14 +19,40 @@ package org.apache.pinot.plugin.ingestion.batch.common; +import java.io.File; +import java.io.IOException; +import java.net.URI; +import java.nio.charset.StandardCharsets; import java.util.HashMap; +import java.util.HashSet; import java.util.Map; +import java.util.Set; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.spi.filesystem.LocalPinotFS; +import org.apache.pinot.spi.filesystem.PinotFS; import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec; import org.testng.Assert; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class SegmentGenerationJobUtilsTest { + private File _tempDir; + + @BeforeMethod + public void setUp() + throws IOException { + _tempDir = new File(FileUtils.getTempDirectory(), "SegmentGenerationJobUtilsTest-" + UUID.randomUUID()); + FileUtils.forceMkdir(_tempDir); + } + + @AfterMethod + public void tearDown() + throws IOException { + FileUtils.deleteDirectory(_tempDir); + } @Test public void testUseGlobalDirectorySequenceId() { @@ -48,4 +74,135 @@ public void testUseGlobalDirectorySequenceId() { spec.setConfigs(Map.of("local.directory.sequence.id", "False")); Assert.assertTrue(SegmentGenerationJobUtils.useGlobalDirectorySequenceId(spec)); } + + @Test + public void testGetStagingCopyParallelism() { + Assert.assertEquals(SegmentGenerationJobUtils.getStagingCopyParallelism(null), + SegmentGenerationJobUtils.DEFAULT_STAGING_COPY_PARALLELISM); + Assert.assertEquals(SegmentGenerationJobUtils.getStagingCopyParallelism(Map.of()), + SegmentGenerationJobUtils.DEFAULT_STAGING_COPY_PARALLELISM); + Assert.assertEquals(SegmentGenerationJobUtils.getStagingCopyParallelism( + Map.of(SegmentGenerationJobUtils.STAGING_COPY_PARALLELISM, "4")), 4); + Assert.assertEquals(SegmentGenerationJobUtils.getStagingCopyParallelism( + Map.of(SegmentGenerationJobUtils.STAGING_COPY_PARALLELISM, "0")), + SegmentGenerationJobUtils.DEFAULT_STAGING_COPY_PARALLELISM); + Assert.assertEquals(SegmentGenerationJobUtils.getStagingCopyParallelism( + Map.of(SegmentGenerationJobUtils.STAGING_COPY_PARALLELISM, "-1")), + SegmentGenerationJobUtils.DEFAULT_STAGING_COPY_PARALLELISM); + Assert.assertEquals(SegmentGenerationJobUtils.getStagingCopyParallelism( + Map.of(SegmentGenerationJobUtils.STAGING_COPY_PARALLELISM, "not-a-number")), + SegmentGenerationJobUtils.DEFAULT_STAGING_COPY_PARALLELISM); + Assert.assertEquals(SegmentGenerationJobUtils.getStagingCopyParallelism( + Map.of(SegmentGenerationJobUtils.STAGING_COPY_PARALLELISM, "999")), + SegmentGenerationJobUtils.MAX_STAGING_COPY_PARALLELISM); + } + + @Test + public void testMoveFilesParallelismOnePreservesRelativeLayout() + throws Exception { + runMoveAndAssertLayout(1); + } + + @Test + public void testMoveFilesParallelismFourPreservesRelativeLayout() + throws Exception { + runMoveAndAssertLayout(4); + } + + @Test + public void testMoveFilesOverwriteFalseSkipsExisting() + throws Exception { + File sourceDir = new File(_tempDir, "src"); + File destDir = new File(_tempDir, "dest"); + FileUtils.forceMkdir(sourceDir); + FileUtils.forceMkdir(destDir); + FileUtils.writeStringToFile(new File(sourceDir, "seg.tar.gz"), "new", StandardCharsets.UTF_8); + FileUtils.writeStringToFile(new File(destDir, "seg.tar.gz"), "old", StandardCharsets.UTF_8); + + LocalPinotFS fs = new LocalPinotFS(); + SegmentGenerationJobUtils.moveFiles(fs, sourceDir.toURI(), destDir.toURI(), false, 2); + + Assert.assertEquals(FileUtils.readFileToString(new File(destDir, "seg.tar.gz"), StandardCharsets.UTF_8), "old"); + // Source left in place when destination was not overwritten + Assert.assertTrue(new File(sourceDir, "seg.tar.gz").exists()); + } + + @Test + public void testMoveFilesOverwriteTrueReplacesExisting() + throws Exception { + File sourceDir = new File(_tempDir, "src"); + File destDir = new File(_tempDir, "dest"); + FileUtils.forceMkdir(sourceDir); + FileUtils.forceMkdir(destDir); + FileUtils.writeStringToFile(new File(sourceDir, "seg.tar.gz"), "new", StandardCharsets.UTF_8); + FileUtils.writeStringToFile(new File(destDir, "seg.tar.gz"), "old", StandardCharsets.UTF_8); + + LocalPinotFS fs = new LocalPinotFS(); + SegmentGenerationJobUtils.moveFiles(fs, sourceDir.toURI(), destDir.toURI(), true, 2); + + Assert.assertEquals(FileUtils.readFileToString(new File(destDir, "seg.tar.gz"), StandardCharsets.UTF_8), "new"); + Assert.assertFalse(new File(sourceDir, "seg.tar.gz").exists()); + } + + @Test + public void testMoveFilesFailsWhenOneMoveFails() + throws Exception { + File sourceDir = new File(_tempDir, "src"); + File destDir = new File(_tempDir, "dest"); + FileUtils.forceMkdir(sourceDir); + FileUtils.forceMkdir(destDir); + FileUtils.writeStringToFile(new File(sourceDir, "ok.tar.gz"), "ok", StandardCharsets.UTF_8); + FileUtils.writeStringToFile(new File(sourceDir, "bad.tar.gz"), "bad", StandardCharsets.UTF_8); + + PinotFS fs = new LocalPinotFS() { + @Override + public boolean move(URI srcUri, URI dstUri, boolean overwrite) + throws IOException { + if (srcUri.getPath().endsWith("bad.tar.gz")) { + throw new IOException("simulated move failure for " + srcUri); + } + return super.move(srcUri, dstUri, overwrite); + } + }; + + try { + SegmentGenerationJobUtils.moveFiles(fs, sourceDir.toURI(), destDir.toURI(), true, 4); + Assert.fail("Expected IOException when one move fails"); + } catch (IOException e) { + Assert.assertTrue(e.getMessage().contains("bad.tar.gz") || (e.getCause() != null && e.getCause().getMessage() + .contains("bad.tar.gz")) || e.getMessage().contains("simulated move failure"), e.getMessage()); + } + } + + private void runMoveAndAssertLayout(int parallelism) + throws Exception { + File sourceDir = new File(_tempDir, "src-" + parallelism); + File destDir = new File(_tempDir, "dest-" + parallelism); + FileUtils.forceMkdir(new File(sourceDir, "part-a")); + FileUtils.forceMkdir(new File(sourceDir, "part-b/nested")); + Set expectedRelativePaths = new HashSet<>(); + for (int i = 0; i < 8; i++) { + String relative; + if (i < 4) { + relative = "part-a/seg-" + i + ".tar.gz"; + } else { + relative = "part-b/nested/seg-" + i + ".tar.gz"; + } + FileUtils.writeStringToFile(new File(sourceDir, relative), "content-" + i, StandardCharsets.UTF_8); + expectedRelativePaths.add(relative); + } + + LocalPinotFS fs = new LocalPinotFS(); + SegmentGenerationJobUtils.moveFiles(fs, sourceDir.toURI(), destDir.toURI(), true, parallelism); + + for (String relative : expectedRelativePaths) { + File destFile = new File(destDir, relative); + Assert.assertTrue(destFile.exists(), "Missing dest file: " + relative); + int start = relative.lastIndexOf("seg-") + 4; + int end = relative.indexOf('.', start); + String expectedContent = "content-" + relative.substring(start, end); + Assert.assertEquals(FileUtils.readFileToString(destFile, StandardCharsets.UTF_8), expectedContent); + Assert.assertFalse(new File(sourceDir, relative).exists(), "Source should be moved: " + relative); + } + } } diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentGenerationJobRunner.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentGenerationJobRunner.java index de173ff93d9b..a70350d6f969 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentGenerationJobRunner.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentGenerationJobRunner.java @@ -275,10 +275,12 @@ public void run() throw new RuntimeException("Job failed: " + job); } - LOGGER.info("Moving segment tars from staging directory [{}] to output directory [{}]", stagingDirURI, - outputDirURI); + int stagingCopyParallelism = SegmentGenerationJobUtils.getStagingCopyParallelism( + _spec.getExecutionFrameworkSpec().getExtraConfigs()); + LOGGER.info("Moving segment tars from staging directory [{}] to output directory [{}] with parallelism {}", + stagingDirURI, outputDirURI, stagingCopyParallelism); SegmentGenerationJobUtils.moveFiles(outputDirFS, new Path(stagingDir, SEGMENT_TAR_SUBDIR_NAME).toUri(), - outputDirURI, _spec.isOverwriteOutput()); + outputDirURI, _spec.isOverwriteOutput(), stagingCopyParallelism); } finally { LOGGER.info("Trying to clean up staging directory: [{}]", stagingDirURI); outputDirFS.delete(stagingDirURI, true); diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java index fccf20175cac..ac7caad98646 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java @@ -327,9 +327,11 @@ public void call(String pathAndIdx) } }); if (stagingDirURI != null) { - LOGGER.info("Trying to move segment tars from staging directory: [{}] to output directory [{}]", stagingDirURI, - outputDirURI); - SegmentGenerationJobUtils.moveFiles(outputDirFS, stagingDirURI, outputDirURI, true); + int stagingCopyParallelism = SegmentGenerationJobUtils.getStagingCopyParallelism( + _spec.getExecutionFrameworkSpec().getExtraConfigs()); + LOGGER.info("Trying to move segment tars from staging directory: [{}] to output directory [{}] " + + "with parallelism {}", stagingDirURI, outputDirURI, stagingCopyParallelism); + SegmentGenerationJobUtils.moveFiles(outputDirFS, stagingDirURI, outputDirURI, true, stagingCopyParallelism); } } finally { if (stagingDirURI != null) {