Add --source_parallelism flag to run multiple input sources concurrently

planetiler-core/src/main/java/com/onthegomap/planetiler/Planetiler.java

@@ -1,5 +1,7 @@
 package com.onthegomap.planetiler;
 
+import static com.onthegomap.planetiler.worker.Worker.joinFutures;
+
 import com.onthegomap.planetiler.archive.TileArchiveConfig;
 import com.onthegomap.planetiler.archive.TileArchiveMetadata;
 import com.onthegomap.planetiler.archive.TileArchiveWriter;
@@ -39,6 +41,7 @@
 import com.onthegomap.planetiler.util.Wikidata;
 import com.onthegomap.planetiler.validator.JavaProfileValidator;
 import com.onthegomap.planetiler.worker.RunnableThatThrows;
+import com.onthegomap.planetiler.worker.Worker.NamedThreadFactory;
 import java.io.IOException;
 import java.nio.file.FileSystem;
 import java.nio.file.Files;
@@ -49,6 +52,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.function.Function;
@@ -954,11 +958,16 @@
         featureGroup.loadStringEncoders(stringEncoderPath);
         featureGroup.initFromManifest(chunkManifestPath);
       } else {
-        for (Stage stage : stages) {
-          try {
-            stage.task.run();
-          } catch (Exception e) {
-            throw new PlanetilerException("Error occurred during stage " + stage.id, e);
+        int requestedParallelism = config.sourceParallelism();
+        if (requestedParallelism > 1 && stages.size() > 1) {
+          runStagesInParallel(stages, Math.min(requestedParallelism, stages.size()));
+        } else {
+          for (Stage stage : stages) {
+            try {
+              stage.task.run();
+            } catch (Exception e) {
+              throw new PlanetilerException("Error occurred during stage " + stage.id, e);
+            }
           }
         }
 
@@ -1135,6 +1144,27 @@
     }
   }
 
+  private void runStagesInParallel(List<Stage> stages, int parallelism) {
+    LOGGER.info("Running {} source stages with parallelism={}", stages.size(), parallelism);
+    try (var es = Executors.newFixedThreadPool(parallelism, new NamedThreadFactory("stage-runner"))) {
+      joinFutures(stages.stream()
+        .<CompletableFuture<?>>map(stage -> CompletableFuture.runAsync(() -> {
+          try {
+            stage.task.run();
+          } catch (Exception e) {
+            throw new PlanetilerException("Error occurred during stage " + stage.id, e);
+          }
+        }, es))
+        .toList()
+      ).join();
+    } catch (CompletionException ce) {
+      if (ce.getCause() instanceof PlanetilerException pe) {
+        throw pe;
+      }
+      throw ce;
+    }
+  }
+
   private record Stage(String id, List<String> details, RunnableThatThrows task) {
 
     Stage(String id, String description, RunnableThatThrows task) {

planetiler-core/src/main/java/com/onthegomap/planetiler/config/PlanetilerConfig.java

@@ -73,7 +73,8 @@ public record PlanetilerConfig(
   int featureSourceIdMultiplier,
   List<String> extraNameTags,
   boolean reuseFeatureDb,
-  boolean parallelTempIO
+  boolean parallelTempIO,
+  int sourceParallelism
 ) {
 
   public static final int MIN_MINZOOM = 0;
@@ -256,7 +257,11 @@ public static PlanetilerConfig from(Arguments arguments) {
       arguments.getBoolean("reuse_featuredb",
         "Reuse existing feature DB on disk, skipping source reading stages (for iterating on post-processing logic)",
         false),
-      parallelTempIO
+      parallelTempIO,
+      Math.max(1, arguments.getInteger("source_parallelism",
+        "number of input source stages to run in parallel (1 = sequential, current behavior). " +
+          "Useful for schemas with many small sources (e.g. shapefiles)",
+        1))
     );
   }
 

planetiler-core/src/main/java/com/onthegomap/planetiler/worker/Worker.java

@@ -144,13 +144,13 @@ public void await() {
   }
 
   /** A thread factory that prepends {@code name-} to all thread names. */
-  private static class NamedThreadFactory implements ThreadFactory {
+  public static class NamedThreadFactory implements ThreadFactory {
 
     private final ThreadGroup group;
     private final AtomicInteger threadNumber = new AtomicInteger(1);
     private final String namePrefix;
 
-    private NamedThreadFactory(String name) {
+    public NamedThreadFactory(String name) {
       group = Thread.currentThread().getThreadGroup();
       namePrefix = name + "-";
     }

planetiler-core/src/test/java/com/onthegomap/planetiler/PlanetilerTests.java

@@ -2880,12 +2880,17 @@ public void processFeature(SourceFeature source, FeatureCollector features) {
     }
   }
 
-  @Test
-  void testPlanetilerRunnerShapefile() throws Exception {
+  @ParameterizedTest
+  @ValueSource(strings = {
+    "",
+    "--source_parallelism=2",
+    "--source_parallelism=4"
+  })
+  void testPlanetilerRunnerShapefile(String args) throws Exception {
     Path mbtiles = tempDir.resolve("output.mbtiles");
     Path resourceDir = TestUtils.pathToResource("");
 
-    Planetiler.create(Arguments.fromArgs("--tmpdir=" + tempDir.resolve("data")))
+    Planetiler.create(Arguments.fromArgs((args + " --tmpdir=" + tempDir.resolve("data")).split("\\s+")))
       .setProfile(new Profile.NullProfile() {
         @Override
         public void processFeature(SourceFeature source, FeatureCollector features) {

planetiler-core/src/test/java/com/onthegomap/planetiler/config/PlanetilerConfigTest.java

@@ -0,0 +1,24 @@
+package com.onthegomap.planetiler.config;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.Test;
+
+class PlanetilerConfigTest {
+
+  @Test
+  void testSourceParallelismDefaultsToOne() {
+    assertEquals(1, PlanetilerConfig.defaults().sourceParallelism());
+  }
+
+  @Test
+  void testSourceParallelismParsesArgument() {
+    assertEquals(4, PlanetilerConfig.from(Arguments.fromArgs("--source_parallelism=4")).sourceParallelism());
+  }
+
+  @Test
+  void testSourceParallelismClampedToOne() {
+    assertEquals(1, PlanetilerConfig.from(Arguments.fromArgs("--source_parallelism=0")).sourceParallelism());
+    assertEquals(1, PlanetilerConfig.from(Arguments.fromArgs("--source_parallelism=-5")).sourceParallelism());
+  }
+}