Limit number of files for a single ExchangeStorageReader

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/FileSystemExchangeConfig.java

@@ -42,6 +42,7 @@ public class FileSystemExchangeConfig
     private int exchangeSinkBuffersPerPartition = 2;
     private DataSize exchangeSinkMaxFileSize = DataSize.of(1, GIGABYTE);
     private int exchangeSourceConcurrentReaders = 4;
+    private int exchangeSourceMaxFilesPerReader = 25;
     private int maxOutputPartitionCount = 50;
     private int exchangeFileListingParallelism = 50;
     private DataSize exchangeSourceHandleTargetDataSize = DataSize.of(256, MEGABYTE);
@@ -151,6 +152,19 @@ public FileSystemExchangeConfig setExchangeSourceConcurrentReaders(int exchangeS
         return this;
     }
 
+    @Min(1)
+    public int getExchangeSourceMaxFilesPerReader()
+    {
+        return exchangeSourceMaxFilesPerReader;
+    }
+
+    @Config("exchange.source-max-files-per-reader")
+    public FileSystemExchangeConfig setExchangeSourceMaxFilesPerReader(int exchangeSourceMaxFilesPerReader)
+    {
+        this.exchangeSourceMaxFilesPerReader = exchangeSourceMaxFilesPerReader;
+        return this;
+    }
+
     @Min(1)
     public int getMaxOutputPartitionCount()
     {

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/FileSystemExchangeManager.java

@@ -57,6 +57,7 @@ public class FileSystemExchangeManager
     private final int exchangeSinkBuffersPerPartition;
     private final long exchangeSinkMaxFileSizeInBytes;
     private final int exchangeSourceConcurrentReaders;
+    private final int exchangeSourceMaxFilesPerReader;
     private final int maxOutputPartitionCount;
     private final int exchangeFileListingParallelism;
     private final long exchangeSourceHandleTargetDataSizeInBytes;
@@ -77,6 +78,7 @@ public FileSystemExchangeManager(
         this.exchangeSinkBuffersPerPartition = fileSystemExchangeConfig.getExchangeSinkBuffersPerPartition();
         this.exchangeSinkMaxFileSizeInBytes = fileSystemExchangeConfig.getExchangeSinkMaxFileSize().toBytes();
         this.exchangeSourceConcurrentReaders = fileSystemExchangeConfig.getExchangeSourceConcurrentReaders();
+        this.exchangeSourceMaxFilesPerReader = fileSystemExchangeConfig.getExchangeSourceMaxFilesPerReader();
         this.maxOutputPartitionCount = fileSystemExchangeConfig.getMaxOutputPartitionCount();
         this.exchangeFileListingParallelism = fileSystemExchangeConfig.getExchangeFileListingParallelism();
         this.exchangeSourceHandleTargetDataSizeInBytes = fileSystemExchangeConfig.getExchangeSourceHandleTargetDataSize().toBytes();
@@ -140,6 +142,7 @@ public ExchangeSource createSource()
                 exchangeStorage,
                 stats,
                 maxPageStorageSizeInBytes,
-                exchangeSourceConcurrentReaders);
+                exchangeSourceConcurrentReaders,
+                exchangeSourceMaxFilesPerReader);
     }
 }

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/FileSystemExchangeSource.java

@@ -57,6 +57,7 @@ public class FileSystemExchangeSource
     private final FileSystemExchangeStats stats;
     private final int maxPageStorageSize;
     private final int exchangeSourceConcurrentReaders;
+    private final int maxFilesPerReader;
 
     private final Queue<ExchangeSourceFile> files = new ConcurrentLinkedQueue<>();
     @GuardedBy("this")
