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 a11ce79d12 [mosaic] Fix a close race in the reader stream pool and add
pool and writer sizing tests (#9744)
a11ce79d12 is described below
commit a11ce79d122844aefa87e6bd429cfdf046b5893e
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Fri Sep 11 23:10:10 2026 +0800
[mosaic] Fix a close race in the reader stream pool and add pool and writer
sizing tests (#9744)
---
.../format/mosaic/MosaicInputFileAdapter.java | 8 +-
.../paimon/format/mosaic/MosaicRecordsReader.java | 3 +-
.../paimon/format/mosaic/MosaicRecordsWriter.java | 2 +-
.../format/mosaic/MosaicInputFileAdapterTest.java | 103 ++++++++++++++++++++-
.../format/mosaic/MosaicReaderWriterTest.java | 40 ++++++++
.../format/mosaic/MosaicRecordsWriterTest.java | 100 ++++++++++++++------
6 files changed, 225 insertions(+), 31 deletions(-)
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 07385fdd36..27e18f75e3 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
@@ -33,7 +33,8 @@ import java.util.ArrayList;
import java.util.List;
/**
- * Adapter that exposes a Paimon {@link SeekableInputStream} as a Mosaic
{@link InputFile}.
+ * Adapter that exposes a Paimon file as a Mosaic {@link InputFile} through a
small pool of {@link
+ * SeekableInputStream}s.
*
* <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.
@@ -119,19 +120,22 @@ public class MosaicInputFileAdapter implements InputFile,
Closeable {
}
}
SeekableInputStream opened = null;
+ boolean closedMeanwhile;
try {
opened = fileIO.newInputStream(path);
} finally {
synchronized (this) {
openingStreams--;
+ closedMeanwhile = closed;
if (opened != null && !closed) {
allStreams.add(opened);
} else {
+ // The slot is free again for another borrower.
notifyAll();
}
}
}
- if (closed) {
+ if (closedMeanwhile) {
opened.close();
throw new IOException("Input file " + path + " is closed");
}
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 c36285baea..0f32cb6b28 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
@@ -438,7 +438,7 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
this.future = future;
}
- /** Waits for the data; the batch is not released here even when the
wait fails. */
+ /** Waits for the data; a failed wait leaves the batch queued so
close() still drains it. */
@Nullable
VectorSchemaRoot await() throws IOException {
if (future == null) {
@@ -487,6 +487,7 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
// The native read still uses the reader handle; it must
complete first.
interrupted = true;
} catch (ExecutionException e) {
+ // A failed read holds no data to release.
return interrupted;
}
}
diff --git
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java
index f4a96937ce..df23bf833b 100644
---
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java
+++
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java
@@ -125,7 +125,7 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
private static void setInitialCapacity(FieldVector vector, int capacity) {
if (vector instanceof BaseRepeatedValueVector) {
- // The plain overload would size the element vector for 5 elements
per row.
+ // Avoid Arrow's 5x estimate for fixed- or variable-width element
vectors.
((BaseRepeatedValueVector) vector).setInitialCapacity(capacity,
1.0);
} else {
vector.setInitialCapacity(capacity);
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
index 6a0e4ffb64..9b00ea57ee 100644
---
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
@@ -82,6 +82,101 @@ class MosaicInputFileAdapterTest {
adapter.close();
}
+ @Test
+ void testFailedExtraStreamOpenReleasesItsSlot() throws Exception {
+ CountDownLatch readsStarted = new CountDownLatch(1);
+ CountDownLatch releaseReads = new CountDownLatch(1);
+ CountDownLatch thirdOpen = new CountDownLatch(1);
+ AtomicInteger opens = new AtomicInteger();
+ LocalFileIO fileIO =
+ new LocalFileIO() {
+ @Override
+ public SeekableInputStream newInputStream(Path path)
throws IOException {
+ // The second open (the first extra stream) fails once.
+ int open = opens.incrementAndGet();
+ if (open == 2) {
+ throw new IOException("open failed");
+ }
+ if (open == 3) {
+ thirdOpen.countDown();
+ }
+ return new BlockingStream(readsStarted, releaseReads);
+ }
+ };
+ MosaicInputFileAdapter adapter =
+ new MosaicInputFileAdapter(fileIO, new
Path("file:/tmp/mosaic-adapter-test"), 2);
+ AtomicReference<Throwable> holderFailure = new AtomicReference<>();
+ Thread holder = new Thread(() -> read(adapter, holderFailure));
+ holder.start();
+ readsStarted.await();
+
+ // The failed open must not keep the second slot reserved.
+ AtomicReference<Throwable> failed = new AtomicReference<>();
+ read(adapter, failed);
+
assertThat(failed.get()).isInstanceOf(IOException.class).hasMessage("open
failed");
+ CountDownLatch secondRead = new CountDownLatch(1);
+ AtomicReference<Throwable> retryFailure = new AtomicReference<>();
+ Thread retry =
+ new Thread(
+ () -> {
+ read(adapter, retryFailure);
+ secondRead.countDown();
+ });
+ retry.start();
+ // The retry must open its own stream while the first one is still
held.
+ thirdOpen.await();
+ releaseReads.countDown();
+ retry.join();
+ holder.join();
+ assertThat(retryFailure.get()).isNull();
+ assertThat(holderFailure.get()).isNull();
+ assertThat(opens.get()).isEqualTo(3);
+ adapter.close();
+ }
+
+ @Test
+ void testCloseClosesEveryStreamAndRejectsLaterReads() throws Exception {
+ CountDownLatch readsStarted = new CountDownLatch(2);
+ CountDownLatch releaseReads = new CountDownLatch(1);
+ List<BlockingStream> streams = new ArrayList<>();
+ LocalFileIO fileIO =
+ new LocalFileIO() {
+ @Override
+ public SeekableInputStream newInputStream(Path path) {
+ BlockingStream stream = new
BlockingStream(readsStarted, releaseReads);
+ synchronized (streams) {
+ streams.add(stream);
+ }
+ return stream;
+ }
+ };
+ MosaicInputFileAdapter adapter =
+ new MosaicInputFileAdapter(fileIO, new
Path("file:/tmp/mosaic-adapter-test"), 3);
+ List<Thread> readers = new ArrayList<>();
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+ for (int i = 0; i < 2; i++) {
+ Thread thread = new Thread(() -> read(adapter, failure));
+ thread.start();
+ readers.add(thread);
+ }
+ readsStarted.await();
+ releaseReads.countDown();
+ for (Thread thread : readers) {
+ thread.join();
+ }
+ assertThat(failure.get()).isNull();
+ assertThat(streams).hasSize(2);
+
+ adapter.close();
+ assertThat(streams).allMatch(stream -> stream.closeCount == 1);
+ AtomicReference<Throwable> afterClose = new AtomicReference<>();
+ read(adapter, afterClose);
+ assertThat(afterClose.get()).isInstanceOf(IOException.class);
+ // Closing again is a no-op.
+ adapter.close();
+ assertThat(streams).allMatch(stream -> stream.closeCount == 1);
+ }
+
@Test
void testCloseWakesWaitingReader() throws Exception {
CountDownLatch readsStarted = new CountDownLatch(1);
@@ -110,6 +205,8 @@ class MosaicInputFileAdapterTest {
assertThat(second.get()).isInstanceOf(IOException.class);
releaseReads.countDown();
holder.join();
+ // The read that held the stream completes; the stream is closed by
close().
+ assertThat(first.get()).isNull();
}
private static void read(MosaicInputFileAdapter adapter,
AtomicReference<Throwable> failure) {
@@ -157,6 +254,10 @@ class MosaicInputFileAdapterTest {
}
@Override
- public void close() {}
+ public void close() {
+ closeCount++;
+ }
+
+ private volatile int closeCount;
}
}
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 aceca5c455..14553219ce 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
@@ -31,6 +31,7 @@ import org.apache.paimon.format.FormatReaderFactory;
import org.apache.paimon.format.FormatWriter;
import org.apache.paimon.format.FormatWriterFactory;
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.options.MemorySize;
@@ -57,6 +58,7 @@ import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
@@ -484,6 +486,44 @@ class MosaicReaderWriterTest {
assertThatCode(reader::close).doesNotThrowAnyException();
}
+ @Test
+ void testReaderOpensAtMostDepthPlusOneStreams() 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);
+
+ AtomicInteger opened = new AtomicInteger();
+ LocalFileIO fileIO =
+ new LocalFileIO() {
+ @Override
+ public SeekableInputStream newInputStream(Path file)
throws IOException {
+ opened.incrementAndGet();
+ return super.newInputStream(file);
+ }
+ };
+ FormatReaderFactory readerFactory = createReaderFactory(rowType, null,
3);
+ int count = 0;
+ try (RecordReader<InternalRow> reader =
+ readerFactory.createReader(
+ new FormatReaderContext(
+ fileIO, path, fileIO.getFileSize(path), null,
null))) {
+ RecordReader.RecordIterator<InternalRow> batch;
+ while ((batch = reader.readBatch()) != null) {
+ while (batch.next() != null) {
+ count++;
+ }
+ batch.releaseBatch();
+ }
+ }
+ assertThat(count).isEqualTo(rows.length);
+ // One stream per row group being opened plus one for the consumer.
+ assertThat(opened.get()).isBetween(1, 4);
+ }
+
private void writeRows(
RowType rowType, Path path, Options options, MemorySize blockSize,
GenericRow... rows)
throws IOException {
diff --git
a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java
index be69fb6207..cc1f78cfdf 100644
---
a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java
+++
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java
@@ -38,7 +38,9 @@ import org.apache.arrow.vector.VectorSchemaRoot;
import org.junit.jupiter.api.Test;
import java.io.ByteArrayOutputStream;
+import java.util.Arrays;
import java.util.Collections;
+import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -210,37 +212,74 @@ class MosaicRecordsWriterTest {
void testInitialCapacityNeverExceedsArrowDefaultAllocation() throws
Exception {
// DOUBLE, STRING and ARRAY<DOUBLE> columns cover the fixed-width,
variable-width and
// repeated vector families, whose default sizing differs.
- RowType rowType = mixedRowType(60);
- for (int batchSize : new int[] {4, 1024, 3969, 3970, 65536}) {
- long baseline = baselineFirstRowAllocation(rowType, batchSize);
- MosaicWriter nativeWriter = mock(MosaicWriter.class);
- try (RootAllocator allocator = new RootAllocator()) {
- MosaicRecordsWriter writer =
- createWriter(
- rowType,
- allocator,
- nativeWriter,
- batchSize,
- MemorySize.VALUE_128_MB);
- try {
- writer.addElement(firstRow(rowType));
- assertThat(allocator.getAllocatedMemory())
- .as("batch size %d", batchSize)
- .isLessThanOrEqualTo(baseline);
- // Arrow rounds buffers to powers of two, so only clearly
smaller batches
- // allocate less.
- if (batchSize <= 1024) {
+ for (RowType rowType : familyRowTypes()) {
+ for (int batchSize : new int[] {4, 1024, 3969, 3970, 65536}) {
+ long baseline = baselineFirstRowAllocation(rowType, batchSize);
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ try (RootAllocator allocator = new RootAllocator()) {
+ MosaicRecordsWriter writer =
+ createWriter(
+ rowType,
+ allocator,
+ nativeWriter,
+ batchSize,
+ MemorySize.VALUE_128_MB);
+ try {
+ writer.addElement(firstRow(rowType));
assertThat(allocator.getAllocatedMemory())
.as("batch size %d", batchSize)
- .isLessThan(baseline);
+ .isLessThanOrEqualTo(baseline);
+ // Arrow rounds buffers to powers of two, so only
clearly smaller batches
+ // allocate less.
+ if (batchSize <= 1024) {
+ assertThat(allocator.getAllocatedMemory())
+ .as("batch size %d", batchSize)
+ .isLessThan(baseline);
+ }
+ } finally {
+ writer.close();
}
- } finally {
- writer.close();
}
}
}
}
+ @Test
+ void testFillingOneBatchDoesNotReallocate() throws Exception {
+ // Fixed-width vectors sized for the batch must hold a full batch
without growing.
+ RowType.Builder builder = RowType.builder();
+ for (int i = 0; i < 100; i++) {
+ builder.field("d" + i, DataTypes.DOUBLE());
+ }
+ RowType rowType = builder.build();
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ try (RootAllocator allocator = new RootAllocator()) {
+ MosaicRecordsWriter writer =
+ createWriter(rowType, allocator, nativeWriter, 1024,
MemorySize.VALUE_128_MB);
+ writer.addElement(firstRow(rowType));
+ long afterFirstRow = allocator.getAllocatedMemory();
+ for (int i = 1; i < 1024; i++) {
+ writer.addElement(firstRow(rowType));
+ }
+
assertThat(allocator.getAllocatedMemory()).isEqualTo(afterFirstRow);
+ verify(nativeWriter, never()).write(any());
+ writer.close();
+ }
+ }
+
+ /** The mixed schema plus one schema per vector family, so a skipped
family shows. */
+ private static List<RowType> familyRowTypes() {
+ RowType.Builder doubles = RowType.builder();
+ RowType.Builder strings = RowType.builder();
+ RowType.Builder arrays = RowType.builder();
+ for (int i = 0; i < 60; i++) {
+ doubles.field("d" + i, DataTypes.DOUBLE());
+ strings.field("s" + i, DataTypes.STRING());
+ arrays.field("a" + i, DataTypes.ARRAY(DataTypes.DOUBLE()));
+ }
+ return Arrays.asList(mixedRowType(60), doubles.build(),
strings.build(), arrays.build());
+ }
+
private static RowType mixedRowType(int columnsPerType) {
RowType.Builder builder = RowType.builder();
for (int i = 0; i < columnsPerType; i++) {
@@ -253,9 +292,18 @@ class MosaicRecordsWriterTest {
private static GenericRow firstRow(RowType rowType) {
GenericRow row = new GenericRow(rowType.getFieldCount());
- row.setField(0, 1.0d);
- row.setField(1, BinaryString.fromString("one"));
- row.setField(2, new GenericArray(new Object[] {1.0d}));
+ for (int i = 0; i < rowType.getFieldCount(); i++) {
+ switch (rowType.getTypeAt(i).getTypeRoot()) {
+ case DOUBLE:
+ row.setField(i, 1.0d);
+ break;
+ case VARCHAR:
+ row.setField(i, BinaryString.fromString("one"));
+ break;
+ default:
+ row.setField(i, new GenericArray(new Object[] {1.0d}));
+ }
+ }
return row;
}
