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 6948dd55d0 [mosaic] Prefetch row groups and read with one input stream
per concurrent read (#9740)
6948dd55d0 is described below
commit 6948dd55d07ab8903362be1a841dfd8282f5fe3f
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Fri Sep 11 17:47:54 2026 +0800
[mosaic] Prefetch row groups and read with one input stream per concurrent
read (#9740)
---
docs/docs/concepts/spec/fileformat.md | 2 +
.../paimon/format/mosaic/MosaicFileFormat.java | 30 +-
.../format/mosaic/MosaicInputFileAdapter.java | 156 ++++++++--
.../paimon/format/mosaic/MosaicReaderFactory.java | 16 +-
.../paimon/format/mosaic/MosaicRecordsReader.java | 337 ++++++++++++++++++---
.../format/mosaic/MosaicInputFileAdapterTest.java | 162 ++++++++++
.../format/mosaic/MosaicReaderWriterTest.java | 122 ++++++++
.../format/mosaic/MosaicRecordsReaderTest.java | 270 +++++++++++++++++
8 files changed, 1025 insertions(+), 70 deletions(-)
diff --git a/docs/docs/concepts/spec/fileformat.md
b/docs/docs/concepts/spec/fileformat.md
index ff5c86f9de..8367d86db2 100644
--- a/docs/docs/concepts/spec/fileformat.md
+++ b/docs/docs/concepts/spec/fileformat.md
@@ -285,6 +285,8 @@ Format Options:
| --- | --- | --- | --- |
| `mosaic.num-buckets` | auto | Integer | Number of column buckets for
parallel I/O. When set to 0 or not specified, the format auto-determines the
bucket count. |
| `mosaic.stats-columns` | (empty) | String | Comma-separated column names to
collect min/max statistics for filter pushdown. Empty means no statistics are
collected. |
+| `mosaic.read.prefetch-row-groups` | 8 | Integer | Number of row groups a
reader opens ahead of the one being consumed. Opening a row group issues
several dependent range reads, so prefetching overlaps that latency with
decoding. Each row group ahead keeps its decoded batch in memory and uses its
own input stream, see `mosaic.read.prefetch-max-bytes`. 0 disables prefetching.
|
+| `mosaic.read.prefetch-max-bytes` | 64 mb | MemorySize | Upper bound on the
estimated decoded size of the row groups a reader keeps ahead, from their row
counts and the projected column types. Wide projections or large row groups
therefore lower the effective `mosaic.read.prefetch-row-groups`. |
Limitations:
1. Mosaic does not support complex types: ARRAY, MAP, MULTISET, ROW, VARIANT,
BLOB, VECTOR.
diff --git
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicFileFormat.java
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicFileFormat.java
index 00843ca572..6a1721d58f 100644
---
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicFileFormat.java
+++
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicFileFormat.java
@@ -25,6 +25,7 @@ import org.apache.paimon.format.FormatWriterFactory;
import org.apache.paimon.format.SimpleStatsExtractor;
import org.apache.paimon.options.ConfigOption;
import org.apache.paimon.options.ConfigOptions;
+import org.apache.paimon.options.MemorySize;
import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.statistics.SimpleColStatsCollector;
import org.apache.paimon.types.ArrayType;
@@ -75,6 +76,28 @@ public class MosaicFileFormat extends FileFormat {
.noDefaultValue()
.withDescription("Number of column buckets for parallel
IO.");
+ public static final ConfigOption<Integer> READ_PREFETCH_ROW_GROUPS =
+ ConfigOptions.key("mosaic.read.prefetch-row-groups")
+ .intType()
+ .defaultValue(8)
+ .withDescription(
+ "Number of row groups a reader opens ahead of the
one being consumed. "
+ + "Opening a row group issues several
dependent range reads, "
+ + "so prefetching overlaps that latency
with decoding. Each "
+ + "row group ahead keeps its decoded batch
in memory and uses "
+ + "its own input stream, see
'mosaic.read.prefetch-max-bytes'. "
+ + "0 disables prefetching.");
+
+ public static final ConfigOption<MemorySize> READ_PREFETCH_MAX_BYTES =
+ ConfigOptions.key("mosaic.read.prefetch-max-bytes")
+ .memoryType()
+ .defaultValue(MemorySize.ofMebiBytes(64))
+ .withDescription(
+ "Upper bound on the estimated decoded size of the
row groups a reader "
+ + "keeps ahead, from their row counts and
the projected column "
+ + "types. Wide projections or large row
groups therefore lower "
+ + "the effective
'mosaic.read.prefetch-row-groups'.");
+
static {
System.setProperty("arrow.enable_unsafe_memory_access", "true");
}
@@ -91,7 +114,12 @@ public class MosaicFileFormat extends FileFormat {
RowType dataSchemaRowType,
RowType projectedRowType,
@Nullable List<Predicate> predicates) {
- return new MosaicReaderFactory(dataSchemaRowType, projectedRowType,
predicates);
+ return new MosaicReaderFactory(
+ dataSchemaRowType,
+ projectedRowType,
+ predicates,
+ formatContext.options().get(READ_PREFETCH_ROW_GROUPS),
+
formatContext.options().get(READ_PREFETCH_MAX_BYTES).getBytes());
}
@Override
diff --git
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java
index 3a307ea0f2..07385fdd36 100644
---
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java
+++
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java
@@ -27,53 +27,155 @@ import org.apache.paimon.mosaic.InputFile;
import java.io.Closeable;
import java.io.EOFException;
import java.io.IOException;
+import java.io.InterruptedIOException;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.List;
/**
- * Adapts Paimon's {@link FileIO} to Mosaic's {@link InputFile} interface.
+ * Adapter that exposes a Paimon {@link SeekableInputStream} as a Mosaic
{@link InputFile}.
*
- * <p>Maintains a single {@link SeekableInputStream}. If the stream implements
{@link
- * VectoredReadable}, reads use {@link VectoredReadable#preadFully} which is
thread-safe. Otherwise,
- * reads are synchronized to protect seek+read sequences.
+ * <p>Each read borrows one of at most {@code maxStreams} input streams, so
concurrent reads do not
+ * serialize on a single stream; a read that finds every stream busy waits for
one.
*/
public class MosaicInputFileAdapter implements InputFile, Closeable {
+ private final FileIO fileIO;
private final Path path;
- private final SeekableInputStream in;
- private final VectoredReadable vectoredReadable;
+ private final int maxStreams;
+
+ private final ArrayDeque<SeekableInputStream> idleStreams = new
ArrayDeque<>();
+ private final List<SeekableInputStream> allStreams = new ArrayList<>();
+ private int openingStreams;
+ private boolean closed;
public MosaicInputFileAdapter(FileIO fileIO, Path path) throws IOException
{
+ this(fileIO, path, 1);
+ }
+
+ public MosaicInputFileAdapter(FileIO fileIO, Path path, int maxStreams)
throws IOException {
+ this.fileIO = fileIO;
this.path = path;
- this.in = fileIO.newInputStream(path);
- this.vectoredReadable = in instanceof VectoredReadable ?
(VectoredReadable) in : null;
+ this.maxStreams = Math.max(1, maxStreams);
+ // Open eagerly so that a missing file fails here rather than in a
native callback.
+ SeekableInputStream first = fileIO.newInputStream(path);
+ allStreams.add(first);
+ idleStreams.push(first);
}
@Override
public void readFully(long position, byte[] buffer, int offset, int
length) throws IOException {
- if (vectoredReadable != null) {
- vectoredReadable.preadFully(position, buffer, offset, length);
- } else {
- synchronized (in) {
- in.seek(position);
- int remaining = length;
- int off = offset;
- while (remaining > 0) {
- int read = in.read(buffer, off, remaining);
- if (read < 0) {
- throw new EOFException(
- "Reached end of file while reading "
- + path
- + " at position "
- + position);
- }
- off += read;
- remaining -= read;
+ SeekableInputStream in = borrow();
+ try {
+ doReadFully(in, position, buffer, offset, length);
+ } finally {
+ release(in);
+ }
+ }
+
+ private void doReadFully(
+ SeekableInputStream in, long position, byte[] buffer, int offset,
int length)
+ throws IOException {
+ if (in instanceof VectoredReadable) {
+ ((VectoredReadable) in).preadFully(position, buffer, offset,
length);
+ return;
+ }
+ // The stream is borrowed exclusively, so seek + read needs no extra
locking.
+ in.seek(position);
+ int remaining = length;
+ int off = offset;
+ while (remaining > 0) {
+ int read = in.read(buffer, off, remaining);
+ if (read < 0) {
+ throw new EOFException(
+ "Reached end of file while reading " + path + " at
position " + position);
+ }
+ off += read;
+ remaining -= read;
+ }
+ }
+
+ private SeekableInputStream borrow() throws IOException {
+ synchronized (this) {
+ while (true) {
+ if (closed) {
+ throw new IOException("Input file " + path + " is closed");
+ }
+ SeekableInputStream idle = idleStreams.poll();
+ if (idle != null) {
+ return idle;
+ }
+ if (allStreams.size() + openingStreams < maxStreams) {
+ openingStreams++;
+ break;
+ }
+ try {
+ wait();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new InterruptedIOException(
+ "Interrupted while waiting for an input stream of
" + path);
+ }
+ }
+ }
+ SeekableInputStream opened = null;
+ try {
+ opened = fileIO.newInputStream(path);
+ } finally {
+ synchronized (this) {
+ openingStreams--;
+ if (opened != null && !closed) {
+ allStreams.add(opened);
+ } else {
+ notifyAll();
}
}
}
+ if (closed) {
+ opened.close();
+ throw new IOException("Input file " + path + " is closed");
+ }
+ return opened;
+ }
+
+ private void release(SeekableInputStream in) throws IOException {
+ synchronized (this) {
+ if (!closed) {
+ idleStreams.push(in);
+ notifyAll();
+ return;
+ }
+ }
+ in.close();
}
@Override
public void close() throws IOException {
- in.close();
+ List<SeekableInputStream> toClose;
+ synchronized (this) {
+ if (closed) {
+ return;
+ }
+ closed = true;
+ toClose = new ArrayList<>(allStreams);
+ allStreams.clear();
+ idleStreams.clear();
+ notifyAll();
+ }
+ IOException failure = null;
+ for (SeekableInputStream in : toClose) {
+ try {
+ in.close();
+ } catch (IOException e) {
+ if (failure == null) {
+ failure = e;
+ } else {
+ failure.addSuppressed(e);
+ }
+ }
+ }
+ if (failure != null) {
+ throw failure;
+ }
}
}
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 5b39c867e2..3c8c2fd38a 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
@@ -35,26 +35,36 @@ public class MosaicReaderFactory implements
FormatReaderFactory {
private final RowType dataSchemaRowType;
private final RowType projectedRowType;
@Nullable private final List<Predicate> predicates;
+ private final int prefetchRowGroups;
+ private final long prefetchMaxBytes;
public MosaicReaderFactory(
RowType dataSchemaRowType,
RowType projectedRowType,
- @Nullable List<Predicate> predicates) {
+ @Nullable List<Predicate> predicates,
+ int prefetchRowGroups,
+ long prefetchMaxBytes) {
this.dataSchemaRowType = dataSchemaRowType;
this.projectedRowType = projectedRowType;
this.predicates = predicates;
+ this.prefetchRowGroups = Math.max(0, prefetchRowGroups);
+ this.prefetchMaxBytes = prefetchMaxBytes;
}
@Override
public FileRecordReader<InternalRow> createReader(Context context) throws
IOException {
+ // One stream per row group being opened, plus one for the consumer's
own reads.
MosaicInputFileAdapter inputFile =
- new MosaicInputFileAdapter(context.fileIO(),
context.filePath());
+ new MosaicInputFileAdapter(
+ context.fileIO(), context.filePath(),
prefetchRowGroups + 1);
return new MosaicRecordsReader(
inputFile,
context.fileSize(),
dataSchemaRowType,
projectedRowType,
predicates,
- context.filePath());
+ context.filePath(),
+ 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 d31dc8c3e8..6168418097 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
@@ -31,22 +31,32 @@ import org.apache.paimon.reader.FileRecordReader;
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.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import javax.annotation.Nullable;
import java.io.IOException;
+import java.io.InterruptedIOException;
+import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
import static org.apache.paimon.format.mosaic.MosaicObjects.convertStatsValue;
@@ -64,17 +74,38 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
private final boolean allProjectedColumnsMissing;
@Nullable private final List<Predicate> predicates;
- private int currentRowGroup;
+ /** Opens upcoming row groups while the current one is consumed; opens are
thread-safe. */
+ private static final ExecutorService PREFETCH_POOL =
+ Executors.newCachedThreadPool(new
ExecutorThreadFactory("mosaic-row-group-prefetch"));
+
+ private final int prefetchDepth;
+ private final long prefetchMaxBytes;
+ private final long estimatedRowBytes;
+ private long pendingBytes;
+ private final ArrayDeque<RowGroupBatch> pending = new ArrayDeque<>();
+ private int nextRowGroupToSchedule;
+ private long scheduledRowCount;
+
private long returnedPosition = -1;
private VectorSchemaRoot currentVsr;
+ private static final Logger LOG =
LoggerFactory.getLogger(MosaicRecordsReader.class);
+
+ // Read statistics reported at debug level when the reader closes.
+ private int openedRowGroups;
+ private long openNanos;
+ private long rowsReturned;
+ private final long createdNanos = System.nanoTime();
+
public MosaicRecordsReader(
MosaicInputFileAdapter inputFileAdapter,
long fileSize,
RowType dataSchemaRowType,
RowType projectedRowType,
@Nullable List<Predicate> predicates,
- Path filePath) {
+ Path filePath,
+ int prefetchRowGroups,
+ long prefetchMaxBytes) {
this(
inputFileAdapter,
fileSize,
@@ -83,7 +114,9 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
predicates,
filePath,
new RootAllocator(),
- MosaicReader::open);
+ MosaicReader::open,
+ prefetchRowGroups,
+ prefetchMaxBytes);
}
MosaicRecordsReader(
@@ -95,6 +128,30 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
Path filePath,
BufferAllocator allocator,
NativeReaderOpener nativeReaderOpener) {
+ this(
+ inputFileAdapter,
+ fileSize,
+ dataSchemaRowType,
+ projectedRowType,
+ predicates,
+ filePath,
+ allocator,
+ nativeReaderOpener,
+ MosaicFileFormat.READ_PREFETCH_ROW_GROUPS.defaultValue(),
+
MosaicFileFormat.READ_PREFETCH_MAX_BYTES.defaultValue().getBytes());
+ }
+
+ MosaicRecordsReader(
+ MosaicInputFileAdapter inputFileAdapter,
+ long fileSize,
+ RowType dataSchemaRowType,
+ RowType projectedRowType,
+ @Nullable List<Predicate> predicates,
+ Path filePath,
+ BufferAllocator allocator,
+ NativeReaderOpener nativeReaderOpener,
+ int prefetchRowGroups,
+ long prefetchMaxBytes) {
this.filePath = filePath;
this.inputFileAdapter = inputFileAdapter;
this.dataSchemaRowType = dataSchemaRowType;
@@ -135,63 +192,237 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
this.reader = createdReader;
this.numRowGroups = createdNumRowGroups;
+ this.prefetchDepth = Math.max(0, prefetchRowGroups);
+ this.prefetchMaxBytes = Math.max(0, prefetchMaxBytes);
+ this.estimatedRowBytes = estimatedRowBytes(projectedRowType);
this.allProjectedColumnsMissing = createdAllProjectedColumnsMissing;
- this.currentRowGroup = 0;
this.arrowBatchReader = createdArrowBatchReader;
}
+ /** Rough decoded size of one row of the projected columns, used for the
prefetch budget. */
+ static long estimatedRowBytes(RowType projectedRowType) {
+ long bytes = 0;
+ for (DataField field : projectedRowType.getFields()) {
+ switch (field.type().getTypeRoot()) {
+ case BOOLEAN:
+ case TINYINT:
+ bytes += 2;
+ break;
+ case SMALLINT:
+ bytes += 3;
+ break;
+ case INTEGER:
+ case FLOAT:
+ case DATE:
+ case TIME_WITHOUT_TIME_ZONE:
+ bytes += 5;
+ break;
+ case BIGINT:
+ case DOUBLE:
+ case TIMESTAMP_WITHOUT_TIME_ZONE:
+ case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
+ bytes += 9;
+ break;
+ case DECIMAL:
+ bytes += 17;
+ break;
+ case CHAR:
+ case VARCHAR:
+ case BINARY:
+ case VARBINARY:
+ bytes += 40;
+ break;
+ default:
+ bytes += 128;
+ }
+ }
+ return Math.max(1, bytes);
+ }
+
+ int prefetchDepth() {
+ return prefetchDepth;
+ }
+
@Nullable
@Override
public FileRecordIterator<InternalRow> readBatch() throws IOException {
- while (currentRowGroup < numRowGroups) {
- int numRows = reader.rowGroupNumRows(currentRowGroup);
- if (!matchesRowGroup(currentRowGroup, numRows)) {
- returnedPosition += numRows;
- currentRowGroup++;
- continue;
+ releaseCurrentVsr();
+
+ RowGroupBatch batch = nextRowGroup();
+ if (batch == null) {
+ return null;
+ }
+ // Rows of skipped row groups still count towards the file position.
+ returnedPosition = batch.startPosition - 1;
+ rowsReturned += batch.numRows;
+
+ if (allProjectedColumnsMissing) {
+ return allNullIterator(batch.numRows);
+ }
+
+ Iterator<InternalRow> rows =
arrowBatchReader.readBatch(currentVsr).iterator();
+
+ return new FileRecordIterator<InternalRow>() {
+ @Override
+ public long returnedPosition() {
+ return returnedPosition;
}
- releaseCurrentVsr();
+ @Override
+ public Path filePath() {
+ return filePath;
+ }
- if (allProjectedColumnsMissing) {
- currentRowGroup++;
- return allNullIterator(numRows);
+ @Nullable
+ @Override
+ public InternalRow next() {
+ if (rows.hasNext()) {
+ returnedPosition++;
+ return rows.next();
+ }
+ return null;
}
- VectorSchemaRoot vsr = reader.readRowGroup(currentRowGroup,
allocator);
- currentRowGroup++;
- this.currentVsr = vsr;
+ @Override
+ public void releaseBatch() {
+ releaseCurrentVsr();
+ }
+ };
+ }
- Iterator<InternalRow> rows =
arrowBatchReader.readBatch(vsr).iterator();
+ /** Returns the next matching row group with its data ready, or null at
end of file. */
+ @Nullable
+ private RowGroupBatch nextRowGroup() throws IOException {
+ if (pending.isEmpty()) {
+ fillPrefetchQueue(Math.max(1, prefetchDepth));
+ }
+ RowGroupBatch head = pending.peek();
+ if (head == null) {
+ return null;
+ }
+ // The head stays queued until its data has arrived, so close() can
still wait for it.
+ long waitStart = System.nanoTime();
+ VectorSchemaRoot vsr = head.await();
+ openNanos += System.nanoTime() - waitStart;
+ openedRowGroups++;
+ pending.poll();
+ pendingBytes -= head.bytes;
+ currentVsr = vsr;
+ if (prefetchDepth > 0) {
+ // currentVsr is owned by this reader, so a failure here leaves
nothing unreleased.
+ fillPrefetchQueue(prefetchDepth);
+ }
+ return head;
+ }
- return new FileRecordIterator<InternalRow>() {
- @Override
- public long returnedPosition() {
- return returnedPosition;
+ /** Schedules matching row groups until {@code wanted} are queued or the
byte budget is used. */
+ private void fillPrefetchQueue(int wanted) {
+ while (pending.size() < wanted && nextRowGroupToSchedule <
numRowGroups) {
+ int index = nextRowGroupToSchedule;
+ int numRows = reader.rowGroupNumRows(index);
+ long startPosition = scheduledRowCount;
+ if (!matchesRowGroup(index, numRows)) {
+ nextRowGroupToSchedule++;
+ scheduledRowCount += numRows;
+ 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)
{
+ return;
+ }
+ nextRowGroupToSchedule++;
+ scheduledRowCount += numRows;
+ Future<VectorSchemaRoot> future = null;
+ if (!allProjectedColumnsMissing) {
+ FutureTask<VectorSchemaRoot> task =
+ new FutureTask<>(() -> reader.readRowGroup(index,
allocator));
+ if (prefetchDepth == 0) {
+ task.run();
+ } else {
+ PREFETCH_POOL.execute(task);
}
+ future = task;
+ }
+ pending.add(new RowGroupBatch(index, numRows, startPosition,
bytes, future));
+ pendingBytes += bytes;
+ }
+ }
- @Override
- public Path filePath() {
- return filePath;
- }
+ /** A row group whose data is being, or has been, loaded. */
+ private static final class RowGroupBatch {
+ final int index;
+ final int numRows;
+ final long startPosition;
+ final long bytes;
+ @Nullable private final Future<VectorSchemaRoot> future;
+
+ RowGroupBatch(
+ int index,
+ int numRows,
+ long startPosition,
+ long bytes,
+ @Nullable Future<VectorSchemaRoot> future) {
+ this.index = index;
+ this.numRows = numRows;
+ this.startPosition = startPosition;
+ this.bytes = bytes;
+ this.future = future;
+ }
- @Nullable
- @Override
- public InternalRow next() {
- if (rows.hasNext()) {
- returnedPosition++;
- return rows.next();
- }
- return null;
+ /** Waits for the data; the batch is not released here even when the
wait fails. */
+ @Nullable
+ VectorSchemaRoot await() throws IOException {
+ if (future == null) {
+ return null;
+ }
+ try {
+ return future.get();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ InterruptedIOException interrupted =
+ new InterruptedIOException("Interrupted while opening
row group " + index);
+ interrupted.initCause(e);
+ throw interrupted;
+ } catch (ExecutionException e) {
+ Throwable cause = e.getCause();
+ if (cause instanceof IOException) {
+ throw (IOException) cause;
+ }
+ if (cause instanceof RuntimeException) {
+ throw (RuntimeException) cause;
}
+ if (cause instanceof Error) {
+ throw (Error) cause;
+ }
+ throw new IOException("Failed to open row group " + index,
cause);
+ }
+ }
- @Override
- public void releaseBatch() {
- releaseCurrentVsr();
+ /**
+ * Waits for the read to finish, even if the current thread is
interrupted, and releases its
+ * data. Returns whether an interrupt was swallowed while waiting.
+ */
+ boolean discard() {
+ if (future == null) {
+ return false;
+ }
+ boolean interrupted = false;
+ while (true) {
+ try {
+ VectorSchemaRoot vsr = future.get();
+ if (vsr != null) {
+ vsr.close();
+ }
+ return interrupted;
+ } catch (InterruptedException e) {
+ // The native read still uses the reader handle; it must
complete first.
+ interrupted = true;
+ } catch (ExecutionException e) {
+ return interrupted;
}
- };
+ }
}
- return null;
}
private FileRecordIterator<InternalRow> allNullIterator(int numRows) {
@@ -283,6 +514,19 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
throwable = t;
}
+ // Prefetched row groups must finish and be released before the native
reader and the
+ // allocator go away, even if this thread is interrupted.
+ boolean interrupted = false;
+ RowGroupBatch batch;
+ while ((batch = pending.poll()) != null) {
+ pendingBytes -= batch.bytes;
+ try {
+ interrupted |= batch.discard();
+ } catch (Throwable t) {
+ throwable = addSuppressed(throwable, t);
+ }
+ }
+
try {
reader.close();
} catch (Throwable t) {
@@ -301,6 +545,21 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
throwable = addSuppressed(throwable, t);
}
+ if (LOG.isDebugEnabled()) {
+ LOG.debug(
+ "Closed mosaic reader for {}: row groups {} (opened {},
prefetch depth {}), "
+ + "rows {}, waited {} ms for row groups, lifetime
{} ms",
+ filePath.getName(),
+ numRowGroups,
+ openedRowGroups,
+ prefetchDepth,
+ rowsReturned,
+ openNanos / 1_000_000,
+ (System.nanoTime() - createdNanos) / 1_000_000);
+ }
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
if (throwable != null) {
rethrow(throwable);
}
diff --git
a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapterTest.java
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapterTest.java
new file mode 100644
index 0000000000..6a0e4ffb64
--- /dev/null
+++
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapterTest.java
@@ -0,0 +1,162 @@
+/*
+ * 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.format.mosaic;
+
+import org.apache.paimon.fs.Path;
+import org.apache.paimon.fs.SeekableInputStream;
+import org.apache.paimon.fs.local.LocalFileIO;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for the stream pool of {@link MosaicInputFileAdapter}. */
+class MosaicInputFileAdapterTest {
+
+ @Test
+ void testConcurrentReadsAreCappedAtMaxStreams() throws Exception {
+ CountDownLatch readsStarted = new CountDownLatch(2);
+ CountDownLatch releaseReads = new CountDownLatch(1);
+ AtomicInteger opened = new AtomicInteger();
+ LocalFileIO fileIO =
+ new LocalFileIO() {
+ @Override
+ public SeekableInputStream newInputStream(Path path) {
+ opened.incrementAndGet();
+ return new BlockingStream(readsStarted, releaseReads);
+ }
+ };
+ MosaicInputFileAdapter adapter =
+ new MosaicInputFileAdapter(fileIO, new
Path("file:/tmp/mosaic-adapter-test"), 2);
+
+ List<Thread> readers = new ArrayList<>();
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+ for (int i = 0; i < 3; i++) {
+ Thread thread =
+ new Thread(
+ () -> {
+ try {
+ adapter.readFully(0, new byte[4], 0, 4);
+ } catch (Throwable t) {
+ failure.set(t);
+ }
+ });
+ thread.start();
+ readers.add(thread);
+ }
+ readsStarted.await();
+ // Two reads hold the two streams; the third must wait instead of
opening another one.
+ readers.get(0).join(200);
+ assertThat(opened.get()).isEqualTo(2);
+
assertThat(readers.stream().filter(Thread::isAlive).count()).isEqualTo(3);
+
+ releaseReads.countDown();
+ for (Thread thread : readers) {
+ thread.join();
+ }
+ assertThat(failure.get()).isNull();
+ assertThat(opened.get()).isEqualTo(2);
+ adapter.close();
+ }
+
+ @Test
+ void testCloseWakesWaitingReader() throws Exception {
+ CountDownLatch readsStarted = new CountDownLatch(1);
+ CountDownLatch releaseReads = new CountDownLatch(1);
+ LocalFileIO fileIO =
+ new LocalFileIO() {
+ @Override
+ public SeekableInputStream newInputStream(Path path) {
+ return new BlockingStream(readsStarted, releaseReads);
+ }
+ };
+ MosaicInputFileAdapter adapter =
+ new MosaicInputFileAdapter(fileIO, new
Path("file:/tmp/mosaic-adapter-test"), 1);
+ AtomicReference<Throwable> first = new AtomicReference<>();
+ AtomicReference<Throwable> second = new AtomicReference<>();
+ Thread holder = new Thread(() -> read(adapter, first));
+ holder.start();
+ readsStarted.await();
+ Thread waiter = new Thread(() -> read(adapter, second));
+ waiter.start();
+ waiter.join(200);
+ assertThat(waiter.isAlive()).isTrue();
+
+ adapter.close();
+ waiter.join();
+ assertThat(second.get()).isInstanceOf(IOException.class);
+ releaseReads.countDown();
+ holder.join();
+ }
+
+ private static void read(MosaicInputFileAdapter adapter,
AtomicReference<Throwable> failure) {
+ try {
+ adapter.readFully(0, new byte[4], 0, 4);
+ } catch (Throwable t) {
+ failure.set(t);
+ }
+ }
+
+ /** A stream whose reads block until released, to hold a pooled stream
busy. */
+ private static class BlockingStream extends SeekableInputStream {
+
+ private final CountDownLatch started;
+ private final CountDownLatch release;
+
+ private BlockingStream(CountDownLatch started, CountDownLatch release)
{
+ this.started = started;
+ this.release = release;
+ }
+
+ @Override
+ public void seek(long desired) {}
+
+ @Override
+ public long getPos() {
+ return 0;
+ }
+
+ @Override
+ public int read(byte[] b, int off, int len) throws IOException {
+ started.countDown();
+ try {
+ release.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException(e);
+ }
+ return len;
+ }
+
+ @Override
+ public int read() {
+ return -1;
+ }
+
+ @Override
+ public void close() {}
+ }
+}
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 ae354bd133..a72001f5a8 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,6 +32,7 @@ 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.options.MemorySize;
import org.apache.paimon.options.Options;
import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.predicate.PredicateBuilder;
@@ -55,6 +56,7 @@ import java.util.List;
import java.util.UUID;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.junit.jupiter.api.Assumptions.assumeTrue;
@@ -340,6 +342,116 @@ class MosaicReaderWriterTest {
assertThat(reached).isTrue();
}
+ @Test
+ void testPrefetchedReadMatchesSequentialRead() throws IOException {
+ RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING());
+ Path path = newPath();
+ int numRows = 20_000;
+ GenericRow[] rows = new GenericRow[numRows];
+ for (int i = 0; i < numRows; i++) {
+ rows[i] = GenericRow.of(i, BinaryString.fromString("value_" + i +
"_padding"));
+ }
+ // A tiny row group size produces many row groups, so prefetching is
exercised.
+ Options writeOptions = new Options();
+ writeOptions.set(MosaicFileFormat.STATS_COLUMNS, "f0");
+ writeRows(rowType, path, writeOptions, MemorySize.ofKibiBytes(32),
rows);
+
+ List<Long> sequentialPositions = new ArrayList<>();
+ List<InternalRow> sequential =
+ readAllWithPrefetch(rowType, path, null, 0,
sequentialPositions);
+ assertThat(sequential).hasSize(numRows);
+
+ List<Long> prefetchedPositions = new ArrayList<>();
+ List<InternalRow> prefetched =
+ readAllWithPrefetch(rowType, path, null, 3,
prefetchedPositions);
+ assertThat(prefetched).hasSize(numRows);
+ assertThat(prefetchedPositions).isEqualTo(sequentialPositions);
+ for (int i = 0; i < numRows; i++) {
+
assertThat(prefetched.get(i).getInt(0)).isEqualTo(sequential.get(i).getInt(0));
+
assertThat(prefetched.get(i).getString(1)).isEqualTo(sequential.get(i).getString(1));
+ }
+
+ // Row groups skipped by the predicate must still advance the returned
position.
+ Predicate predicate = new PredicateBuilder(rowType).greaterOrEqual(0,
numRows - 1000);
+ List<Long> filteredPositions = new ArrayList<>();
+ List<InternalRow> filtered =
+ readAllWithPrefetch(
+ rowType, path, Collections.singletonList(predicate),
3, filteredPositions);
+ assertThat(filtered).isNotEmpty();
+ assertThat(filtered.size()).isLessThan(numRows);
+ assertThat(filteredPositions.get(0)).isEqualTo((long)
filtered.get(0).getInt(0));
+ assertThat(filteredPositions.get(filteredPositions.size() - 1))
+ .isEqualTo((long) numRows - 1);
+ }
+
+ @Test
+ void testClosingReaderWithPendingPrefetchReleasesResources() throws
IOException {
+ RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING());
+ Path path = newPath();
+ GenericRow[] rows = new GenericRow[20_000];
+ for (int i = 0; i < rows.length; i++) {
+ rows[i] = GenericRow.of(i, BinaryString.fromString("value_" + i +
"_padding"));
+ }
+ writeRows(rowType, path, new Options(), MemorySize.ofKibiBytes(32),
rows);
+
+ FormatReaderFactory readerFactory = createReaderFactory(rowType, null,
4);
+ LocalFileIO fileIO = new LocalFileIO();
+ RecordReader<InternalRow> reader =
+ readerFactory.createReader(
+ new FormatReaderContext(
+ fileIO, path, fileIO.getFileSize(path), null,
null));
+ // Consume one batch only; the prefetched row groups are still in
flight or buffered.
+ assertThat(reader.readBatch()).isNotNull();
+ // Closing must wait for and release them (an Arrow allocator leak
would throw here).
+ assertThatCode(reader::close).doesNotThrowAnyException();
+ }
+
+ private void writeRows(
+ RowType rowType, Path path, Options options, MemorySize blockSize,
GenericRow... rows)
+ throws IOException {
+ MosaicFileFormat format =
+ new MosaicFileFormat(
+ new FileFormatFactory.FormatContext(
+ options, 1024, 1024, MemorySize.VALUE_128_MB,
1, blockSize));
+ FormatWriterFactory writerFactory =
format.createWriterFactory(rowType);
+ LocalFileIO fileIO = new LocalFileIO();
+ FormatWriter writer =
writerFactory.create(fileIO.newOutputStream(path, false), "zstd");
+ for (GenericRow row : rows) {
+ writer.addElement(row);
+ }
+ writer.close();
+ }
+
+ private List<InternalRow> readAllWithPrefetch(
+ RowType rowType,
+ Path path,
+ List<Predicate> predicates,
+ int prefetchRowGroups,
+ List<Long> positions)
+ throws IOException {
+ FormatReaderFactory readerFactory =
+ createReaderFactory(rowType, predicates, prefetchRowGroups);
+ LocalFileIO fileIO = new LocalFileIO();
+ RecordReader<InternalRow> reader =
+ readerFactory.createReader(
+ new FormatReaderContext(
+ fileIO, path, fileIO.getFileSize(path), null,
null));
+ InternalRowSerializer serializer = new InternalRowSerializer(rowType);
+ List<InternalRow> result = new ArrayList<>();
+ RecordReader.RecordIterator<InternalRow> batch;
+ while ((batch = reader.readBatch()) != null) {
+ FileRecordIterator<InternalRow> fileIterator =
(FileRecordIterator<InternalRow>) batch;
+ InternalRow row;
+ while ((row = fileIterator.next()) != null) {
+ positions.add(fileIterator.returnedPosition());
+ result.add(serializer.copy(row));
+ }
+ batch.releaseBatch();
+ }
+ reader.close();
+ return result;
+ }
+
private Path newPath() {
return new Path(tempDir.toUri().toString(), UUID.randomUUID() +
".mosaic");
}
@@ -379,6 +491,16 @@ class MosaicReaderWriterTest {
return result;
}
+ private static FormatReaderFactory createReaderFactory(
+ RowType rowType, List<Predicate> predicates, int
prefetchRowGroups) {
+ return new MosaicReaderFactory(
+ rowType,
+ rowType,
+ predicates,
+ prefetchRowGroups,
+
MosaicFileFormat.READ_PREFETCH_MAX_BYTES.defaultValue().getBytes());
+ }
+
private static MosaicFileFormat createFormat() {
return createFormat("");
}
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 45bd36c46b..0e88b42909 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
@@ -18,6 +18,7 @@
package org.apache.paimon.format.mosaic;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.fs.Path;
import org.apache.paimon.fs.SeekableInputStream;
@@ -27,20 +28,35 @@ import org.apache.paimon.reader.FileRecordIterator;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
+import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+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 java.io.IOException;
+import java.io.InterruptedIOException;
+import java.util.ArrayList;
import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+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;
@@ -186,6 +202,242 @@ class MosaicRecordsReaderTest {
recordsReader.close();
}
+ @Test
+ void testDisabledPrefetchReadsRowGroupsOnDemand() throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = createProjectedReader(allocator, 3);
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader, 0);
+
+ FileRecordIterator<InternalRow> first = recordsReader.readBatch();
+ assertThat(first).isNotNull();
+ // Depth 0 must not read the next row group before the first batch is
consumed.
+ verify(reader, times(1)).readRowGroup(anyInt(), any());
+ assertThat(first.next().getInt(0)).isEqualTo(0);
+ assertThat(first.next()).isNull();
+ first.releaseBatch();
+
+ List<Integer> values = new ArrayList<>();
+ FileRecordIterator<InternalRow> batch;
+ while ((batch = recordsReader.readBatch()) != null) {
+ InternalRow row;
+ while ((row = batch.next()) != null) {
+ values.add(row.getInt(0));
+ }
+ batch.releaseBatch();
+ }
+ assertThat(values).containsExactly(1, 2);
+ verify(reader, times(3)).readRowGroup(anyInt(), any());
+
+ recordsReader.close();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testInterruptedReadKeepsInFlightRowGroupUntilClose() throws Exception
{
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = createProjectedReader(allocator, 1);
+ CountDownLatch readStarted = new CountDownLatch(1);
+ CountDownLatch releaseRead = new CountDownLatch(1);
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ doAnswer(
+ invocation -> {
+ readStarted.countDown();
+ releaseRead.await();
+ events.add("read-finished");
+ return rowGroup(allocator, 0);
+ })
+ .when(reader)
+ .readRowGroup(eq(0), any());
+ doAnswer(
+ invocation -> {
+ events.add("reader-closed");
+ return null;
+ })
+ .when(reader)
+ .close();
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader, 2);
+ AtomicReference<Throwable> readFailure = new AtomicReference<>();
+ Thread consumer =
+ new Thread(
+ () -> {
+ try {
+ recordsReader.readBatch();
+ } catch (Throwable t) {
+ readFailure.set(t);
+ }
+ });
+ consumer.start();
+ readStarted.await();
+ consumer.interrupt();
+ consumer.join();
+
assertThat(readFailure.get()).isInstanceOf(InterruptedIOException.class);
+
+ // The interrupted read is still running natively: close() has to wait
for it.
+ AtomicReference<Throwable> closeFailure = new AtomicReference<>();
+ Thread closer = new Thread(() -> closeQuietly(recordsReader,
closeFailure));
+ closer.start();
+ closer.join(200);
+ assertThat(closer.isAlive()).isTrue();
+ assertThat(events).isEmpty();
+ releaseRead.countDown();
+ closer.join();
+
+ assertThat(closeFailure.get()).isNull();
+ assertThat(events).containsExactly("read-finished", "reader-closed");
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testCloseWithInterruptFlagStillWaitsForInFlightRowGroups() throws
Exception {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = createProjectedReader(allocator, 2);
+ CountDownLatch readStarted = new CountDownLatch(1);
+ CountDownLatch releaseRead = new CountDownLatch(1);
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ doAnswer(
+ invocation -> {
+ readStarted.countDown();
+ releaseRead.await();
+ events.add("read-finished");
+ return rowGroup(allocator, 1);
+ })
+ .when(reader)
+ .readRowGroup(eq(1), any());
+ doAnswer(
+ invocation -> {
+ events.add("reader-closed");
+ return null;
+ })
+ .when(reader)
+ .close();
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader, 1);
+ // Consuming row group 0 schedules row group 1, which now blocks in
the background.
+ assertThat(recordsReader.readBatch()).isNotNull();
+ readStarted.await();
+
+ AtomicReference<Throwable> closeFailure = new AtomicReference<>();
+ AtomicBoolean interruptedAfterClose = new AtomicBoolean();
+ Thread closer =
+ new Thread(
+ () -> {
+ Thread.currentThread().interrupt();
+ closeQuietly(recordsReader, closeFailure);
+
interruptedAfterClose.set(Thread.currentThread().isInterrupted());
+ });
+ closer.start();
+ closer.join(200);
+ assertThat(closer.isAlive()).isTrue();
+ assertThat(events).isEmpty();
+ releaseRead.countDown();
+ closer.join();
+
+ assertThat(closeFailure.get()).isNull();
+ assertThat(events).containsExactly("read-finished", "reader-closed");
+ assertThat(interruptedAfterClose).isTrue();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testRefillFailureLeavesCurrentRowGroupReleasable() throws IOException
{
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = createProjectedReader(allocator, 3);
+ RuntimeException failure = new RuntimeException("row group 2 metadata
failed");
+ when(reader.rowGroupNumRows(2)).thenThrow(failure);
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader, 1);
+ // Row group 0 is handed over; scheduling row group 2 fails while
refilling behind it.
+ assertThat(recordsReader.readBatch()).isNotNull();
+ assertThatThrownBy(recordsReader::readBatch).isSameAs(failure);
+
+ // Row group 0 must still be released, otherwise the allocator reports
a leak here.
+ recordsReader.close();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testPrefetchIsBoundedByEstimatedDecodedBytes() 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());
+ }
+ recordsReader.close();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ }
+ }
+
+ private static void closeQuietly(
+ MosaicRecordsReader recordsReader, AtomicReference<Throwable>
failure) {
+ try {
+ recordsReader.close();
+ } catch (Throwable t) {
+ failure.set(t);
+ }
+ }
+
+ /** A mocked native reader whose file schema contains the projected
column. */
+ private static MosaicReader createProjectedReader(BufferAllocator
allocator, int numRowGroups) {
+ MosaicReader reader = mock(MosaicReader.class);
+ when(reader.getSchema())
+ .thenReturn(
+ new Schema(
+ Collections.singletonList(
+ Field.nullable("f0", new
ArrowType.Int(32, true)))));
+ when(reader.numRowGroups()).thenReturn(numRowGroups);
+ when(reader.rowGroupNumRows(anyInt())).thenReturn(1);
+ when(reader.readRowGroup(anyInt(), any()))
+ .thenAnswer(invocation -> rowGroup(allocator,
invocation.getArgument(0)));
+ return reader;
+ }
+
+ private static VectorSchemaRoot rowGroup(BufferAllocator allocator, int
value) {
+ VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(rowType(),
allocator);
+ IntVector vector = (IntVector) root.getVector(0);
+ vector.allocateNew(1);
+ vector.set(0, value);
+ vector.setValueCount(1);
+ root.setRowCount(1);
+ return root;
+ }
+
private static MosaicInputFileAdapter createInputFileAdapter(
CloseCountingSeekableInputStream inputStream) throws IOException {
return new MosaicInputFileAdapter(
@@ -207,6 +459,24 @@ class MosaicRecordsReaderTest {
(inputFile, fileSize, bufferAllocator) -> reader);
}
+ private static MosaicRecordsReader createRecordsReader(
+ MosaicInputFileAdapter inputFileAdapter,
+ CloseCountingRootAllocator allocator,
+ MosaicReader reader,
+ int prefetchRowGroups) {
+ return new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType(),
+ rowType(),
+ null,
+ new Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator) -> reader,
+ prefetchRowGroups,
+
MosaicFileFormat.READ_PREFETCH_MAX_BYTES.defaultValue().getBytes());
+ }
+
private static MosaicReader createReader() {
MosaicReader reader = mock(MosaicReader.class);
when(reader.getSchema()).thenReturn(new
Schema(Collections.emptyList()));