@@ -74,12 +75,14 @@ public FileSystemExchangeSource(
             FileSystemExchangeStorage exchangeStorage,
             FileSystemExchangeStats stats,
             int maxPageStorageSize,
-            int exchangeSourceConcurrentReaders)
+            int exchangeSourceConcurrentReaders,
+            int maxFilesPerReader)
     {
         this.exchangeStorage = requireNonNull(exchangeStorage, "exchangeStorage is null");
         this.stats = requireNonNull(stats, "stats is null");
         this.maxPageStorageSize = maxPageStorageSize;
         this.exchangeSourceConcurrentReaders = exchangeSourceConcurrentReaders;
+        this.maxFilesPerReader = maxFilesPerReader;
     }
 
     @Override
@@ -272,6 +275,7 @@ private void closeAndCreateReadersIfNecessary()
             try {
                 while (activeReaders.size() < exchangeSourceConcurrentReaders && !files.isEmpty()) {
                     ImmutableList.Builder<ExchangeSourceFile> readerFiles = ImmutableList.builder();
+                    int readerFileCount = 0;
                     long readerFileSize = 0;
                     while (!files.isEmpty()) {
                         ExchangeSourceFile file = files.peek();
@@ -280,9 +284,10 @@ private void closeAndCreateReadersIfNecessary()
                                 file.getExchangeId(),
                                 file.getSourceTaskPartitionId(),
                                 file.getSourceTaskAttemptId());
-                        if (readerFileSize == 0 || readerFileSize + file.getFileSize() <= maxPageStorageSize + exchangeStorage.getWriteBufferSize()) {
+                        if (readerFileCount == 0 || ((readerFileSize + file.getFileSize() <= maxPageStorageSize + exchangeStorage.getWriteBufferSize()) && readerFileCount < maxFilesPerReader)) {
                             readerFiles.add(file);
                             readerFileSize += file.getFileSize();
+                            readerFileCount++;
                             files.poll();
                         }
                         else {

plugin/trino-exchange-filesystem/src/test/java/io/trino/plugin/exchange/filesystem/TestFileSystemExchangeConfig.java

@@ -38,6 +38,7 @@ public void testDefaults()
                 .setExchangeSinkBuffersPerPartition(2)
                 .setExchangeSinkMaxFileSize(DataSize.of(1, GIGABYTE))
                 .setExchangeSourceConcurrentReaders(4)
+                .setExchangeSourceMaxFilesPerReader(25)
                 .setMaxOutputPartitionCount(50)
                 .setExchangeFileListingParallelism(50)
                 .setExchangeSourceHandleTargetDataSize(DataSize.of(256, MEGABYTE)));
@@ -54,6 +55,7 @@ public void testExplicitPropertyMappings()
                 .put("exchange.sink-buffers-per-partition", "3")
                 .put("exchange.sink-max-file-size", "2GB")
                 .put("exchange.source-concurrent-readers", "10")
+                .put("exchange.source-max-files-per-reader", "111")
                 .put("exchange.max-output-partition-count", "53")
                 .put("exchange.file-listing-parallelism", "20")
                 .put("exchange.source-handle-target-data-size", "1GB")
@@ -67,6 +69,7 @@ public void testExplicitPropertyMappings()
                 .setExchangeSinkBuffersPerPartition(3)
                 .setExchangeSinkMaxFileSize(DataSize.of(2, GIGABYTE))
                 .setExchangeSourceConcurrentReaders(10)
+                .setExchangeSourceMaxFilesPerReader(111)
                 .setMaxOutputPartitionCount(53)
                 .setExchangeFileListingParallelism(20)
                 .setExchangeSourceHandleTargetDataSize(DataSize.of(1, GIGABYTE));

plugin/trino-exchange-filesystem/src/test/java/io/trino/plugin/exchange/filesystem/TestFileSystemExchangeSource.java

@@ -29,7 +29,8 @@ public void testIsBlockedNonCancellable()
                 new LocalFileSystemExchangeStorage(),
                 new FileSystemExchangeStats(),
                 1024,
-                2)) {
+                2,
+                1)) {
             CompletableFuture<Void> first = source.isBlocked();
             CompletableFuture<Void> second = source.isBlocked();
             assertThat(first)