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 8548906653 [mosaic] Respect decoded byte budget when refilling prefetch
8548906653 is described below
commit 85489066537858239156d295e9737bb0be459a32
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(
