diff --git a/paimon-common/src/main/java/org/apache/paimon/reader/LimitRecordReader.java b/paimon-common/src/main/java/org/apache/paimon/reader/LimitRecordReader.java index 256453497f51..2f0766e29a34 100644 --- a/paimon-common/src/main/java/org/apache/paimon/reader/LimitRecordReader.java +++ b/paimon-common/src/main/java/org/apache/paimon/reader/LimitRecordReader.java @@ -18,6 +18,8 @@ package org.apache.paimon.reader; +import org.apache.paimon.fs.Path; + import javax.annotation.Nullable; import java.io.IOException; @@ -59,6 +61,9 @@ public RecordIterator readBatch() throws IOException { if (iterator instanceof ScoreRecordIterator) { return new LimitScoreRecordIterator<>((ScoreRecordIterator) iterator); } + if (iterator instanceof FileRecordIterator) { + return new LimitFileRecordIterator<>((FileRecordIterator) iterator); + } return new LimitRecordIterator<>(iterator); } @@ -114,4 +119,25 @@ public long returnedRowId() { return iterator.returnedRowId(); } } + + private class LimitFileRecordIterator extends LimitRecordIterator + implements FileRecordIterator { + + private final FileRecordIterator iterator; + + private LimitFileRecordIterator(FileRecordIterator iterator) { + super(iterator); + this.iterator = iterator; + } + + @Override + public long returnedPosition() { + return iterator.returnedPosition(); + } + + @Override + public Path filePath() { + return iterator.filePath(); + } + } } diff --git a/paimon-common/src/test/java/org/apache/paimon/reader/LimitRecordReaderTest.java b/paimon-common/src/test/java/org/apache/paimon/reader/LimitRecordReaderTest.java new file mode 100644 index 000000000000..6e5224d98d73 --- /dev/null +++ b/paimon-common/src/test/java/org/apache/paimon/reader/LimitRecordReaderTest.java @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.reader; + +import org.apache.paimon.fs.Path; + +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link LimitRecordReader}. */ +public class LimitRecordReaderTest { + + @Test + public void testPreservesFileRecordIterator() throws Exception { + FileRecordIterator fileIterator = + new FileRecordIterator() { + private int position = -1; + + @Override + public long returnedPosition() { + return position; + } + + @Override + public Path filePath() { + return new Path("test-file.parquet"); + } + + @Nullable + @Override + public Integer next() { + position++; + return position < 3 ? position : null; + } + + @Override + public void releaseBatch() {} + }; + + FileRecordReader fileReader = + new FileRecordReader() { + private boolean batchReturned; + + @Nullable + @Override + public FileRecordIterator readBatch() { + if (batchReturned) { + return null; + } + batchReturned = true; + return fileIterator; + } + + @Override + public void close() {} + }; + + try (RecordReader reader = LimitRecordReader.limit(fileReader, 2)) { + RecordReader.RecordIterator batch = reader.readBatch(); + assertThat(batch).isInstanceOf(FileRecordIterator.class); + + FileRecordIterator limited = (FileRecordIterator) batch; + assertThat(limited.filePath()).isEqualTo(new Path("test-file.parquet")); + assertThat(limited.next()).isEqualTo(0); + assertThat(limited.returnedPosition()).isEqualTo(0); + assertThat(limited.next()).isEqualTo(1); + assertThat(limited.returnedPosition()).isEqualTo(1); + assertThat(limited.next()).isNull(); + limited.releaseBatch(); + + assertThat(reader.readBatch()).isNull(); + } + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/io/FileIndexEvaluator.java b/paimon-core/src/main/java/org/apache/paimon/io/FileIndexEvaluator.java index 1a7862271770..a82194cc7616 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/FileIndexEvaluator.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/FileIndexEvaluator.java @@ -18,6 +18,7 @@ package org.apache.paimon.io; +import org.apache.paimon.deletionvectors.Bitmap64DeletionVector; import org.apache.paimon.deletionvectors.BitmapDeletionVector; import org.apache.paimon.deletionvectors.DeletionVector; import org.apache.paimon.fileindex.FileIndexPredicate; @@ -52,6 +53,12 @@ public static FileIndexResult evaluate( DataFileMeta file, @Nullable DeletionVector dv) throws IOException { + // File index selections use 32-bit positions. Fall back when they cannot safely represent + // the file or its deletion vector. + if (file.rowCount() > RoaringBitmap32.MAX_VALUE || dv instanceof Bitmap64DeletionVector) { + return FileIndexResult.REMAIN; + } + if (isNullOrEmpty(dataFilter) && topN == null) { if (limit == null) { return FileIndexResult.REMAIN; diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java index 87f7063a81c0..63d0ab884f9f 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java @@ -41,6 +41,7 @@ import org.apache.paimon.predicate.TopN; import org.apache.paimon.reader.EmptyFileRecordReader; import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.LimitRecordReader; import org.apache.paimon.reader.ReaderSupplier; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.SchemaManager; @@ -213,7 +214,12 @@ public RecordReader createReader( null)); } - return ConcatRecordReader.create(suppliers); + RecordReader reader = ConcatRecordReader.create(suppliers); + // Apply the final limit after deletion vectors when no later predicate can drop rows. + if (topN == null && (filters == null || filters.isEmpty())) { + return LimitRecordReader.limit(reader, limit); + } + return reader; } FileRecordReader createFileReader( diff --git a/paimon-core/src/test/java/org/apache/paimon/io/FileIndexEvaluatorTest.java b/paimon-core/src/test/java/org/apache/paimon/io/FileIndexEvaluatorTest.java index 5b91c648134a..891ae8380a8c 100644 --- a/paimon-core/src/test/java/org/apache/paimon/io/FileIndexEvaluatorTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/io/FileIndexEvaluatorTest.java @@ -18,12 +18,14 @@ package org.apache.paimon.io; +import org.apache.paimon.deletionvectors.Bitmap64DeletionVector; import org.apache.paimon.deletionvectors.BitmapDeletionVector; import org.apache.paimon.fileindex.FileIndexFormat; import org.apache.paimon.fileindex.FileIndexResult; import org.apache.paimon.fileindex.FileIndexWriter; import org.apache.paimon.fileindex.bitmap.BitmapFileIndex; import org.apache.paimon.fileindex.bitmap.BitmapFileIndexFactory; +import org.apache.paimon.fileindex.bitmap.BitmapIndexResult; import org.apache.paimon.options.Options; import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.schema.TableSchema; @@ -32,6 +34,7 @@ import org.apache.paimon.types.DataType; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.RoaringBitmap32; import org.junit.jupiter.api.Test; @@ -46,6 +49,64 @@ /** Tests for {@link FileIndexEvaluator}. */ public class FileIndexEvaluatorTest { + @Test + public void testLimitIntersectsBitmapDeletionVector() throws Exception { + BitmapDeletionVector deletionVector = new BitmapDeletionVector(); + for (int position = 0; position < 5; position++) { + deletionVector.delete(position); + } + + FileIndexResult result = + FileIndexEvaluator.evaluate( + null, + null, + Collections.emptyList(), + null, + 10, + null, + DataFileTestUtils.newFile("data.avro", 0, 0, 19, 0L), + deletionVector); + + assertThat(result).isInstanceOf(BitmapIndexResult.class); + assertThat(((BitmapIndexResult) result).get()) + .isEqualTo(RoaringBitmap32.bitmapOfRange(5, 15)); + } + + @Test + public void testBitmap64DeletionVectorFallsBack() throws Exception { + Bitmap64DeletionVector deletionVector = new Bitmap64DeletionVector(); + deletionVector.delete(0); + + FileIndexResult result = + FileIndexEvaluator.evaluate( + null, + null, + Collections.emptyList(), + null, + 10, + null, + DataFileTestUtils.newFile("data.avro", 0, 0, 19, 0L), + deletionVector); + + assertThat(result).isSameAs(FileIndexResult.REMAIN); + } + + @Test + public void testLargeFileFallsBack() throws Exception { + FileIndexResult result = + FileIndexEvaluator.evaluate( + null, + null, + Collections.emptyList(), + null, + 10, + null, + fileWithRowCount(Integer.MAX_VALUE + 2L), + null); + + assertThat(result).isSameAs(FileIndexResult.REMAIN); + } + @Test public void testDataFilterIntersectsDeletionVector() throws Exception { BitmapDeletionVector deletionVector = new BitmapDeletionVector(); @@ -95,6 +156,24 @@ private static TableSchema tableSchema() { null); } + private static DataFileMeta fileWithRowCount(long rowCount) { + return DataFileMeta.forAppend( + "data.avro", + 0, + rowCount, + SimpleStats.EMPTY_STATS, + 0, + 0, + 0, + Collections.emptyList(), + null, + null, + null, + null, + null, + null); + } + private static byte[] embeddedBitmapIndex() throws IOException { BitmapFileIndex bitmapFileIndex = new BitmapFileIndex(DataTypes.INT(), new Options()); FileIndexWriter indexWriter = bitmapFileIndex.createWriter(); diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/RawFileSplitReadTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/RawFileSplitReadTest.java index fd4fcb2d4853..57de4d266ec4 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/RawFileSplitReadTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/RawFileSplitReadTest.java @@ -18,10 +18,13 @@ package org.apache.paimon.operation; +import org.apache.paimon.AppendOnlyFileStore; import org.apache.paimon.CoreOptions; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; +import org.apache.paimon.deletionvectors.Bitmap64DeletionVector; +import org.apache.paimon.deletionvectors.DeletionVector; import org.apache.paimon.format.FileFormat; import org.apache.paimon.format.FlushingFileFormat; import org.apache.paimon.format.FormatReaderFactory; @@ -43,11 +46,15 @@ import org.apache.paimon.table.source.InnerTableRead; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.IOExceptionSupplier; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; @@ -133,7 +140,47 @@ public FormatReaderFactory createReaderFactory( } } + @Test + void testLimitAfterBitmap64DeletionVector() throws Exception { + List rows = new ArrayList<>(); + for (int i = 0; i < 20; i++) { + rows.add(GenericRow.of(BinaryString.fromString("value-" + i), i)); + } + FileStoreTable table = createTable("bitmap64-limit", rows); + DataSplit split = singleSplit(table); + assertThat(split.dataFiles()).hasSize(1); + + DeletionVector deletionVector = new Bitmap64DeletionVector(); + for (int position = 0; position < 5; position++) { + deletionVector.delete(position); + } + String fileName = split.dataFiles().get(0).fileName(); + Map> deletionVectorFactories = + Collections.singletonMap(fileName, () -> deletionVector); + + RawFileSplitRead read = ((AppendOnlyFileStore) table.store()).newRead(); + read.withLimit(10); + AtomicInteger count = new AtomicInteger(); + try (RecordReader reader = + read.createReader( + split.partition(), + split.bucket(), + split.dataFiles(), + deletionVectorFactories)) { + reader.forEachRemaining(ignored -> count.incrementAndGet()); + } + + assertThat(count).hasValue(10); + } + private FileStoreTable createTable(String directory) throws Exception { + return createTable( + directory, + Collections.singletonList(GenericRow.of(BinaryString.fromString("value"), 42))); + } + + private FileStoreTable createTable(String directory, List rows) + throws Exception { Path tablePath = new Path(tempDir.resolve(directory).toUri()); Options options = new Options(); options.set(CoreOptions.PATH, tablePath.toString()); @@ -153,7 +200,9 @@ private FileStoreTable createTable(String directory) throws Exception { BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); try (BatchTableWrite write = writeBuilder.newWrite(); BatchTableCommit commit = writeBuilder.newCommit()) { - write.write(GenericRow.of(BinaryString.fromString("value"), 42)); + for (InternalRow row : rows) { + write.write(row); + } commit.commit(write.prepareCommit()); } return table;