Improve HashBuilderOperator unspill parallelism

sopel39 · sopel39 · commit 04aceb00b82f · 2025-10-14T18:33:03.000+02:00
HashBuilderOperator is unspilling sequentially partition by partition.
This commit improves join unspilling performance by making FileSingleStreamSpiller
unspill single partition in parallel. FileSingleStreamSpiller is enhanced
so it can spill to/unspill from multiple files.
diff --git a/core/trino-main/src/main/java/io/trino/operator/join/HashBuilderOperator.java b/core/trino-main/src/main/java/io/trino/operator/join/HashBuilderOperator.java
@@ -411,7 +411,8 @@ private ListenableFuture<Void> spillIndex()
         spiller = Optional.of(singleStreamSpillerFactory.create(
                 index.getTypes(),
                 operatorContext.getSpillContext().newLocalSpillContext(),
-                operatorContext.newLocalUserMemoryContext(HashBuilderOperator.class.getSimpleName())));
+                operatorContext.newLocalUserMemoryContext(HashBuilderOperator.class.getSimpleName()),
+                true));
         long spillStartNanos = System.nanoTime();
         ListenableFuture<DataSize> spillFuture = getSpiller().spill(index.getPages());
         addSuccessCallback(spillFuture, dataSize -> {
diff --git a/core/trino-main/src/main/java/io/trino/spiller/FileSingleStreamSpiller.java b/core/trino-main/src/main/java/io/trino/spiller/FileSingleStreamSpiller.java
@@ -42,13 +42,15 @@
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.Iterators.transform;
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static io.trino.spiller.FileSingleStreamSpillerFactory.SPILL_FILE_PREFIX;
@@ -63,7 +65,9 @@ public class FileSingleStreamSpiller
     @VisibleForTesting
     static final int BUFFER_SIZE = 4 * 1024;
 
-    private final FileHolder targetFile;
+    private final List<FileHolder> targetFiles;
+    private volatile int currentFileIndex;
+
     private final Closer closer = Closer.create();
     private final PagesSerdeFactory serdeFactory;
     private volatile Optional<SecretKey> encryptionKey;
@@ -84,7 +88,7 @@ public FileSingleStreamSpiller(
             PagesSerdeFactory serdeFactory,
             Optional<SecretKey> encryptionKey,
             ListeningExecutorService executor,
-            Path spillPath,
+            List<Path> spillPaths,
             SpillerStats spillerStats,
             SpillContext spillContext,
             LocalMemoryContext memoryContext,
@@ -109,8 +113,14 @@ public FileSingleStreamSpiller(
         // middle of execution when close() is called (note that this applies to both readPages() and writePages() methods).
         this.memoryContext.setBytes(BUFFER_SIZE);
         this.fileSystemErrorHandler = requireNonNull(fileSystemErrorHandler, "filesystemErrorHandler is null");
+        requireNonNull(spillPaths, "spillPaths is null");
+        checkState(!spillPaths.isEmpty(), "spillPaths is empty");
         try {
-            this.targetFile = closer.register(new FileHolder(Files.createTempFile(spillPath, SPILL_FILE_PREFIX, SPILL_FILE_SUFFIX)));
+            ImmutableList.Builder<FileHolder> builder = ImmutableList.builder();
+            for (Path path : spillPaths) {
+                builder.add(closer.register(new FileHolder(Files.createTempFile(path, SPILL_FILE_PREFIX, SPILL_FILE_SUFFIX))));
+            }
+            this.targetFiles = builder.build();
         }
         catch (IOException e) {
             this.fileSystemErrorHandler.run();
@@ -137,61 +147,129 @@ public long getSpilledPagesInMemorySize()
     public Iterator<Page> getSpilledPages()
     {
         checkNoSpillInProgress();
-        return readPages();
+        checkState(writable.getAndSet(false), "Repeated reads are disallowed to prevent potential resource leaks");
+
+        try {
+            Optional<SecretKey> encryptionKey = this.encryptionKey;
+            checkState(encrypted == encryptionKey.isPresent(), "encryptionKey has been discarded");
+            PageDeserializer deserializer = serdeFactory.createDeserializer(encryptionKey);
+            this.encryptionKey = Optional.empty();
+
+            int fileCount = targetFiles.size();
+            List<Iterator<Page>> iterators = new ArrayList<>(targetFiles.size());
+            for (FileHolder file : targetFiles) {
+                iterators.add(readFilePages(deserializer, file, closer));
+            }
+
+            return new AbstractIterator<>()
+            {
+                int fileIndex;
+
+                @Override
+                protected Page computeNext()
+                {
+                    Iterator<Page> iterator = iterators.get(fileIndex);
+                    if (!iterator.hasNext()) {
+                        return endOfData();
+                    }
+
+                    Page page = iterator.next();
+                    fileIndex = (fileIndex + 1) % fileCount;
+                    return page;
+                }
+            };
+        }
+        catch (IOException e) {
+            fileSystemErrorHandler.run();
+            throw new TrinoException(GENERIC_INTERNAL_ERROR, "Failed to read spilled pages", e);
+        }
     }
 
     @Override
     public ListenableFuture<List<Page>> getAllSpilledPages()
     {
-        return executor.submit(() -> ImmutableList.copyOf(getSpilledPages()));
+        checkNoSpillInProgress();
+        checkState(writable.getAndSet(false), "Repeated reads are disallowed to prevent potential resource leaks");
+
+        Optional<SecretKey> encryptionKey = this.encryptionKey;
+        checkState(encrypted == encryptionKey.isPresent(), "encryptionKey has been discarded");
+        this.encryptionKey = Optional.empty();
+
+        List<ListenableFuture<List<Page>>> futures = new ArrayList<>();
+        for (FileHolder file : targetFiles) {
+            futures.add(executor.submit(() -> {
+                PageDeserializer deserializer = serdeFactory.createDeserializer(encryptionKey);
+                ImmutableList.Builder<Page> pages = ImmutableList.builder();
+                try (Closer closer = Closer.create()) {
+                    readFilePages(deserializer, file, closer).forEachRemaining(pages::add);
+                }
+                return pages.build();
+            }));
+        }
+
+        // Combine pages from all spill files according to the round-robin order.
+        return Futures.transform(Futures.allAsList(futures), pagesPerFile -> {
+            ImmutableList.Builder<Page> builder = ImmutableList.builder();
+            int fileCount = targetFiles.size();
+
+            List<Iterator<Page>> iterators = new ArrayList<>(fileCount);
+            for (List<Page> pages : pagesPerFile) {
+                iterators.add(pages.iterator());
+            }
+
+            int fileIndex = 0;
+            while (true) {
+                Iterator<Page> iterator = iterators.get(fileIndex);
+                if (!iterator.hasNext()) {
+                    break;
+                }
+                builder.add(iterator.next());
+                fileIndex = (fileIndex + 1) % fileCount;
+            }
+            return builder.build();
+        }, executor);
     }
 
-    private DataSize writePages(Iterator<Page> pageIterator)
+    private DataSize writePages(Iterator<Page> pages)
     {
         checkState(writable.get(), "Spilling no longer allowed. The spiller has been made non-writable on first read for subsequent reads to be consistent");
 
         Optional<SecretKey> encryptionKey = this.encryptionKey;
         checkState(encrypted == encryptionKey.isPresent(), "encryptionKey has been discarded");
         PageSerializer serializer = serdeFactory.createSerializer(encryptionKey);
+
         long spilledPagesBytes = 0;
-        try (SliceOutput output = new OutputStreamSliceOutput(targetFile.newOutputStream(APPEND), BUFFER_SIZE)) {
-            while (pageIterator.hasNext()) {
-                Page page = pageIterator.next();
+        int fileIndex = currentFileIndex;
+        int fileCount = targetFiles.size();
+
+        try {
+            while (pages.hasNext()) {
+                Page page = pages.next();
                 long pageSizeInBytes = page.getSizeInBytes();
+                Slice serialized = serializer.serialize(page);
+                long serializedPageSize = serialized.length();
+
+                try (SliceOutput out = newSliceOutput(fileIndex)) {
+                    out.writeBytes(serialized);
+                }
+
                 spilledPagesBytes += pageSizeInBytes;
+
                 spilledPagesInMemorySize.addAndGet(pageSizeInBytes);
-                Slice serializedPage = serializer.serialize(page);
-                long pageSize = serializedPage.length();
-                localSpillContext.updateBytes(pageSize);
-                spillerStats.addToTotalSpilledBytes(pageSize);
-                output.writeBytes(serializedPage);
+                localSpillContext.updateBytes(serializedPageSize);
+                spillerStats.addToTotalSpilledBytes(serializedPageSize);
+
+                fileIndex = (fileIndex + 1) % fileCount;
             }
+
+            currentFileIndex = fileIndex;
         }
         catch (UncheckedIOException | IOException e) {
             fileSystemErrorHandler.run();
             throw new TrinoException(GENERIC_INTERNAL_ERROR, "Failed to spill pages", e);
         }
-        return DataSize.ofBytes(spilledPagesBytes);
-    }
 
-    private Iterator<Page> readPages()
-    {
-        checkState(writable.getAndSet(false), "Repeated reads are disallowed to prevent potential resource leaks");
-
-        try {
-            Optional<SecretKey> encryptionKey = this.encryptionKey;
-            checkState(encrypted == encryptionKey.isPresent(), "encryptionKey has been discarded");
-            PageDeserializer deserializer = serdeFactory.createDeserializer(encryptionKey);
-            // encryption key is safe to discard since it now belongs to the PageDeserializer and repeated reads are disallowed
-            this.encryptionKey = Optional.empty();
-            InputStream input = closer.register(targetFile.newInputStream());
-            Iterator<Page> pages = PagesSerdeUtil.readPages(deserializer, input);
-            return closeWhenExhausted(pages, input);
-        }
-        catch (IOException e) {
-            fileSystemErrorHandler.run();
-            throw new TrinoException(GENERIC_INTERNAL_ERROR, "Failed to read spilled pages", e);
-        }
+        return DataSize.ofBytes(spilledPagesBytes);
     }
 
     @Override
@@ -215,6 +293,23 @@ private void checkNoSpillInProgress()
         checkState(spillInProgress.isDone(), "spill in progress");
     }
 
+    private SliceOutput newSliceOutput(int fileIndex)
+            throws IOException
+    {
+        return new OutputStreamSliceOutput(targetFiles.get(fileIndex).newOutputStream(APPEND), BUFFER_SIZE);
+    }
+
+    /**
+     * Returns an iterator that exposes all pages stored in the given file.
+     * Pages are lazily deserialized as the iterator is consumed.
+     */
+    private Iterator<Page> readFilePages(PageDeserializer deserializer, FileHolder file, Closer closer)
+            throws IOException
+    {
+        InputStream input = closer.register(file.newInputStream());
+        return transform(closeWhenExhausted(PagesSerdeUtil.readSerializedPages(input), input), deserializer::deserialize);
+    }
+
     private static <T> Iterator<T> closeWhenExhausted(Iterator<T> iterator, Closeable resource)
     {
         requireNonNull(iterator, "iterator is null");
diff --git a/core/trino-main/src/main/java/io/trino/spiller/FileSingleStreamSpillerFactory.java b/core/trino-main/src/main/java/io/trino/spiller/FileSingleStreamSpillerFactory.java
@@ -22,6 +22,7 @@
 import io.airlift.log.Logger;
 import io.trino.FeaturesConfig;
 import io.trino.cache.NonKeyEvictableLoadingCache;
+import io.trino.execution.TaskManagerConfig;
 import io.trino.execution.buffer.CompressionCodec;
 import io.trino.execution.buffer.PagesSerdeFactory;
 import io.trino.memory.context.LocalMemoryContext;
@@ -80,11 +81,12 @@ public class FileSingleStreamSpillerFactory
     private final SpillerStats spillerStats;
     private final double maxUsedSpaceThreshold;
     private final boolean spillEncryptionEnabled;
+    private final int spillFileCount;
     private int roundRobinIndex;
     private final NonKeyEvictableLoadingCache<Path, Boolean> spillPathHealthCache;
 
     @Inject
-    public FileSingleStreamSpillerFactory(BlockEncodingSerde blockEncodingSerde, SpillerStats spillerStats, FeaturesConfig featuresConfig, NodeSpillConfig nodeSpillConfig)
+    public FileSingleStreamSpillerFactory(BlockEncodingSerde blockEncodingSerde, SpillerStats spillerStats, FeaturesConfig featuresConfig, NodeSpillConfig nodeSpillConfig, TaskManagerConfig taskManagerConfig)
     {
         this(
                 listeningDecorator(newFixedThreadPool(
@@ -93,6 +95,7 @@ public FileSingleStreamSpillerFactory(BlockEncodingSerde blockEncodingSerde, Spi
                 requireNonNull(blockEncodingSerde, "blockEncodingSerde is null"),
                 spillerStats,
                 featuresConfig.getSpillerSpillPaths(),
+                Math.min(featuresConfig.getSpillerThreads(), taskManagerConfig.getTaskConcurrency()),
                 featuresConfig.getSpillMaxUsedSpaceThreshold(),
                 nodeSpillConfig.getSpillCompressionCodec(),
                 nodeSpillConfig.isSpillEncryptionEnabled());
@@ -104,6 +107,7 @@ public FileSingleStreamSpillerFactory(
             BlockEncodingSerde blockEncodingSerde,
             SpillerStats spillerStats,
             List<Path> spillPaths,
+            int spillFileCount,
             double maxUsedSpaceThreshold,
             CompressionCodec compressionCodec,
             boolean spillEncryptionEnabled)
@@ -124,6 +128,7 @@ public FileSingleStreamSpillerFactory(
                 throw new IllegalArgumentException(format("spill path %s is not accessible, it must be +rwx; adjust %s config property or filesystem permissions", path, SPILLER_SPILL_PATH));
             }
         });
+        this.spillFileCount = spillFileCount;
         this.maxUsedSpaceThreshold = maxUsedSpaceThreshold;
         this.spillEncryptionEnabled = spillEncryptionEnabled;
         this.roundRobinIndex = 0;
@@ -165,14 +170,19 @@ private static void cleanupOldSpillFiles(Path path)
     }
 
     @Override
-    public SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext)
+    public SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext, boolean parallelSpill)
     {
         Optional<SecretKey> encryptionKey = spillEncryptionEnabled ? Optional.of(createRandomAesEncryptionKey()) : Optional.empty();
+        ImmutableList.Builder<Path> paths = ImmutableList.builder();
+        int spillFileCount = parallelSpill ? this.spillFileCount : 1;
+        for (int i = 0; i < spillFileCount; i++) {
+            paths.add(getNextSpillPath());
+        }
         return new FileSingleStreamSpiller(
                 serdeFactory,
                 encryptionKey,
                 executor,
-                getNextSpillPath(),
+                paths.build(),
                 spillerStats,
                 spillContext,
                 memoryContext,
diff --git a/core/trino-main/src/main/java/io/trino/spiller/SingleStreamSpillerFactory.java b/core/trino-main/src/main/java/io/trino/spiller/SingleStreamSpillerFactory.java
@@ -21,11 +21,16 @@
 
 public interface SingleStreamSpillerFactory
 {
-    SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext);
+    default SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext)
+    {
+        return create(types, spillContext, memoryContext, false);
+    }
+
+    SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext, boolean parallelSpill);
 
     static SingleStreamSpillerFactory unsupportedSingleStreamSpillerFactory()
     {
-        return (types, spillContext, memoryContext) -> {
+        return (types, spillContext, memoryContext, parallelSpill) -> {
             throw new UnsupportedOperationException();
         };
     }
diff --git a/core/trino-main/src/test/java/io/trino/execution/TaskTestUtils.java b/core/trino-main/src/test/java/io/trino/execution/TaskTestUtils.java
@@ -177,12 +177,12 @@ public static LocalExecutionPlanner createTestingPlanner()
                 new JoinFilterFunctionCompiler(PLANNER_CONTEXT.getFunctionManager()),
                 new IndexJoinLookupStats(),
                 new TaskManagerConfig(),
-                new GenericSpillerFactory((types, spillContext, memoryContext) -> {
+                new GenericSpillerFactory((types, spillContext, memoryContext, parallelSpill) -> {
                     throw new UnsupportedOperationException();
                 }),
                 new QueryDataEncoders(new SpoolingEnabledConfig(), Set.of()),
                 Optional.empty(),
-                (types, spillContext, memoryContext) -> {
+                (types, spillContext, memoryContext, parallelSpill) -> {
                     throw new UnsupportedOperationException();
                 },
                 (types, partitionFunction, spillContext, memoryContext) -> {
diff --git a/core/trino-main/src/test/java/io/trino/operator/join/JoinTestUtils.java b/core/trino-main/src/test/java/io/trino/operator/join/JoinTestUtils.java
@@ -317,7 +317,7 @@ public DummySpillerFactory failUnspill()
         }
 
         @Override
-        public SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext)
+        public SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext, boolean parallelSpill)
         {
             return new SingleStreamSpiller()
             {
diff --git a/core/trino-main/src/test/java/io/trino/operator/spiller/BenchmarkBinaryFileSpiller.java b/core/trino-main/src/test/java/io/trino/operator/spiller/BenchmarkBinaryFileSpiller.java
@@ -117,6 +117,7 @@ public void setup()
                     BLOCK_ENCODING_SERDE,
                     spillerStats,
                     ImmutableList.of(SPILL_PATH),
+                    1,
                     1.0,
                     compressionCodec,
                     encryptionEnabled);
diff --git a/core/trino-main/src/test/java/io/trino/spiller/TestBinaryFileSpiller.java b/core/trino-main/src/test/java/io/trino/spiller/TestBinaryFileSpiller.java
@@ -17,6 +17,7 @@
 import io.airlift.slice.Slices;
 import io.trino.FeaturesConfig;
 import io.trino.RowPagesBuilder;
+import io.trino.execution.TaskManagerConfig;
 import io.trino.execution.buffer.PageSerializer;
 import io.trino.execution.buffer.PagesSerdeFactory;
 import io.trino.memory.context.AggregatedMemoryContext;
@@ -82,7 +83,7 @@ public void setUp()
         featuresConfig.setSpillMaxUsedSpaceThreshold(1.0);
         NodeSpillConfig nodeSpillConfig = new NodeSpillConfig();
         BlockEncodingSerde blockEncodingSerde = new TestingBlockEncodingSerde();
-        singleStreamSpillerFactory = new FileSingleStreamSpillerFactory(blockEncodingSerde, spillerStats, featuresConfig, nodeSpillConfig);
+        singleStreamSpillerFactory = new FileSingleStreamSpillerFactory(blockEncodingSerde, spillerStats, featuresConfig, nodeSpillConfig, new TaskManagerConfig());
         factory = new GenericSpillerFactory(singleStreamSpillerFactory);
         PagesSerdeFactory pagesSerdeFactory = createSpillingPagesSerdeFactory(blockEncodingSerde, nodeSpillConfig.getSpillCompressionCodec());
         serializer = pagesSerdeFactory.createSerializer(Optional.empty());
diff --git a/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpiller.java b/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpiller.java
diff --git a/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpillerFactory.java b/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpillerFactory.java
diff --git a/core/trino-main/src/test/java/io/trino/spiller/TestGenericPartitioningSpiller.java b/core/trino-main/src/test/java/io/trino/spiller/TestGenericPartitioningSpiller.java

Original file line number	Diff line number	Diff line change
`@@ -21,11 +21,16 @@`
`21`	`21`
`22`	`22`	`public interface SingleStreamSpillerFactory`
`23`	`23`	`{`
`24`		`- SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext);`
	`24`	`+ default SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext)`
	`25`	`+ {`
	`26`	`+ return create(types, spillContext, memoryContext, false);`
	`27`	`+ }`
	`28`	`+`
	`29`	`+ SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext, boolean parallelSpill);`
`25`	`30`
`26`	`31`	`static SingleStreamSpillerFactory unsupportedSingleStreamSpillerFactory()`
`27`	`32`	`{`
`28`		`- return (types, spillContext, memoryContext) -> {`
	`33`	`+ return (types, spillContext, memoryContext, parallelSpill) -> {`
`29`	`34`	`throw new UnsupportedOperationException();`
`30`	`35`	`};`
`31`	`36`	`}`
Original file line number	Diff line number	Diff line change
`@@ -317,7 +317,7 @@ public DummySpillerFactory failUnspill()`
`317`	`317`	`}`
`318`	`318`
`319`	`319`	`@Override`
`320`		`- public SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext)`
	`320`	`+ public SingleStreamSpiller create(List<Type> types, SpillContext spillContext, LocalMemoryContext memoryContext, boolean parallelSpill)`
`321`	`321`	`{`
`322`	`322`	`return new SingleStreamSpiller()`
`323`	`323`	`{`