Parallelize segment staging copy to deep store

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 <overwrite> 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.
+   * <p>
+   * 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 <sourceDir> to the <destDir> 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.
+   * <p>
+   * 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 &lt;= 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<URI> 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<Future<Void>> 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<Void> 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<String, String> 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<URI> listSourceFiles(PinotFS fs, URI sourceDir)
+      throws IOException, URISyntaxException {
+    List<URI> 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);
     }
   }
 }

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<String> 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);
+    }
+  }
 }

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);