This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 552cb2602a [mosaic] Push row selection down to row-group reads (#9742)
552cb2602a is described below
commit 552cb2602a1fabf91259950245ff18863503f842
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(
