This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit f41cbcf860a9516688f5ff371d7a5b3bb595e9d5 Author: jianguotian <[email protected]> AuthorDate: Fri Sep 11 23:08:29 2026 +0800 [mosaic] Push row selection down to row-group reads (#9742) --- .../paimon/format/mosaic/MosaicReaderFactory.java | 1 + .../paimon/format/mosaic/MosaicRecordsReader.java | 70 ++++++- .../format/mosaic/MosaicReaderWriterTest.java | 78 ++++++++ .../format/mosaic/MosaicRecordsReaderTest.java | 209 +++++++++++++++++++++ 4 files changed, 357 insertions(+), 1 deletion(-) diff --git a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicReaderFactory.java b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicReaderFactory.java index 3c8c2fd38a..c71b55c394 100644 --- a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicReaderFactory.java +++ b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicReaderFactory.java @@ -64,6 +64,7 @@ public class MosaicReaderFactory implements FormatReaderFactory { projectedRowType, predicates, context.filePath(), + context.selection(), prefetchRowGroups, prefetchMaxBytes); } diff --git a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java index e9a1712b77..c36285baea 100644 --- a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java +++ b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java @@ -32,6 +32,7 @@ import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataType; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.ExecutorThreadFactory; +import org.apache.paimon.utils.RoaringBitmap32; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -73,6 +74,7 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { private final int projectedFieldCount; private final boolean allProjectedColumnsMissing; @Nullable private final List<Predicate> predicates; + @Nullable private final RoaringBitmap32 selection; /** Opens upcoming row groups while the current one is consumed; opens are thread-safe. */ private static final ExecutorService PREFETCH_POOL = @@ -113,6 +115,29 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { projectedRowType, predicates, filePath, + null, + prefetchRowGroups, + prefetchMaxBytes); + } + + MosaicRecordsReader( + MosaicInputFileAdapter inputFileAdapter, + long fileSize, + RowType dataSchemaRowType, + RowType projectedRowType, + @Nullable List<Predicate> predicates, + Path filePath, + @Nullable RoaringBitmap32 selection, + int prefetchRowGroups, + long prefetchMaxBytes) { + this( + inputFileAdapter, + fileSize, + dataSchemaRowType, + projectedRowType, + predicates, + filePath, + selection, new RootAllocator(), MosaicReader::open, prefetchRowGroups, @@ -135,6 +160,7 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { projectedRowType, predicates, filePath, + null, allocator, nativeReaderOpener, MosaicFileFormat.READ_PREFETCH_ROW_GROUPS.defaultValue(), @@ -152,11 +178,38 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { NativeReaderOpener nativeReaderOpener, int prefetchRowGroups, long prefetchMaxBytes) { + this( + inputFileAdapter, + fileSize, + dataSchemaRowType, + projectedRowType, + predicates, + filePath, + null, + allocator, + nativeReaderOpener, + prefetchRowGroups, + prefetchMaxBytes); + } + + MosaicRecordsReader( + MosaicInputFileAdapter inputFileAdapter, + long fileSize, + RowType dataSchemaRowType, + RowType projectedRowType, + @Nullable List<Predicate> predicates, + Path filePath, + @Nullable RoaringBitmap32 selection, + BufferAllocator allocator, + NativeReaderOpener nativeReaderOpener, + int prefetchRowGroups, + long prefetchMaxBytes) { this.filePath = filePath; this.inputFileAdapter = inputFileAdapter; this.dataSchemaRowType = dataSchemaRowType; this.projectedFieldCount = projectedRowType.getFieldCount(); this.predicates = predicates; + this.selection = selection; this.allocator = allocator; MosaicReader createdReader = null; @@ -321,7 +374,7 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { int index = nextRowGroupToSchedule; int numRows = reader.rowGroupNumRows(index); long startPosition = scheduledRowCount; - if (!matchesRowGroup(index, numRows)) { + if (!matchesSelection(startPosition, numRows) || !matchesRowGroup(index, numRows)) { nextRowGroupToSchedule++; scheduledRowCount += numRows; continue; @@ -349,6 +402,21 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { } } + private boolean matchesSelection(long rowGroupStart, long rowCount) { + if (selection == null) { + return true; + } + if (rowCount <= 0 || rowGroupStart < 0 || rowGroupStart > RoaringBitmap32.MAX_VALUE) { + return false; + } + + long maxSupremum = (long) RoaringBitmap32.MAX_VALUE + 1; + long remainingAddressableRows = maxSupremum - rowGroupStart; + long rowGroupEnd = + rowCount > remainingAddressableRows ? maxSupremum : rowGroupStart + rowCount; + return selection.intersects(rowGroupStart, rowGroupEnd); + } + /** A row group whose data is being, or has been, loaded. */ private static final class RowGroupBatch { final int index; diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java index a72001f5a8..aceca5c455 100644 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java @@ -32,14 +32,17 @@ import org.apache.paimon.format.FormatWriter; import org.apache.paimon.format.FormatWriterFactory; import org.apache.paimon.fs.Path; import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.io.DataFileRecordReader; import org.apache.paimon.options.MemorySize; import org.apache.paimon.options.Options; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.reader.FileRecordReader; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.RoaringBitmap32; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; @@ -246,6 +249,81 @@ class MosaicReaderWriterTest { reader.close(); } + @Test + void testEmptySelectionSkipsAllRowGroups() throws IOException { + RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING()); + Path path = newPath(); + + writeRows( + rowType, + path, + GenericRow.of(1, BinaryString.fromString("a")), + GenericRow.of(2, BinaryString.fromString("b"))); + + MosaicFileFormat format = createFormat(); + FormatReaderFactory readerFactory = format.createReaderFactory(rowType, rowType, null); + LocalFileIO fileIO = new LocalFileIO(); + try (RecordReader<InternalRow> reader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, + path, + fileIO.getFileSize(path), + new RoaringBitmap32(), + null))) { + assertThat(reader.readBatch()).isNull(); + } + } + + @Test + void testNonEmptySelectionThroughReaderFactory() throws IOException { + RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING()); + Path path = newPath(); + + writeRows( + rowType, + path, + GenericRow.of(1, BinaryString.fromString("a")), + GenericRow.of(2, BinaryString.fromString("b")), + GenericRow.of(3, BinaryString.fromString("c"))); + + MosaicFileFormat format = createFormat(); + FormatReaderFactory readerFactory = format.createReaderFactory(rowType, rowType, null); + LocalFileIO fileIO = new LocalFileIO(); + RoaringBitmap32 selection = RoaringBitmap32.bitmapOf(1); + FileRecordReader<InternalRow> mosaicReader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, path, fileIO.getFileSize(path), selection, null)); + try (DataFileRecordReader reader = + new DataFileRecordReader( + rowType, + mosaicReader, + false, + false, + null, + null, + null, + false, + null, + 0, + Collections.emptyMap(), + selection, + path)) { + FileRecordIterator<InternalRow> batch = + (FileRecordIterator<InternalRow>) reader.readBatch(); + assertThat(batch).isNotNull(); + InternalRow row = batch.next(); + assertThat(row).isNotNull(); + assertThat(row.getInt(0)).isEqualTo(2); + assertThat(row.getString(1).toString()).isEqualTo("b"); + assertThat(batch.returnedPosition()).isEqualTo(1); + assertThat(batch.next()).isNull(); + batch.releaseBatch(); + assertThat(reader.readBatch()).isNull(); + } + } + @Test void testProjectionWithMissingColumns() throws IOException { RowType writeType = diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsReaderTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsReaderTest.java index ca5af5f7e4..b540f8c1da 100644 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsReaderTest.java +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsReaderTest.java @@ -23,10 +23,12 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.fs.Path; import org.apache.paimon.fs.SeekableInputStream; import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.io.DataFileRecordReader; import org.apache.paimon.mosaic.MosaicReader; import org.apache.paimon.reader.FileRecordIterator; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.RoaringBitmap32; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -203,6 +205,192 @@ class MosaicRecordsReaderTest { recordsReader.close(); } + @Test + void testAllProjectedColumnsMissingPreservesSelectedPositionsAcrossRowGroups() + throws IOException { + CloseCountingSeekableInputStream inputStream = new CloseCountingSeekableInputStream(); + MosaicInputFileAdapter inputFileAdapter = createInputFileAdapter(inputStream); + CloseCountingRootAllocator allocator = new CloseCountingRootAllocator(); + MosaicReader reader = createReader(); + when(reader.numRowGroups()).thenReturn(3); + when(reader.rowGroupNumRows(0)).thenReturn(2); + when(reader.rowGroupNumRows(1)).thenReturn(2); + when(reader.rowGroupNumRows(2)).thenReturn(2); + + Path filePath = new Path("file:/tmp/mosaic-reader-test"); + RoaringBitmap32 selection = RoaringBitmap32.bitmapOf(1, 5); + MosaicRecordsReader recordsReader = + new MosaicRecordsReader( + inputFileAdapter, + 0, + rowType(), + rowType(), + null, + filePath, + selection, + allocator, + (inputFile, fileSize, bufferAllocator) -> reader, + 3, + MosaicFileFormat.READ_PREFETCH_MAX_BYTES.defaultValue().getBytes()); + DataFileRecordReader dataFileReader = + new DataFileRecordReader( + rowType(), + recordsReader, + false, + false, + null, + null, + null, + false, + null, + 0, + Collections.emptyMap(), + selection, + filePath); + + FileRecordIterator<InternalRow> firstBatch = dataFileReader.readBatch(); + assertThat(firstBatch).isNotNull(); + InternalRow firstRow = firstBatch.next(); + assertThat(firstRow).isNotNull(); + assertThat(firstRow.isNullAt(0)).isTrue(); + assertThat(firstBatch.returnedPosition()).isEqualTo(1); + assertThat(firstBatch.next()).isNull(); + firstBatch.releaseBatch(); + + FileRecordIterator<InternalRow> secondBatch = dataFileReader.readBatch(); + assertThat(secondBatch).isNotNull(); + InternalRow secondRow = secondBatch.next(); + assertThat(secondRow).isNotNull(); + assertThat(secondRow.isNullAt(0)).isTrue(); + assertThat(secondBatch.returnedPosition()).isEqualTo(5); + assertThat(secondBatch.next()).isNull(); + secondBatch.releaseBatch(); + + assertThat(dataFileReader.readBatch()).isNull(); + verify(reader, never()).readRowGroup(anyInt(), any()); + + dataFileReader.close(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + + @Test + void testPrefetchSkipsUnselectedRowGroupsAndPreservesPositions() throws IOException { + CloseCountingSeekableInputStream inputStream = new CloseCountingSeekableInputStream(); + MosaicInputFileAdapter inputFileAdapter = createInputFileAdapter(inputStream); + CloseCountingRootAllocator allocator = new CloseCountingRootAllocator(); + MosaicReader reader = mock(MosaicReader.class); + VectorSchemaRoot firstRoot = intRoot(allocator, 0, 1); + VectorSchemaRoot thirdRoot = intRoot(allocator, 5, 6); + VectorSchemaRoot fourthRoot = intRoot(allocator, 7, 8, 9); + when(reader.getSchema()).thenReturn(firstRoot.getSchema()); + when(reader.numRowGroups()).thenReturn(4); + when(reader.rowGroupNumRows(0)).thenReturn(2); + when(reader.rowGroupNumRows(1)).thenReturn(3); + when(reader.rowGroupNumRows(2)).thenReturn(2); + when(reader.rowGroupNumRows(3)).thenReturn(3); + when(reader.readRowGroup(0, allocator)).thenReturn(firstRoot); + when(reader.readRowGroup(2, allocator)).thenReturn(thirdRoot); + when(reader.readRowGroup(3, allocator)).thenReturn(fourthRoot); + + Path filePath = new Path("file:/tmp/mosaic-reader-test"); + RoaringBitmap32 selection = RoaringBitmap32.bitmapOf(1, 5, 9); + MosaicRecordsReader recordsReader = + new MosaicRecordsReader( + inputFileAdapter, + 0, + rowType(), + rowType(), + null, + filePath, + selection, + allocator, + (inputFile, fileSize, bufferAllocator) -> reader, + 4, + MosaicFileFormat.READ_PREFETCH_MAX_BYTES.defaultValue().getBytes()); + DataFileRecordReader dataFileReader = + new DataFileRecordReader( + rowType(), + recordsReader, + false, + false, + null, + null, + null, + false, + null, + 0, + Collections.emptyMap(), + selection, + filePath); + + assertSelectedRow(dataFileReader.readBatch(), 1, 1); + assertSelectedRow(dataFileReader.readBatch(), 5, 5); + assertSelectedRow(dataFileReader.readBatch(), 9, 9); + assertThat(dataFileReader.readBatch()).isNull(); + + verify(reader).readRowGroup(0, allocator); + verify(reader, never()).readRowGroup(1, allocator); + verify(reader).readRowGroup(2, allocator); + verify(reader).readRowGroup(3, allocator); + + dataFileReader.close(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + + @Test + void testPrefetchPreservesPositionAfterPartiallyConsumedBatch() throws IOException { + CloseCountingSeekableInputStream inputStream = new CloseCountingSeekableInputStream(); + MosaicInputFileAdapter inputFileAdapter = createInputFileAdapter(inputStream); + CloseCountingRootAllocator allocator = new CloseCountingRootAllocator(); + MosaicReader reader = mock(MosaicReader.class); + VectorSchemaRoot secondRoot = intRoot(allocator, 3, 4, 5, 6); + VectorSchemaRoot thirdRoot = intRoot(allocator, 7, 8); + when(reader.getSchema()).thenReturn(secondRoot.getSchema()); + when(reader.numRowGroups()).thenReturn(3); + when(reader.rowGroupNumRows(0)).thenReturn(3); + when(reader.rowGroupNumRows(1)).thenReturn(4); + when(reader.rowGroupNumRows(2)).thenReturn(2); + when(reader.readRowGroup(1, allocator)).thenReturn(secondRoot); + when(reader.readRowGroup(2, allocator)).thenReturn(thirdRoot); + + RoaringBitmap32 selection = RoaringBitmap32.bitmapOf(3, 7); + MosaicRecordsReader recordsReader = + new MosaicRecordsReader( + inputFileAdapter, + 0, + rowType(), + rowType(), + null, + new Path("file:/tmp/mosaic-reader-test"), + selection, + allocator, + (inputFile, fileSize, bufferAllocator) -> reader, + 2, + MosaicFileFormat.READ_PREFETCH_MAX_BYTES.defaultValue().getBytes()); + + FileRecordIterator<InternalRow> secondBatch = + recordsReader.readBatch().selection(RoaringBitmap32.bitmapOf(3)); + assertThat(secondBatch.next().getInt(0)).isEqualTo(3); + assertThat(secondBatch.returnedPosition()).isEqualTo(3); + assertThat(secondBatch.next()).isNull(); + secondBatch.releaseBatch(); + + FileRecordIterator<InternalRow> thirdBatch = + recordsReader.readBatch().selection(RoaringBitmap32.bitmapOf(7)); + assertThat(thirdBatch.next().getInt(0)).isEqualTo(7); + assertThat(thirdBatch.returnedPosition()).isEqualTo(7); + assertThat(thirdBatch.next()).isNull(); + thirdBatch.releaseBatch(); + assertThat(recordsReader.readBatch()).isNull(); + + verify(reader, never()).readRowGroup(0, allocator); + verify(reader).readRowGroup(1, allocator); + verify(reader).readRowGroup(2, allocator); + + recordsReader.close(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + @Test void testDisabledPrefetchReadsRowGroupsOnDemand() throws IOException { CloseCountingSeekableInputStream inputStream = new CloseCountingSeekableInputStream(); @@ -447,6 +635,27 @@ class MosaicRecordsReaderTest { return root; } + private static void assertSelectedRow( + FileRecordIterator<InternalRow> batch, int value, long position) throws IOException { + assertThat(batch).isNotNull(); + assertThat(batch.next().getInt(0)).isEqualTo(value); + assertThat(batch.returnedPosition()).isEqualTo(position); + assertThat(batch.next()).isNull(); + batch.releaseBatch(); + } + + private static VectorSchemaRoot intRoot(RootAllocator allocator, int... values) { + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(rowType(), allocator); + IntVector vector = (IntVector) root.getVector(0); + vector.allocateNew(values.length); + for (int i = 0; i < values.length; i++) { + vector.setSafe(i, values[i]); + } + vector.setValueCount(values.length); + root.setRowCount(values.length); + return root; + } + private static MosaicInputFileAdapter createInputFileAdapter( CloseCountingSeekableInputStream inputStream) throws IOException { return new MosaicInputFileAdapter(
