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 332d63ca5bdfeb3ff1634917929e19df78d65ad7 Author: JingsongLi <[email protected]> AuthorDate: Fri Sep 11 22:10:54 2026 +0800 [mosaic] Respect decoded byte budget when refilling prefetch --- .../paimon/format/mosaic/MosaicRecordsReader.java | 10 ++-- .../format/mosaic/MosaicRecordsReaderTest.java | 69 ++++++++++++---------- 2 files changed, 44 insertions(+), 35 deletions(-) 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 6168418097..e9a1712b77 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 @@ -294,7 +294,7 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { @Nullable private RowGroupBatch nextRowGroup() throws IOException { if (pending.isEmpty()) { - fillPrefetchQueue(Math.max(1, prefetchDepth)); + fillPrefetchQueue(Math.max(1, prefetchDepth), true); } RowGroupBatch head = pending.peek(); if (head == null) { @@ -310,13 +310,13 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { currentVsr = vsr; if (prefetchDepth > 0) { // currentVsr is owned by this reader, so a failure here leaves nothing unreleased. - fillPrefetchQueue(prefetchDepth); + fillPrefetchQueue(prefetchDepth, false); } return head; } /** Schedules matching row groups until {@code wanted} are queued or the byte budget is used. */ - private void fillPrefetchQueue(int wanted) { + private void fillPrefetchQueue(int wanted, boolean readOnDemand) { while (pending.size() < wanted && nextRowGroupToSchedule < numRowGroups) { int index = nextRowGroupToSchedule; int numRows = reader.rowGroupNumRows(index); @@ -327,8 +327,8 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { continue; } long bytes = numRows * estimatedRowBytes; - // The first queued row group is always read; the rest must fit the decoded budget. - if (!pending.isEmpty() && pendingBytes + bytes > prefetchMaxBytes) { + // Only an on-demand read may exceed the decoded budget to make progress. + if ((!readOnDemand || !pending.isEmpty()) && pendingBytes + bytes > prefetchMaxBytes) { return; } nextRowGroupToSchedule++; 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 0e88b42909..ca5af5f7e4 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 @@ -36,6 +36,8 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import java.io.IOException; import java.io.InterruptedIOException; @@ -55,7 +57,6 @@ import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; -import static org.mockito.Mockito.timeout; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -369,39 +370,47 @@ class MosaicRecordsReaderTest { assertThat(allocator.closeCount()).isEqualTo(1); } - @Test - void testPrefetchIsBoundedByEstimatedDecodedBytes() throws IOException { + @ParameterizedTest + @CsvSource({ + "0, 1, 1", + "4000, 1, 1", + "4000, 4, 4", + "5000, 1, 2", + "9999, 1, 2", + "10000, 1, 3", + "100000, 1, 4" + }) + void testPrefetchIsBoundedByEstimatedDecodedBytes( + long budget, int batchesToRead, int expectedReads) throws IOException { // One INT column: 5 bytes per row; 1,000 rows per row group is 5,000 bytes. assertThat(MosaicRecordsReader.estimatedRowBytes(rowType())).isEqualTo(5); - for (long budget : new long[] {4_000L, 100_000L}) { - CloseCountingSeekableInputStream inputStream = new CloseCountingSeekableInputStream(); - MosaicInputFileAdapter inputFileAdapter = createInputFileAdapter(inputStream); - CloseCountingRootAllocator allocator = new CloseCountingRootAllocator(); - MosaicReader reader = createProjectedReader(allocator, 4); - when(reader.rowGroupNumRows(anyInt())).thenReturn(1000); - MosaicRecordsReader recordsReader = - new MosaicRecordsReader( - inputFileAdapter, - 0, - rowType(), - rowType(), - null, - new Path("file:/tmp/mosaic-reader-test"), - allocator, - (inputFile, fileSize, bufferAllocator) -> reader, - 8, - budget); - assertThat(recordsReader.readBatch()).isNotNull(); - if (budget < 5_000L) { - // Below one row group: nothing is read ahead of the batch being consumed. - verify(reader, times(1)).readRowGroup(anyInt(), any()); - } else { - // The three remaining row groups fit the budget and are read ahead. - verify(reader, timeout(5_000).times(4)).readRowGroup(anyInt(), any()); + CloseCountingSeekableInputStream inputStream = new CloseCountingSeekableInputStream(); + MosaicInputFileAdapter inputFileAdapter = createInputFileAdapter(inputStream); + CloseCountingRootAllocator allocator = new CloseCountingRootAllocator(); + MosaicReader reader = createProjectedReader(allocator, 4); + when(reader.rowGroupNumRows(anyInt())).thenReturn(1000); + try (MosaicRecordsReader recordsReader = + new MosaicRecordsReader( + inputFileAdapter, + 0, + rowType(), + rowType(), + null, + new Path("file:/tmp/mosaic-reader-test"), + allocator, + (inputFile, fileSize, bufferAllocator) -> reader, + 8, + budget)) { + for (int i = 0; i < batchesToRead; i++) { + FileRecordIterator<InternalRow> batch = recordsReader.readBatch(); + assertThat(batch).isNotNull(); + assertThat(batch.next().getInt(0)).isEqualTo(i); + batch.releaseBatch(); } - recordsReader.close(); - assertThat(allocator.closeCount()).isEqualTo(1); } + // Closing drains scheduled reads, so verification cannot race with background tasks. + verify(reader, times(expectedReads)).readRowGroup(anyInt(), any()); + assertThat(allocator.closeCount()).isEqualTo(1); } private static void closeQuietly(
