[core] optimize read append table with limit (#6848)

Author: wwj6591812

paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreScan.java

@@ -92,6 +92,7 @@ public abstract class AbstractFileStoreScan implements FileStoreScan {
     private ScanMetrics scanMetrics = null;
     private boolean dropStats;
     @Nullable protected List<Range> rowRanges;
+    @Nullable protected Long limit;
 
     public AbstractFileStoreScan(
             ManifestsReader manifestsReader,
@@ -251,6 +252,12 @@ public FileStoreScan withReadType(RowType readType) {
         return this;
     }
 
+    @Override
+    public FileStoreScan withLimit(long limit) {
+        this.limit = limit;
+        return this;
+    }
+
     @Nullable
     @Override
     public Integer parallelism() {
@@ -379,7 +386,7 @@ public Iterator<ManifestEntry> readFileIterator() {
         return readManifestEntries(readManifests().filteredManifests, true);
     }
 
-    private Iterator<ManifestEntry> readManifestEntries(
+    protected Iterator<ManifestEntry> readManifestEntries(
             List<ManifestFileMeta> manifests, boolean useSequential) {
         return scanMode == ScanMode.ALL
                 ? readAndMergeFileEntries(manifests, Function.identity(), useSequential)

paimon-core/src/main/java/org/apache/paimon/operation/AppendOnlyFileStoreScan.java

@@ -22,6 +22,7 @@
 import org.apache.paimon.fileindex.FileIndexPredicate;
 import org.apache.paimon.manifest.ManifestEntry;
 import org.apache.paimon.manifest.ManifestFile;
+import org.apache.paimon.manifest.ManifestFileMeta;
 import org.apache.paimon.predicate.Predicate;
 import org.apache.paimon.schema.SchemaManager;
 import org.apache.paimon.schema.TableSchema;
@@ -33,6 +34,8 @@
 import javax.annotation.Nullable;
 
 import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -80,6 +83,18 @@ public AppendOnlyFileStoreScan withFilter(Predicate predicate) {
         return this;
     }
 
+    @Override
+    protected Iterator<ManifestEntry> readManifestEntries(
+            List<ManifestFileMeta> manifests, boolean useSequential) {
+        Iterator<ManifestEntry> baseIterator = super.readManifestEntries(manifests, useSequential);
+
+        if (limit != null && limit > 0) {
+            return new LimitAwareManifestEntryIterator(baseIterator, limit);
+        }
+
+        return baseIterator;
+    }
+
     /** Note: Keep this thread-safe. */
     @Override
     protected boolean filterByStats(ManifestEntry entry) {
@@ -142,4 +157,64 @@ private boolean testFileIndex(@Nullable byte[] embeddedIndexBytes, ManifestEntry
             throw new RuntimeException("Exception happens while checking predicate.", e);
         }
     }
+
+    /**
+     * Iterator that applies limit pushdown by stopping early when enough rows have been
+     * accumulated.
+     */
+    private static class LimitAwareManifestEntryIterator implements Iterator<ManifestEntry> {
+        private final Iterator<ManifestEntry> baseIterator;
+        private final long limit;
+
+        private long accumulatedRowCount = 0;
+        private ManifestEntry nextEntry = null;
+        private boolean hasNext = false;
+
+        LimitAwareManifestEntryIterator(Iterator<ManifestEntry> baseIterator, long limit) {
+            this.baseIterator = baseIterator;
+            this.limit = limit;
+            advance();
+        }
+
+        private void advance() {
+            // If we've already accumulated enough rows, stop reading more entries
+            if (accumulatedRowCount >= limit) {
+                hasNext = false;
+                nextEntry = null;
+                return;
+            }
+
+            if (baseIterator.hasNext()) {
+                nextEntry = baseIterator.next();
+                hasNext = true;
+
+                long fileRowCount = nextEntry.file().rowCount();
+                if (fileRowCount > 0) {
+                    accumulatedRowCount += fileRowCount;
+                }
+
+                return;
+            }
+
+            // No more base entries
+            hasNext = false;
+            nextEntry = null;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return hasNext;
+        }
+
+        @Override
+        public ManifestEntry next() {
+            // This exception is only thrown if next() is called when hasNext() returns false.
+            if (!hasNext) {
+                throw new java.util.NoSuchElementException();
+            }
+            ManifestEntry current = nextEntry;
+            advance();
+            return current;
+        }
+    }
 }

paimon-core/src/main/java/org/apache/paimon/operation/FileStoreScan.java

@@ -95,6 +95,8 @@ public interface FileStoreScan {
 
     FileStoreScan withReadType(RowType readType);
 
+    FileStoreScan withLimit(long limit);
+
     @Nullable
     Integer parallelism();
 

paimon-core/src/main/java/org/apache/paimon/table/source/DataTableBatchScan.java

@@ -79,6 +79,7 @@ public InnerTableScan withFilter(Predicate predicate) {
     @Override
     public InnerTableScan withLimit(int limit) {
         this.pushDownLimit = limit;
+        snapshotReader.withLimit(limit);
         return this;
     }
 

paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReader.java

@@ -116,6 +116,8 @@ public interface SnapshotReader {
 
     SnapshotReader withReadType(RowType readType);
 
+    SnapshotReader withLimit(int limit);
+
     /** Get splits plan from snapshot. */
     Plan read();
 

paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReaderImpl.java

@@ -322,6 +322,12 @@ public SnapshotReader withDataFileNameFilter(Filter<String> fileNameFilter) {
         return this;
     }
 
+    @Override
+    public SnapshotReader withLimit(int limit) {
+        scan.withLimit(limit);
+        return this;
+    }
+
     @Override
     public SnapshotReader dropStats() {
         scan.dropStats();

paimon-core/src/main/java/org/apache/paimon/table/system/AuditLogTable.java

@@ -438,6 +438,12 @@ public SnapshotReader withReadType(RowType readType) {
             return this;
         }
 
+        @Override
+        public SnapshotReader withLimit(int limit) {
+            wrapped.withLimit(limit);
+            return this;
+        }
+
         @Override
         public Plan read() {
             return wrapped.read();

paimon-core/src/test/java/org/apache/paimon/table/source/TableScanTest.java

@@ -19,6 +19,8 @@
 package org.apache.paimon.table.source;
 
 import org.apache.paimon.predicate.FieldRef;
+import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.predicate.PredicateBuilder;
 import org.apache.paimon.predicate.TopN;
 import org.apache.paimon.stats.SimpleStatsEvolutions;
 import org.apache.paimon.table.FileStoreTable;
@@ -111,6 +113,167 @@ public void testPushDownLimit() throws Exception {
         commit.close();
     }
 
+    @Test
+    public void testLimitPushdownWithFilter() throws Exception {
+        createAppendOnlyTable();
+
+        StreamTableWrite write = table.newWrite(commitUser);
+        StreamTableCommit commit = table.newCommit(commitUser);
+
+        // Write 50 files, each with 1 row. Rows 0-24 have 'a' = 10, rows 25-49 have 'a' = 20.
+        for (int i = 0; i < 25; i++) {
+            write.write(rowData(i, 10, (long) i * 100));
+            commit.commit(i, write.prepareCommit(true, i));
+        }
+        for (int i = 25; i < 50; i++) {
+            write.write(rowData(i, 20, (long) i * 100));
+            commit.commit(i, write.prepareCommit(true, i));
+        }
+
+        // Without limit, should read all 50 files
+        TableScan.Plan planWithoutLimit = table.newScan().plan();
+        int totalSplits = planWithoutLimit.splits().size();
+        assertThat(totalSplits).isEqualTo(50);
+
+        // With filter (a = 20) and limit (10)
+        // filterByStats has already been applied in baseIterator, so only files 25-49 will be
+        // returned
+        // To get 10 rows, it should read 10 files (from index 25 to 34)
+        Predicate filter =
+                new PredicateBuilder(table.schema().logicalRowType())
+                        .equal(1, 20); // Filter on 'a' = 20
+        TableScan.Plan planWithFilterAndLimit =
+                table.newScan().withFilter(filter).withLimit(10).plan();
+        int splitsWithFilterAndLimit = planWithFilterAndLimit.splits().size();
+
+        // Should read exactly 10 files (from index 25 to 34) to get 10 rows
+        assertThat(splitsWithFilterAndLimit).isLessThanOrEqualTo(10);
+        assertThat(splitsWithFilterAndLimit).isGreaterThan(0);
+        assertThat(splitsWithFilterAndLimit).isLessThan(totalSplits);
+
+        write.close();
+        commit.close();
+    }
+
+    @Test
+    public void testLimitPushdownWhenDataLessThanLimit() throws Exception {
+        createAppendOnlyTable();
+
+        StreamTableWrite write = table.newWrite(commitUser);
+        StreamTableCommit commit = table.newCommit(commitUser);
+
+        // Write only 3 rows
+        write.write(rowData(1, 10, 100L));
+        write.write(rowData(2, 20, 200L));
+        write.write(rowData(3, 30, 300L));
+        commit.commit(0, write.prepareCommit(true, 0));
+
+        // With limit 10, but only 3 rows exist
+        // Should return all 3 rows without throwing exception
+        TableScan.Plan plan = table.newScan().withLimit(10).plan();
+        assertThat(plan.splits().size()).isEqualTo(3);
+
+        // Verify we can read all data correctly
+        List<String> result = getResult(table.newRead(), plan.splits());
+        assertThat(result.size()).isEqualTo(3);
+        assertThat(result).containsExactlyInAnyOrder("+I 1|10|100", "+I 2|20|200", "+I 3|30|300");
+
+        write.close();
+        commit.close();
+    }
+
+    @Test
+    public void testLimitPushdownWithFilterWhenDataLessThanLimit() throws Exception {
+        createAppendOnlyTable();
+
+        StreamTableWrite write = table.newWrite(commitUser);
+        StreamTableCommit commit = table.newCommit(commitUser);
+
+        // Write 10 rows, but only 3 rows have 'a' = 20
+        for (int i = 0; i < 7; i++) {
+            write.write(rowData(i, 10, (long) i * 100));
+            commit.commit(i, write.prepareCommit(true, i));
+        }
+        for (int i = 7; i < 10; i++) {
+            write.write(rowData(i, 20, (long) i * 100));
+            commit.commit(i, write.prepareCommit(true, i));
+        }
+
+        // With filter (a = 20) and limit (10), but only 3 rows match
+        // Should return all 3 matching rows without throwing exception
+        Predicate filter = new PredicateBuilder(table.schema().logicalRowType()).equal(1, 20);
+        TableScan.Plan plan = table.newScan().withFilter(filter).withLimit(10).plan();
+
+        // Should read all 3 files that match the filter
+        assertThat(plan.splits().size()).isEqualTo(3);
+
+        // Verify we can read all matching data correctly
+        List<String> result = getResult(table.newRead(), plan.splits());
+        assertThat(result.size()).isEqualTo(3);
+        assertThat(result).containsExactlyInAnyOrder("+I 7|20|700", "+I 8|20|800", "+I 9|20|900");
+
+        write.close();
+        commit.close();
+    }
+
+    @Test
+    public void testLimitPushdownEarlyStop() throws Exception {
+        createAppendOnlyTable();
+
+        StreamTableWrite write = table.newWrite(commitUser);
+        StreamTableCommit commit = table.newCommit(commitUser);
+
+        // Write 100 files, each with 1 row
+        for (int i = 0; i < 100; i++) {
+            write.write(rowData(i, i, (long) i * 100));
+            commit.commit(i, write.prepareCommit(true, i));
+        }
+
+        // Without limit, should read all 100 files
+        TableScan.Plan planWithoutLimit = table.newScan().plan();
+        assertThat(planWithoutLimit.splits().size()).isEqualTo(100);
+
+        // With limit 10, should only read 10 files (early stop optimization)
+        TableScan.Plan planWithLimit = table.newScan().withLimit(10).plan();
+        assertThat(planWithLimit.splits().size()).isEqualTo(10);
+
+        // Verify the returned data is correct
+        List<String> result = getResult(table.newRead(), planWithLimit.splits());
+        assertThat(result.size()).isEqualTo(10);
+
+        write.close();
+        commit.close();
+    }
+
+    @Test
+    public void testLimitPushdownBoundaryCases() throws Exception {
+        createAppendOnlyTable();
+
+        StreamTableWrite write = table.newWrite(commitUser);
+        StreamTableCommit commit = table.newCommit(commitUser);
+
+        // Write 5 rows
+        for (int i = 0; i < 5; i++) {
+            write.write(rowData(i, i, (long) i * 100));
+            commit.commit(i, write.prepareCommit(true, i));
+        }
+
+        // Test limit = 1
+        TableScan.Plan plan1 = table.newScan().withLimit(1).plan();
+        assertThat(plan1.splits().size()).isEqualTo(1);
+
+        // Test limit = 5 (exactly the number of rows)
+        TableScan.Plan plan5 = table.newScan().withLimit(5).plan();
+        assertThat(plan5.splits().size()).isEqualTo(5);
+
+        // Test limit = 10 (more than the number of rows, should return all 5)
+        TableScan.Plan plan10 = table.newScan().withLimit(10).plan();
+        assertThat(plan10.splits().size()).isEqualTo(5);
+
+        write.close();
+        commit.close();
+    }
+
     @Test
     public void testPushDownTopN() throws Exception {
         createAppendOnlyTable();