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 7246bdcb03 [mosaic] Expose column-batch reads and writer extension
hooks (#10516)
7246bdcb03 is described below
commit 7246bdcb03ac9717c60506487286a1f4ef60db02
Author: jianguotian <[email protected]>
AuthorDate: Mon Oct 12 21:39:56 2026 +0800
[mosaic] Expose column-batch reads and writer extension hooks (#10516)
---
.../paimon/arrow/reader/ArrowBatchReader.java | 68 ++-
.../paimon/arrow/reader/ArrowBatchReaderTest.java | 89 +++
.../paimon/format/mosaic/MosaicRecordsReader.java | 232 ++++---
.../paimon/format/mosaic/MosaicRecordsWriter.java | 68 ++-
.../paimon/format/mosaic/MosaicWriterFactory.java | 13 +-
.../format/mosaic/MosaicColumnBatchReaderTest.java | 668 +++++++++++++++++++++
.../format/mosaic/MosaicReaderWriterTest.java | 51 ++
.../format/mosaic/MosaicRecordsWriterTest.java | 68 +++
.../format/mosaic/MosaicRowTrackingTest.java | 208 +++++++
.../extension/MosaicWriterExtensionTest.java | 150 +++++
10 files changed, 1503 insertions(+), 112 deletions(-)
diff --git
a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java
index ab939d70b4..5548dfb2fc 100644
---
a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java
+++
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java
@@ -41,7 +41,7 @@ import static
org.apache.paimon.utils.StringUtils.toLowerCaseIfNeed;
/** Reader from a {@link VectorSchemaRoot} to paimon rows. */
public class ArrowBatchReader {
- private final VectorizedColumnBatch batch;
+ private final VectorizedColumnBatch reusableBatch;
private final Arrow2PaimonVectorConverter[] convertors;
private final RowType projectedRowType;
private final boolean caseSensitive;
@@ -57,12 +57,11 @@ public class ArrowBatchReader {
RowType rowType,
boolean caseSensitive,
Arrow2PaimonVectorConverter.Arrow2PaimonVectorConvertorVisitor
visitor) {
- ColumnVector[] columnVectors = new
ColumnVector[rowType.getFieldCount()];
+ this.reusableBatch = new VectorizedColumnBatch(new
ColumnVector[rowType.getFieldCount()]);
this.convertors = new
Arrow2PaimonVectorConverter[rowType.getFieldCount()];
- this.batch = new VectorizedColumnBatch(columnVectors);
this.projectedRowType = rowType;
- for (int i = 0; i < columnVectors.length; i++) {
+ for (int i = 0; i < convertors.length; i++) {
this.convertors[i] =
Arrow2PaimonVectorConverter.construct(visitor,
rowType.getTypeAt(i));
}
@@ -70,6 +69,41 @@ public class ArrowBatchReader {
}
public Iterable<InternalRow> readBatch(VectorSchemaRoot vsr) {
+ populateBatch(vsr, reusableBatch);
+ int rowCount = reusableBatch.getNumRows();
+ final ColumnarRow columnarRow = new ColumnarRow(reusableBatch);
+ return () ->
+ new Iterator<InternalRow>() {
+ private int position = 0;
+
+ @Override
+ public boolean hasNext() {
+ return position < rowCount;
+ }
+
+ @Override
+ public InternalRow next() {
+ columnarRow.setRowId(position);
+ position++;
+ return columnarRow;
+ }
+ };
+ }
+
+ /**
+ * Wraps an Arrow batch as Paimon column vectors without materializing
rows.
+ *
+ * <p>The returned batch container is not reused, but its columns borrow
vectors owned by {@code
+ * vsr} and must not be used after the root is released.
+ */
+ public VectorizedColumnBatch readVectorizedBatch(VectorSchemaRoot vsr) {
+ VectorizedColumnBatch resultBatch =
+ new VectorizedColumnBatch(new
ColumnVector[projectedRowType.getFieldCount()]);
+ populateBatch(vsr, resultBatch);
+ return resultBatch;
+ }
+
+ private void populateBatch(VectorSchemaRoot vsr, VectorizedColumnBatch
targetBatch) {
int[] mapping = new int[projectedRowType.getFieldCount()];
Schema arrowSchema = vsr.getSchema();
Map<String, Integer> arrowFieldIndex = new HashMap<>();
@@ -83,32 +117,14 @@ public class ArrowBatchReader {
mapping[i] = arrowFieldIndex.getOrDefault(fieldName, -1);
}
- for (int i = 0; i < batch.columns.length; i++) {
+ for (int i = 0; i < targetBatch.columns.length; i++) {
if (mapping[i] >= 0) {
- batch.columns[i] =
convertors[i].convertVector(vsr.getVector(mapping[i]));
+ targetBatch.columns[i] =
convertors[i].convertVector(vsr.getVector(mapping[i]));
} else {
- batch.columns[i] = AllNullColumnVector.INSTANCE;
+ targetBatch.columns[i] = AllNullColumnVector.INSTANCE;
}
}
- int rowCount = vsr.getRowCount();
- batch.setNumRows(vsr.getRowCount());
- final ColumnarRow columnarRow = new ColumnarRow(batch);
- return () ->
- new Iterator<InternalRow>() {
- private int position = 0;
-
- @Override
- public boolean hasNext() {
- return position < rowCount;
- }
-
- @Override
- public InternalRow next() {
- columnarRow.setRowId(position);
- position++;
- return columnarRow;
- }
- };
+ targetBatch.setNumRows(vsr.getRowCount());
}
}
diff --git
a/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java
b/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java
new file mode 100644
index 0000000000..c70f21ad44
--- /dev/null
+++
b/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java
@@ -0,0 +1,89 @@
+/*
+ * 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.arrow.reader;
+
+import org.apache.paimon.arrow.ArrowUtils;
+import org.apache.paimon.data.columnar.ColumnarRow;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link ArrowBatchReader}. */
+class ArrowBatchReaderTest {
+
+ @Test
+ void testReadBatchWrapperIsReused() {
+ RowType rowType = RowType.builder().field("id",
DataTypes.INT()).build();
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot firstRoot = intRoot(rowType, allocator, 11);
+ VectorSchemaRoot secondRoot = intRoot(rowType, allocator, 22))
{
+ ArrowBatchReader reader = new ArrowBatchReader(rowType, true);
+
+ ColumnarRow first = (ColumnarRow)
reader.readBatch(firstRoot).iterator().next();
+ VectorizedColumnBatch reusableBatch = first.batch();
+ assertThat(first.getInt(0)).isEqualTo(11);
+
+ ColumnarRow second = (ColumnarRow)
reader.readBatch(secondRoot).iterator().next();
+
+ assertThat(second.batch()).isSameAs(reusableBatch);
+ assertThat(second.getInt(0)).isEqualTo(22);
+ }
+ }
+
+ @Test
+ void testVectorizedBatchWrappersAreNotReused() {
+ RowType rowType = RowType.builder().field("id",
DataTypes.INT()).build();
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot firstRoot = intRoot(rowType, allocator, 11);
+ VectorSchemaRoot secondRoot = intRoot(rowType, allocator, 22,
33)) {
+ ArrowBatchReader reader = new ArrowBatchReader(rowType, true);
+
+ VectorizedColumnBatch first =
reader.readVectorizedBatch(firstRoot);
+ VectorizedColumnBatch second =
reader.readVectorizedBatch(secondRoot);
+
+ assertThat(first).isNotSameAs(second);
+ assertThat(first.columns).isNotSameAs(second.columns);
+ assertThat(first.getNumRows()).isEqualTo(1);
+ assertThat(second.getNumRows()).isEqualTo(2);
+ assertThat(first.getInt(0, 0)).isEqualTo(11);
+ assertThat(second.getInt(0, 0)).isEqualTo(22);
+ assertThat(second.getInt(1, 0)).isEqualTo(33);
+ }
+ }
+
+ private static VectorSchemaRoot intRoot(
+ RowType rowType, 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;
+ }
+}
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 0f32cb6b28..f0e1018b90 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
@@ -22,12 +22,18 @@ import org.apache.paimon.arrow.reader.ArrowBatchReader;
import org.apache.paimon.data.GenericArray;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.columnar.ColumnVector;
+import org.apache.paimon.data.columnar.ColumnarRow;
+import org.apache.paimon.data.columnar.ColumnarRowIterator;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
+import org.apache.paimon.data.columnar.VectorizedRowIterator;
import org.apache.paimon.fs.Path;
import org.apache.paimon.mosaic.ColumnStatistics;
import org.apache.paimon.mosaic.MosaicReader;
import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.reader.FileRecordIterator;
import org.apache.paimon.reader.FileRecordReader;
+import org.apache.paimon.table.SpecialFields;
import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataType;
import org.apache.paimon.types.RowType;
@@ -49,7 +55,6 @@ 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;
@@ -60,6 +65,7 @@ import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;
import static org.apache.paimon.format.mosaic.MosaicObjects.convertStatsValue;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
/** File reader for Mosaic format. */
public class MosaicRecordsReader implements FileRecordReader<InternalRow> {
@@ -88,8 +94,8 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
private int nextRowGroupToSchedule;
private long scheduledRowCount;
- private long returnedPosition = -1;
- private VectorSchemaRoot currentVsr;
+ private final Object batchLock = new Object();
+ private final List<BatchRecycler> activeBatches = new ArrayList<>();
private static final Logger LOG =
LoggerFactory.getLogger(MosaicRecordsReader.class);
@@ -299,73 +305,47 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
@Nullable
@Override
public FileRecordIterator<InternalRow> readBatch() throws IOException {
- releaseCurrentVsr();
-
- RowGroupBatch batch = nextRowGroup();
+ if (pending.isEmpty()) {
+ fillPrefetchQueue(Math.max(1, prefetchDepth), true);
+ }
+ RowGroupBatch batch = pending.peek();
if (batch == null) {
return null;
}
- // Rows of skipped row groups still count towards the file position.
- returnedPosition = batch.startPosition - 1;
+ // Leave a failed or interrupted read queued so close() waits for it.
+ long waitStart = System.nanoTime();
+ VectorSchemaRoot vsr = batch.await();
+ openNanos += System.nanoTime() - waitStart;
+ openedRowGroups++;
+ pending.poll();
+ pendingBytes -= batch.bytes;
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;
+ // Register the borrowed batch before refilling, which can fail while
scheduling a read.
+ BatchRecycler recycler = vsr == null ? null : recycler(vsr);
+ try {
+ if (prefetchDepth > 0) {
+ fillPrefetchQueue(prefetchDepth, false);
}
-
- @Override
- public Path filePath() {
- return filePath;
+ if (allProjectedColumnsMissing) {
+ return allNullIterator(batch.startPosition, batch.numRows);
}
-
- @Nullable
- @Override
- public InternalRow next() {
- if (rows.hasNext()) {
- returnedPosition++;
- return rows.next();
+ return new MosaicVectorizedRecordIterator(
+ filePath,
+ batch.startPosition,
+ arrowBatchReader.readVectorizedBatch(vsr),
+ recycler);
+ } catch (Throwable t) {
+ try {
+ if (recycler != null) {
+ recycler.run();
}
- return null;
+ } catch (Throwable closeFailure) {
+ t.addSuppressed(closeFailure);
}
-
- @Override
- public void releaseBatch() {
- releaseCurrentVsr();
- }
- };
- }
-
- /** 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), true);
- }
- RowGroupBatch head = pending.peek();
- if (head == null) {
+ rethrow(t);
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, false);
- }
- return head;
}
/** Schedules matching row groups until {@code wanted} are queued or the
byte budget is used. */
@@ -494,14 +474,18 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
}
}
- private FileRecordIterator<InternalRow> allNullIterator(int numRows) {
+ private FileRecordIterator<InternalRow> allNullIterator(long
batchStartPosition, int numRows) {
GenericRow row = new GenericRow(projectedFieldCount);
return new FileRecordIterator<InternalRow>() {
private int position;
+ private long batchReturnedPosition = -1;
@Override
public long returnedPosition() {
- return returnedPosition;
+ if (batchReturnedPosition < 0) {
+ throw new IllegalStateException("returnedPosition() is
called before next()");
+ }
+ return batchReturnedPosition;
}
@Override
@@ -513,8 +497,8 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
@Override
public InternalRow next() {
if (position < numRows) {
+ batchReturnedPosition = batchStartPosition + position;
position++;
- returnedPosition++;
return row;
}
return null;
@@ -566,10 +550,37 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
return true;
}
- private void releaseCurrentVsr() {
- if (currentVsr != null) {
- currentVsr.close();
- currentVsr = null;
+ private BatchRecycler recycler(VectorSchemaRoot vsr) {
+ try {
+ BatchRecycler recycler = new BatchRecycler(vsr);
+ synchronized (batchLock) {
+ activeBatches.add(recycler);
+ }
+ return recycler;
+ } catch (Throwable t) {
+ try {
+ vsr.close();
+ } catch (Throwable closeFailure) {
+ t.addSuppressed(closeFailure);
+ }
+ throw rethrowUnchecked(t);
+ }
+ }
+
+ private void releaseActiveBatches() {
+ Throwable failure = null;
+ synchronized (batchLock) {
+ while (!activeBatches.isEmpty()) {
+ BatchRecycler recycler =
activeBatches.get(activeBatches.size() - 1);
+ try {
+ recycler.releaseLocked();
+ } catch (Throwable t) {
+ failure = addSuppressed(failure, t);
+ }
+ }
+ }
+ if (failure != null) {
+ throw rethrowUnchecked(failure);
}
}
@@ -578,7 +589,7 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
Throwable throwable = null;
try {
- releaseCurrentVsr();
+ releaseActiveBatches();
} catch (Throwable t) {
throwable = t;
}
@@ -602,10 +613,12 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
throwable = addSuppressed(throwable, t);
}
- try {
- allocator.close();
- } catch (Throwable t) {
- throwable = addSuppressed(throwable, t);
+ synchronized (batchLock) {
+ try {
+ allocator.close();
+ } catch (Throwable t) {
+ throwable = addSuppressed(throwable, t);
+ }
}
try {
@@ -695,6 +708,81 @@ public class MosaicRecordsReader implements
FileRecordReader<InternalRow> {
interface NativeReaderOpener {
MosaicReader open(
- MosaicInputFileAdapter inputFileAdapter, long fileSize,
BufferAllocator allocator);
+ MosaicInputFileAdapter inputFileAdapter, long fileSize,
BufferAllocator allocator)
+ throws IOException;
+ }
+
+ private class BatchRecycler implements Runnable {
+
+ private final VectorSchemaRoot vsr;
+ private boolean released;
+
+ private BatchRecycler(VectorSchemaRoot vsr) {
+ this.vsr = vsr;
+ }
+
+ @Override
+ public void run() {
+ synchronized (batchLock) {
+ releaseLocked();
+ }
+ }
+
+ private void releaseLocked() {
+ if (released) {
+ return;
+ }
+
+ released = true;
+ activeBatches.remove(this);
+ vsr.close();
+ }
+ }
+
+ private static class MosaicVectorizedRecordIterator extends
VectorizedRowIterator {
+
+ private final BatchRecycler batchRecycler;
+ private final long startPosition;
+
+ private MosaicVectorizedRecordIterator(
+ Path filePath,
+ long startPosition,
+ VectorizedColumnBatch batch,
+ BatchRecycler recycler) {
+ super(filePath, new ColumnarRow(batch), recycler);
+ this.batchRecycler = recycler;
+ this.startPosition = startPosition;
+ reset(startPosition);
+ }
+
+ @Override
+ public MosaicVectorizedRecordIterator copy(ColumnVector[] vectors) {
+ checkArgument(returnedPositionIndex == 0, "copy() should not be
called after next()");
+ // Projection must retain the format's tracking fallback and batch
ownership.
+ MosaicVectorizedRecordIterator copied =
+ new MosaicVectorizedRecordIterator(
+ filePath, startPosition, batch().copy(vectors),
batchRecycler);
+ copied.reset(positionIterator);
+ return copied;
+ }
+
+ @Override
+ public ColumnarRowIterator assignRowTracking(
+ Long firstRowId, Long snapshotId, Map<String, Integer> meta) {
+ ColumnVector[] columns = batch().columns.clone();
+ VectorizedColumnBatch copiedBatch = batch().copy(columns);
+ ColumnarRowIterator fallback;
+ if (firstRowId != null &&
meta.containsKey(SpecialFields.ROW_ID.name())) {
+ fallback =
+ new ColumnarRowIterator(
+ filePath, new ColumnarRow(copiedBatch),
batchRecycler);
+ } else {
+ fallback =
+ new VectorizedRowIterator(
+ filePath, new ColumnarRow(copiedBatch),
batchRecycler);
+ }
+ fallback.reset(startPosition);
+ return fallback.assignRowTracking(firstRowId, snapshotId, meta);
+ }
}
}
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 df23bf833b..101e9842a2 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
@@ -47,6 +47,8 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Map;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+
/** Mosaic records writer. */
public class MosaicRecordsWriter implements BundleFormatWriter {
@@ -54,6 +56,7 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
private final MosaicWriter nativeWriter;
private final BufferAllocator allocator;
private final List<String> statsColumnNames;
+ private boolean failed;
@Nullable private MosaicWriterMetadata metadata;
public MosaicRecordsWriter(
@@ -101,6 +104,10 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
ArrowFormatWriter createdArrowWriter = null;
MosaicWriter createdNativeWriter = null;
try {
+ checkArgument(
+ writeBatchSize > 0,
+ "'write.batch-size' must be greater than 0, but was %s.",
+ writeBatchSize);
createdArrowWriter =
ArrowFormatWriter.forBorrowedAllocator(
rowType, writeBatchSize, true, allocator,
writeBatchMemory);
@@ -134,6 +141,7 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
@Override
public void addElement(InternalRow internalRow) {
+ checkNotFailed();
if (!arrowFormatWriter.write(internalRow)) {
flush();
if (!arrowFormatWriter.write(internalRow)) {
@@ -144,24 +152,26 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
@Override
public void writeBundle(BundleRecords bundleRecords) {
+ checkNotFailed();
if (bundleRecords instanceof ArrowBundleRecords) {
ArrowBundleRecords arrowBundle = (ArrowBundleRecords)
bundleRecords;
VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot();
- // Mosaic exports the borrowed vectors through the writer
allocator, so direct writes
- // require every source vector to share its root; otherwise
preserve semantics via rows.
if (arrowFormatWriter.isArrowBundleSchemaCompatible(arrowBundle)
&& ArrowUtils.hasSameRootAllocator(root, allocator)) {
- flush();
- nativeWriter.write(root);
+ writeArrowBatch(root);
return;
}
}
-
for (InternalRow row : bundleRecords) {
addElement(row);
}
}
+ /** The conversion buffer owned and closed by this writer. */
+ protected final ArrowFormatWriter arrowWriter() {
+ return arrowFormatWriter;
+ }
+
@Override
public boolean reachTargetSize(boolean suggestedCheck, long targetSize) {
if (!suggestedCheck) {
@@ -174,10 +184,12 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
public void close() throws IOException {
Throwable throwable = null;
- try {
- flush();
- } catch (Throwable t) {
- throwable = t;
+ if (!failed) {
+ try {
+ flush();
+ } catch (Throwable t) {
+ throwable = t;
+ }
}
try {
@@ -224,13 +236,42 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
this.metadata = new MosaicWriterMetadata(numRowGroups, allStats,
statsColumnNames);
}
- private void flush() {
+ /**
+ * Write a borrowed batch synchronously after flushing buffered rows. The
caller must validate
+ * schema and allocator compatibility for the native implementation it
uses. Ownership remains
+ * with the caller, which must keep the batch alive until this method
returns.
+ */
+ private void writeArrowBatch(VectorSchemaRoot root) {
+ checkNotFailed();
+ flush();
+ try {
+ nativeWriter.write(root);
+ } catch (RuntimeException | Error e) {
+ failed = true;
+ throw e;
+ }
+ }
+
+ protected final void flush() {
+ if (arrowFormatWriter.empty()) {
+ return;
+ }
arrowFormatWriter.flush();
- if (!arrowFormatWriter.empty()) {
- VectorSchemaRoot vsr = arrowFormatWriter.getVectorSchemaRoot();
+ VectorSchemaRoot vsr = arrowFormatWriter.getVectorSchemaRoot();
+ try {
nativeWriter.write(vsr);
+ } catch (RuntimeException | Error e) {
+ failed = true;
+ throw e;
+ } finally {
+ arrowFormatWriter.reset();
+ }
+ }
+
+ protected final void checkNotFailed() {
+ if (failed) {
+ throw new IllegalStateException("Mosaic writer has failed");
}
- arrowFormatWriter.reset();
}
private static Throwable addSuppressed(Throwable throwable, Throwable
suppressed) {
@@ -292,6 +333,7 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
}
}
+ /** Creates the native writer using the allocator owned by the records
writer. */
@FunctionalInterface
interface NativeWriterFactory {
diff --git
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicWriterFactory.java
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicWriterFactory.java
index ca67647f8c..8e85eff328 100644
---
a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicWriterFactory.java
+++
b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicWriterFactory.java
@@ -26,6 +26,7 @@ import org.apache.paimon.types.RowType;
import javax.annotation.Nullable;
+import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
@@ -59,7 +60,17 @@ public class MosaicWriterFactory implements
FormatWriterFactory {
@Override
public FormatWriter create(PositionOutputStream out, String compression) {
validateCompression(compression);
- return new MosaicRecordsWriter(out, rowType, formatContext,
statsColumnNames, numBuckets);
+ return createWriter(out, rowType, formatContext, statsColumnNames,
numBuckets);
+ }
+
+ /** Create the records writer while retaining the factory's option
validation. */
+ protected MosaicRecordsWriter createWriter(
+ OutputStream out,
+ RowType type,
+ FileFormatFactory.FormatContext context,
+ List<String> statsColumns,
+ @Nullable Integer buckets) {
+ return new MosaicRecordsWriter(out, type, context, statsColumns,
buckets);
}
private static void validateCompression(String compression) {
diff --git
a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicColumnBatchReaderTest.java
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicColumnBatchReaderTest.java
new file mode 100644
index 0000000000..1bbea3cbfd
--- /dev/null
+++
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicColumnBatchReaderTest.java
@@ -0,0 +1,668 @@
+/*
+ * 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.arrow.ArrowUtils;
+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.reader.VectorizedRecordIterator;
+import org.apache.paimon.table.SpecialFields;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+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.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/** Test for {@link MosaicRecordsReader}. */
+class MosaicColumnBatchReaderTest {
+
+ @Test
+ void testConstructorRuntimeExceptionClosesCreatedResources() throws
IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ RuntimeException failure = new RuntimeException("native reader
failed");
+
+ assertThatThrownBy(
+ () ->
+ new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType(),
+ rowType(),
+ null,
+ new
Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator)
-> {
+ throw failure;
+ }))
+ .isSameAs(failure);
+
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ assertThat(inputStream.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testConstructorIOExceptionClosesCreatedResources() throws IOException
{
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ IOException failure = new IOException("native reader failed");
+
+ assertThatThrownBy(
+ () ->
+ new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType(),
+ rowType(),
+ null,
+ new
Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator)
-> {
+ throw failure;
+ }))
+ .isInstanceOf(RuntimeException.class)
+ .hasCause(failure);
+
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ assertThat(inputStream.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testConstructorErrorClosesCreatedResources() throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ UnsatisfiedLinkError failure = new UnsatisfiedLinkError("native
library failed");
+
+ assertThatThrownBy(
+ () ->
+ new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType(),
+ rowType(),
+ null,
+ new
Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator)
-> {
+ throw failure;
+ }))
+ .isSameAs(failure);
+
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ assertThat(inputStream.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void
testConstructorFailureAfterReaderCreatedClosesReaderAndOtherResources()
+ throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = mock(MosaicReader.class);
+ RuntimeException failure = new RuntimeException("schema failed");
+ doThrow(failure).when(reader).getSchema();
+
+ assertThatThrownBy(
+ () ->
+ new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType(),
+ rowType(),
+ null,
+ new
Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator)
-> reader))
+ .isSameAs(failure);
+
+ verify(reader).close();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ assertThat(inputStream.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testCloseContinuesWhenReaderCloseThrows() throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = createReader();
+ RuntimeException failure = new RuntimeException("reader close failed");
+ doThrow(failure).when(reader).close();
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader);
+
+ assertThatThrownBy(recordsReader::close).isSameAs(failure);
+
+ verify(reader).close();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ assertThat(inputStream.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testCloseAddsSuppressedExceptionsFromLaterResources() throws
IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ RuntimeException allocatorFailure = new RuntimeException("allocator
close failed");
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator(allocatorFailure);
+ MosaicReader reader = createReader();
+ RuntimeException readerFailure = new RuntimeException("reader close
failed");
+ doThrow(readerFailure).when(reader).close();
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader);
+
+ assertThatThrownBy(recordsReader::close)
+ .isSameAs(readerFailure)
+ .satisfies(t ->
assertThat(t.getSuppressed()).containsExactly(allocatorFailure));
+
+ verify(reader).close();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ assertThat(inputStream.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testAllProjectedColumnsMissingSkipsRowGroupRead() throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = createReader();
+ when(reader.numRowGroups()).thenReturn(1);
+ when(reader.rowGroupNumRows(0)).thenReturn(3);
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader);
+
+ assertThat(readerBatchSize(recordsReader)).isEqualTo(3);
+ verify(reader, never()).readRowGroup(anyInt(), any());
+
+ recordsReader.close();
+ }
+
+ @Test
+ void testAllMissingBatchesTrackPositionsIndependently() throws IOException
{
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = createReader();
+ when(reader.numRowGroups()).thenReturn(2);
+ when(reader.rowGroupNumRows(0)).thenReturn(2);
+ when(reader.rowGroupNumRows(1)).thenReturn(2);
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader);
+ FileRecordIterator<InternalRow> first = recordsReader.readBatch();
+ FileRecordIterator<InternalRow> second = recordsReader.readBatch();
+
+ second.next();
+ assertThat(second.returnedPosition()).isEqualTo(2);
+ first.next();
+ assertThat(first.returnedPosition()).isZero();
+ second.next();
+ assertThat(second.returnedPosition()).isEqualTo(3);
+ first.next();
+ assertThat(first.returnedPosition()).isEqualTo(1);
+
+ first.releaseBatch();
+ second.releaseBatch();
+ verify(reader, never()).readRowGroup(anyInt(), any());
+ recordsReader.close();
+ }
+
+ @Test
+ void testMosaicBatchExposesVectorizedRecordIterator() throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = mock(MosaicReader.class);
+ IntVector vector = new IntVector("f0", allocator);
+ vector.allocateNew(3);
+ vector.setSafe(0, 10);
+ vector.setNull(1);
+ vector.setSafe(2, 30);
+ vector.setValueCount(3);
+ VectorSchemaRoot root = new
VectorSchemaRoot(Collections.singletonList(vector));
+ root.setRowCount(3);
+ when(reader.getSchema()).thenReturn(root.getSchema());
+ when(reader.numRowGroups()).thenReturn(1);
+ when(reader.rowGroupNumRows(0)).thenReturn(3);
+ when(reader.readRowGroup(0, allocator)).thenReturn(root);
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader);
+ FileRecordIterator<InternalRow> records = recordsReader.readBatch();
+
+ assertThat(records).isInstanceOf(VectorizedRecordIterator.class);
+ assertThat(((VectorizedRecordIterator)
records).batch().getNumRows()).isEqualTo(3);
+ assertThat(records.next().getInt(0)).isEqualTo(10);
+ assertThat(records.returnedPosition()).isZero();
+ assertThat(records.next().isNullAt(0)).isTrue();
+ assertThat(records.returnedPosition()).isEqualTo(1);
+ records.releaseBatch();
+ assertThat(allocator.getAllocatedMemory()).isZero();
+
+ recordsReader.close();
+ }
+
+ @Test
+ void testReadingNextBatchDoesNotInvalidateOutstandingBatch() throws
IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = mock(MosaicReader.class);
+ VectorSchemaRoot firstRoot = intRoot(allocator, 10);
+ VectorSchemaRoot secondRoot = intRoot(allocator, 20, 30);
+ when(reader.getSchema()).thenReturn(firstRoot.getSchema());
+ when(reader.numRowGroups()).thenReturn(2);
+ when(reader.rowGroupNumRows(0)).thenReturn(1);
+ when(reader.rowGroupNumRows(1)).thenReturn(2);
+ when(reader.readRowGroup(0, allocator)).thenReturn(firstRoot);
+ when(reader.readRowGroup(1, allocator)).thenReturn(secondRoot);
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader);
+ FileRecordIterator<InternalRow> first = recordsReader.readBatch();
+ FileRecordIterator<InternalRow> second = recordsReader.readBatch();
+
+ assertThat(first).isInstanceOf(VectorizedRecordIterator.class);
+ assertThat(second).isInstanceOf(VectorizedRecordIterator.class);
+ assertThat(((VectorizedRecordIterator) first).batch())
+ .isNotSameAs(((VectorizedRecordIterator) second).batch());
+ assertThat(((VectorizedRecordIterator)
first).batch().getNumRows()).isEqualTo(1);
+ assertThat(((VectorizedRecordIterator)
second).batch().getNumRows()).isEqualTo(2);
+ assertThat(first.next().getInt(0)).isEqualTo(10);
+ assertThat(second.next().getInt(0)).isEqualTo(20);
+ assertThat(second.next().getInt(0)).isEqualTo(30);
+ assertThat(second.next()).isNull();
+ assertThat(firstRoot.getVector(0).getObject(0)).isEqualTo(10);
+
+ second.releaseBatch();
+ assertThat(firstRoot.getVector(0).getObject(0)).isEqualTo(10);
+ assertThat(allocator.getAllocatedMemory()).isGreaterThan(0);
+ first.releaseBatch();
+ assertThat(allocator.getAllocatedMemory()).isZero();
+ recordsReader.close();
+ }
+
+ @Test
+ void testCloseReleasesOutstandingBatch() throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = mock(MosaicReader.class);
+ VectorSchemaRoot root = intRoot(allocator, 10);
+ when(reader.getSchema()).thenReturn(root.getSchema());
+ when(reader.numRowGroups()).thenReturn(1);
+ when(reader.rowGroupNumRows(0)).thenReturn(1);
+ when(reader.readRowGroup(0, allocator)).thenReturn(root);
+
+ MosaicRecordsReader recordsReader =
+ createRecordsReader(inputFileAdapter, allocator, reader);
+ FileRecordIterator<InternalRow> batch = recordsReader.readBatch();
+
+ assertThat(allocator.getAllocatedMemory()).isGreaterThan(0);
+ recordsReader.close();
+ assertThat(allocator.getAllocatedMemory()).isZero();
+ batch.releaseBatch();
+ }
+
+ @Test
+ void testCloseWaitsForConcurrentBatchRelease() throws Exception {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ BlockingRootAllocator allocator = new BlockingRootAllocator();
+ MosaicReader reader = mock(MosaicReader.class);
+ VectorSchemaRoot root = intRoot(allocator, 10);
+ when(reader.getSchema()).thenReturn(root.getSchema());
+ when(reader.numRowGroups()).thenReturn(1);
+ when(reader.rowGroupNumRows(0)).thenReturn(1);
+ when(reader.readRowGroup(0, allocator)).thenReturn(root);
+
+ MosaicRecordsReader recordsReader =
+ new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType(),
+ rowType(),
+ null,
+ new Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator) -> reader);
+ FileRecordIterator<InternalRow> batch = recordsReader.readBatch();
+ allocator.blockNextRelease();
+
+ AtomicReference<Throwable> releaseFailure = new AtomicReference<>();
+ Thread releaseThread =
+ new Thread(
+ () -> {
+ try {
+ batch.releaseBatch();
+ } catch (Throwable t) {
+ releaseFailure.set(t);
+ }
+ });
+ releaseThread.start();
+ allocator.awaitReleaseStarted();
+
+ CountDownLatch closeStarted = new CountDownLatch(1);
+ AtomicReference<Throwable> closeFailure = new AtomicReference<>();
+ Thread closeThread =
+ new Thread(
+ () -> {
+ closeStarted.countDown();
+ try {
+ recordsReader.close();
+ } catch (Throwable t) {
+ closeFailure.set(t);
+ }
+ });
+ closeThread.start();
+ assertThat(closeStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ awaitBlocked(closeThread);
+ assertThat(allocator.closeCount()).isZero();
+
+ allocator.allowRelease();
+ releaseThread.join(TimeUnit.SECONDS.toMillis(5));
+ closeThread.join(TimeUnit.SECONDS.toMillis(5));
+
+ assertThat(releaseThread.isAlive()).isFalse();
+ assertThat(closeThread.isAlive()).isFalse();
+ assertThat(releaseFailure.get()).isNull();
+ assertThat(closeFailure.get()).isNull();
+ assertThat(allocator.getAllocatedMemory()).isZero();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ }
+
+ @Test
+ void testRowTrackingFallsBackFromColumnBatch() throws IOException {
+ CloseCountingSeekableInputStream inputStream = new
CloseCountingSeekableInputStream();
+ MosaicInputFileAdapter inputFileAdapter =
createInputFileAdapter(inputStream);
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ MosaicReader reader = mock(MosaicReader.class);
+ RowType rowType =
+ RowType.builder()
+ .field("id", DataTypes.INT())
+ .field(SpecialFields.ROW_ID.name(), DataTypes.BIGINT())
+ .field(SpecialFields.SEQUENCE_NUMBER.name(),
DataTypes.BIGINT())
+ .build();
+ VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(rowType,
allocator);
+ IntVector idVector = (IntVector) root.getVector(0);
+ idVector.allocateNew(1);
+ idVector.setSafe(0, 7);
+ idVector.setValueCount(1);
+ BigIntVector rowIdVector = (BigIntVector) root.getVector(1);
+ rowIdVector.allocateNew(1);
+ rowIdVector.setNull(0);
+ rowIdVector.setValueCount(1);
+ BigIntVector sequenceVector = (BigIntVector) root.getVector(2);
+ sequenceVector.allocateNew(1);
+ sequenceVector.setNull(0);
+ sequenceVector.setValueCount(1);
+ root.setRowCount(1);
+ when(reader.getSchema()).thenReturn(root.getSchema());
+ when(reader.numRowGroups()).thenReturn(1);
+ when(reader.rowGroupNumRows(0)).thenReturn(1);
+ when(reader.readRowGroup(0, allocator)).thenReturn(root);
+
+ MosaicRecordsReader recordsReader =
+ new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType,
+ rowType,
+ null,
+ new Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator) -> reader);
+ Map<String, Integer> systemFields = new LinkedHashMap<>();
+ systemFields.put(SpecialFields.ROW_ID.name(), 1);
+ systemFields.put(SpecialFields.SEQUENCE_NUMBER.name(), 2);
+ DataFileRecordReader dataFileReader =
+ new DataFileRecordReader(
+ rowType,
+ recordsReader,
+ false,
+ false,
+ new int[] {0, 1, 2},
+ null,
+ null,
+ true,
+ 1000L,
+ 123L,
+ systemFields,
+ null,
+ new Path("file:/tmp/mosaic-reader-test"));
+ FileRecordIterator<InternalRow> tracked = dataFileReader.readBatch();
+
+ assertThat(tracked).isNotInstanceOf(VectorizedRecordIterator.class);
+ InternalRow row = tracked.next();
+ assertThat(row.getInt(0)).isEqualTo(7);
+ assertThat(row.getLong(1)).isEqualTo(1000L);
+ assertThat(row.getLong(2)).isEqualTo(123L);
+ assertThat(root.getVector(1).isNull(0)).isTrue();
+ assertThat(root.getVector(2).isNull(0)).isTrue();
+
+ tracked.releaseBatch();
+ dataFileReader.close();
+ }
+
+ private static MosaicInputFileAdapter createInputFileAdapter(
+ CloseCountingSeekableInputStream inputStream) throws IOException {
+ return new MosaicInputFileAdapter(
+ new CloseCountingFileIO(inputStream), new
Path("file:/tmp/mosaic-reader-test"));
+ }
+
+ private static MosaicRecordsReader createRecordsReader(
+ MosaicInputFileAdapter inputFileAdapter,
+ CloseCountingRootAllocator allocator,
+ MosaicReader reader) {
+ return new MosaicRecordsReader(
+ inputFileAdapter,
+ 0,
+ rowType(),
+ rowType(),
+ null,
+ new Path("file:/tmp/mosaic-reader-test"),
+ allocator,
+ (inputFile, fileSize, bufferAllocator) -> reader);
+ }
+
+ private static MosaicReader createReader() {
+ MosaicReader reader = mock(MosaicReader.class);
+ when(reader.getSchema()).thenReturn(new
Schema(Collections.emptyList()));
+ return reader;
+ }
+
+ private static int readerBatchSize(MosaicRecordsReader recordsReader)
throws IOException {
+ int count = 0;
+ while (true) {
+ FileRecordIterator<InternalRow> batch = recordsReader.readBatch();
+ if (batch == null) {
+ return count;
+ }
+ InternalRow row;
+ while ((row = batch.next()) != null) {
+ assertThat(row.isNullAt(0)).isTrue();
+ count++;
+ }
+ batch.releaseBatch();
+ }
+ }
+
+ private static RowType rowType() {
+ return DataTypes.ROW(DataTypes.INT());
+ }
+
+ 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 class CloseCountingFileIO extends LocalFileIO {
+
+ private final CloseCountingSeekableInputStream inputStream;
+
+ private CloseCountingFileIO(CloseCountingSeekableInputStream
inputStream) {
+ this.inputStream = inputStream;
+ }
+
+ @Override
+ public SeekableInputStream newInputStream(Path path) {
+ return inputStream;
+ }
+ }
+
+ private static class CloseCountingSeekableInputStream extends
SeekableInputStream {
+
+ private int closeCount;
+
+ @Override
+ public void seek(long desired) {}
+
+ @Override
+ public long getPos() {
+ return 0;
+ }
+
+ @Override
+ public int read(byte[] b, int off, int len) {
+ return -1;
+ }
+
+ @Override
+ public int read() {
+ return -1;
+ }
+
+ @Override
+ public void close() {
+ closeCount++;
+ }
+
+ int closeCount() {
+ return closeCount;
+ }
+ }
+
+ private static class CloseCountingRootAllocator extends RootAllocator {
+
+ private final RuntimeException closeFailure;
+ private int closeCount;
+
+ private CloseCountingRootAllocator() {
+ this(null);
+ }
+
+ private CloseCountingRootAllocator(RuntimeException closeFailure) {
+ this.closeFailure = closeFailure;
+ }
+
+ @Override
+ public void close() {
+ closeCount++;
+ if (closeFailure != null) {
+ throw closeFailure;
+ }
+ super.close();
+ }
+
+ int closeCount() {
+ return closeCount;
+ }
+ }
+
+ private static class BlockingRootAllocator extends
CloseCountingRootAllocator {
+
+ private final CountDownLatch releaseStarted = new CountDownLatch(1);
+ private final CountDownLatch allowRelease = new CountDownLatch(1);
+ private volatile boolean blockNextRelease;
+
+ private void blockNextRelease() {
+ blockNextRelease = true;
+ }
+
+ private void awaitReleaseStarted() throws InterruptedException {
+ assertThat(releaseStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ }
+
+ private void allowRelease() {
+ allowRelease.countDown();
+ }
+
+ @Override
+ public void releaseBytes(long size) {
+ if (blockNextRelease) {
+ blockNextRelease = false;
+ releaseStarted.countDown();
+ try {
+ assertThat(allowRelease.await(5,
TimeUnit.SECONDS)).isTrue();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ }
+ super.releaseBytes(size);
+ }
+ }
+
+ private static void awaitBlocked(Thread thread) {
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (thread.getState() != Thread.State.BLOCKED && System.nanoTime()
< deadline) {
+ Thread.yield();
+ }
+ assertThat(thread.getState()).isEqualTo(Thread.State.BLOCKED);
+ }
+}
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 14553219ce..3037987232 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
@@ -34,6 +34,7 @@ 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.io.VectorizedBundleRecords;
import org.apache.paimon.options.MemorySize;
import org.apache.paimon.options.Options;
import org.apache.paimon.predicate.Predicate;
@@ -41,6 +42,7 @@ 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.reader.VectorizedRecordIterator;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.RoaringBitmap32;
@@ -75,6 +77,55 @@ class MosaicReaderWriterTest {
assumeTrue(isNativeAvailable(), "Mosaic native library not available");
}
+ @Test
+ void testColumnBatchRoundTripThroughPublicWriter() throws IOException {
+ RowType rowType =
+ RowType.builder()
+ .field("id", DataTypes.INT())
+ .field("name", DataTypes.STRING())
+ .build();
+ Path sourcePath = newPath();
+ Path targetPath = newPath();
+ writeRows(
+ rowType,
+ sourcePath,
+ GenericRow.of(1, BinaryString.fromString("one")),
+ GenericRow.of(2, null));
+ LocalFileIO fileIO = new LocalFileIO();
+ MosaicFileFormat format = createFormat();
+ try (RecordReader<InternalRow> reader =
+ format.createReaderFactory(rowType, rowType, null)
+ .createReader(
+ new FormatReaderContext(
+ fileIO,
+ sourcePath,
+ fileIO.getFileSize(sourcePath),
+ null,
+ null));
+ FormatWriter writer =
+ format.createWriterFactory(rowType)
+ .create(fileIO.newOutputStream(targetPath,
false), "zstd")) {
+ RecordReader.RecordIterator<InternalRow> batch;
+ while ((batch = reader.readBatch()) != null) {
+ try {
+
assertThat(batch).isInstanceOf(VectorizedRecordIterator.class);
+ ((BundleFormatWriter) writer)
+ .writeBundle(
+ new VectorizedBundleRecords(
+ ((VectorizedRecordIterator)
batch).batch(), null));
+ } finally {
+ batch.releaseBatch();
+ }
+ }
+ }
+ List<InternalRow> rows = readAll(rowType, rowType, targetPath, null);
+ assertThat(rows).hasSize(2);
+ assertThat(rows.get(0).getInt(0)).isEqualTo(1);
+ assertThat(rows.get(0).getString(1).toString()).isEqualTo("one");
+ assertThat(rows.get(1).getInt(0)).isEqualTo(2);
+ assertThat(rows.get(1).isNullAt(1)).isTrue();
+ }
+
@Test
void testWriteAndRead() throws IOException {
RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING());
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 cc1f78cfdf..817ac3c328 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
@@ -36,6 +36,8 @@ import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import java.io.ByteArrayOutputStream;
import java.util.Arrays;
@@ -46,9 +48,11 @@ 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.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.same;
+import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
/** Test for {@link MosaicRecordsWriter}. */
@@ -373,4 +377,68 @@ class MosaicRecordsWriterTest {
return closeCount;
}
}
+
+ @Test
+ void testRejectsNonPositiveWriteBatchSize() {
+ RowType rowType = RowType.builder().field("a",
DataTypes.INT()).build();
+
+ for (int writeBatchSize : new int[] {0, -1}) {
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ FileFormatFactory.FormatContext formatContext =
+ new FileFormatFactory.FormatContext(new Options(), 1024,
writeBatchSize);
+
+ assertThatThrownBy(
+ () -> {
+ try (MosaicRecordsWriter ignored =
+ createWriter(
+ rowType,
+ allocator,
+ mock(MosaicWriter.class),
+ writeBatchSize,
+ MemorySize.VALUE_128_MB)) {
+ throw new AssertionError(
+ "Accepted write batch size " +
writeBatchSize);
+ }
+ })
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("write.batch-size")
+ .hasMessageContaining("greater than 0");
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testNativeWriteFailureIsNotRetriedDuringClose(boolean directWrite)
throws Exception {
+ RowType type = RowType.builder().field("a", DataTypes.INT()).build();
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ RuntimeException failure = new RuntimeException("native write failed");
+ doThrow(failure).when(nativeWriter).write(any(VectorSchemaRoot.class));
+ CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
+ try (MosaicRecordsWriter writer = createWriter(type, allocator,
nativeWriter)) {
+ try (BufferAllocator sourceAllocator =
+ allocator.newChildAllocator("mosaic-failure-test",
0, Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(type,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("a"), 1);
+ root.setRowCount(1);
+ ArrowBundleRecords bundle = new ArrowBundleRecords(root, type,
true);
+ if (directWrite) {
+ assertThatThrownBy(() ->
writer.writeBundle(bundle)).isSameAs(failure);
+ verify(nativeWriter).write(same(root));
+ } else {
+ writer.addElement(GenericRow.of(1));
+ assertThatThrownBy(writer::flush).isSameAs(failure);
+ }
+ assertThatThrownBy(() -> writer.addElement(GenericRow.of(2)))
+ .isInstanceOf(IllegalStateException.class);
+ assertThatThrownBy(() -> writer.writeBundle(bundle))
+ .isInstanceOf(IllegalStateException.class);
+ }
+ }
+ verify(nativeWriter, times(1)).write(any(VectorSchemaRoot.class));
+ verify(nativeWriter).close();
+ assertThat(allocator.closeCount()).isEqualTo(1);
+ assertThat(allocator.getAllocatedMemory()).isZero();
+ }
}
diff --git
a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRowTrackingTest.java
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRowTrackingTest.java
new file mode 100644
index 0000000000..048e58340e
--- /dev/null
+++
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRowTrackingTest.java
@@ -0,0 +1,208 @@
+/*
+ * 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.arrow.ArrowUtils;
+import org.apache.paimon.arrow.converter.ArrowBatchConverter;
+import org.apache.paimon.arrow.converter.ArrowPerRowBatchConverter;
+import org.apache.paimon.arrow.converter.ArrowVectorizedBatchConverter;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.columnar.ColumnarRowIterator;
+import org.apache.paimon.deletionvectors.ApplyDeletionFileRecordIterator;
+import org.apache.paimon.fs.Path;
+import org.apache.paimon.io.DataFileRecordReader;
+import org.apache.paimon.mosaic.MosaicReader;
+import org.apache.paimon.reader.FileRecordIterator;
+import org.apache.paimon.reader.VectorizedRecordIterator;
+import org.apache.paimon.table.SpecialFields;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.Collections;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+class MosaicRowTrackingTest {
+ @Test
+ void synthesizedRowIdsRequireRowFallback() throws Exception {
+ checkRowIds(false, false, false);
+ }
+
+ @Test
+ void synthesizedRowIdsRemainDistinctAfterDeletionSelection() throws
Exception {
+ checkRowIds(true, false, false);
+ }
+
+ @Test
+ void synthesizedRowIdsRemainCorrectThroughRowFallback() throws Exception {
+ checkRowIds(true, true, false);
+ }
+
+ @Test
+ void directDecorationDoesNotExposePositionDependentVectors() throws
Exception {
+ checkRowIds(false, false, true);
+ checkRowIds(true, false, true);
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {0, 1})
+ void synthesizedRowIdsRemainCorrectAfterReorderedProjection(int
rowGroupIndex)
+ throws Exception {
+ checkRowIds(true, false, false, true, rowGroupIndex);
+ }
+
+ private void checkRowIds(boolean deletion, boolean forceRowFallback,
boolean directDecoration)
+ throws Exception {
+ checkRowIds(deletion, forceRowFallback, directDecoration, false, 0);
+ }
+
+ private void checkRowIds(
+ boolean deletion,
+ boolean forceRowFallback,
+ boolean directDecoration,
+ boolean reorderedProjection,
+ int rowGroupIndex)
+ throws Exception {
+ RowType type =
+ RowType.builder()
+ .field("id", DataTypes.INT())
+ .field(SpecialFields.ROW_ID.name(), DataTypes.BIGINT())
+ .build();
+ RootAllocator allocator = new RootAllocator();
+ MosaicReader nativeReader = mock(MosaicReader.class);
+ when(nativeReader.numRowGroups()).thenReturn(rowGroupIndex + 1);
+ for (int group = 0; group <= rowGroupIndex; group++) {
+ VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(type,
allocator);
+ root.allocateNew();
+ for (int row = 0; row < 3; row++) {
+ ((IntVector) root.getVector(0)).setSafe(row, 7 + group * 3 +
row);
+ ((BigIntVector) root.getVector(1)).setNull(row);
+ }
+ root.setRowCount(3);
+ when(nativeReader.getSchema()).thenReturn(root.getSchema());
+ when(nativeReader.rowGroupNumRows(group)).thenReturn(3);
+ when(nativeReader.readRowGroup(group, allocator)).thenReturn(root);
+ }
+ Path path = new Path("file:/tmp/pr9-review");
+ MosaicRecordsReader reader =
+ new MosaicRecordsReader(
+ mock(MosaicInputFileAdapter.class),
+ 0,
+ type,
+ type,
+ null,
+ path,
+ allocator,
+ (file, size, alloc) -> nativeReader);
+ int[] indexMapping = reorderedProjection ? new int[] {1, 0} : new
int[] {0, 1};
+ int rowIdIndex = reorderedProjection ? 0 : 1;
+ RowType outputType = reorderedProjection ? type.project(indexMapping)
: type;
+ try (DataFileRecordReader dataReader =
+ new DataFileRecordReader(
+ outputType,
+ reader,
+ false,
+ false,
+ indexMapping,
+ null,
+ null,
+ true,
+ 1000L,
+ 123L,
+
Collections.singletonMap(SpecialFields.ROW_ID.name(), rowIdIndex),
+ null,
+ path);
+ RootAllocator outputAllocator = new RootAllocator();
+ VectorSchemaRoot output =
+ ArrowUtils.createVectorSchemaRoot(outputType,
outputAllocator)) {
+ for (int group = 0; group < rowGroupIndex; group++) {
+ FileRecordIterator<InternalRow> previousBatch =
dataReader.readBatch();
+ try {
+ for (int row = 0; row < 3; row++) {
+ assertThat(previousBatch.next().getLong(rowIdIndex))
+ .isEqualTo(1000L + group * 3 + row);
+ }
+ assertThat(previousBatch.next()).isNull();
+ } finally {
+ previousBatch.releaseBatch();
+ }
+ }
+ long startPosition = rowGroupIndex * 3L;
+ FileRecordIterator<InternalRow> batch =
+ directDecoration
+ ? ((ColumnarRowIterator) reader.readBatch())
+ .assignRowTracking(
+ 1000L,
+ 123L,
+ Collections.singletonMap(
+
SpecialFields.ROW_ID.name(), 1))
+ : dataReader.readBatch();
+ assertThat(batch).isNotInstanceOf(VectorizedRecordIterator.class);
+ ArrowBatchConverter converter;
+ if (!forceRowFallback && batch instanceof
VectorizedRecordIterator) {
+ ArrowVectorizedBatchConverter vectorConverter =
+ new ArrowVectorizedBatchConverter(
+ output,
ArrowUtils.createArrowFieldWriters(output, outputType));
+ if (deletion) {
+ vectorConverter.reset(
+ new ApplyDeletionFileRecordIterator(
+ batch, pos -> pos == startPosition + 1));
+ } else {
+ vectorConverter.reset((VectorizedRecordIterator) batch);
+ }
+ converter = vectorConverter;
+ } else {
+ ArrowPerRowBatchConverter rowConverter =
+ new ArrowPerRowBatchConverter(
+ output,
ArrowUtils.createArrowFieldWriters(output, outputType));
+ rowConverter.reset(
+ deletion
+ ? new ApplyDeletionFileRecordIterator(
+ batch, pos -> pos == startPosition + 1)
+ : batch);
+ converter = rowConverter;
+ }
+ converter.next(1024);
+ long[] observed = new long[output.getRowCount()];
+ for (int i = 0; i < observed.length; i++) {
+ observed[i] = ((BigIntVector)
output.getVector(rowIdIndex)).get(i);
+ }
+ batch.releaseBatch();
+ assertThat(observed)
+ .containsExactly(
+ deletion
+ ? new long[] {1000L + startPosition, 1002L
+ startPosition}
+ : new long[] {
+ 1000L + startPosition,
+ 1001L + startPosition,
+ 1002L + startPosition
+ });
+ }
+ }
+}
diff --git
a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/extension/MosaicWriterExtensionTest.java
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/extension/MosaicWriterExtensionTest.java
new file mode 100644
index 0000000000..4492c73768
--- /dev/null
+++
b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/extension/MosaicWriterExtensionTest.java
@@ -0,0 +1,150 @@
+/*
+ * 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.extension;
+
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.columnar.ColumnVector;
+import org.apache.paimon.data.columnar.heap.HeapIntVector;
+import org.apache.paimon.format.FileFormatFactory;
+import org.apache.paimon.format.FormatReaderContext;
+import org.apache.paimon.format.FormatWriter;
+import org.apache.paimon.format.mosaic.MosaicFileFormat;
+import org.apache.paimon.format.mosaic.MosaicRecordsWriter;
+import org.apache.paimon.format.mosaic.MosaicWriterFactory;
+import org.apache.paimon.fs.Path;
+import org.apache.paimon.fs.PositionOutputStream;
+import org.apache.paimon.fs.local.LocalFileIO;
+import org.apache.paimon.options.Options;
+import org.apache.paimon.reader.RecordReader;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.mock;
+
+/** Verifies external writer extension through the normal public constructor.
*/
+class MosaicWriterExtensionTest {
+
+ private static final RowType TYPE = RowType.of(DataTypes.INT());
+ private static final FileFormatFactory.FormatContext CONTEXT =
+ new FileFormatFactory.FormatContext(new Options(), 1024, 1024);
+
+ @TempDir java.nio.file.Path tempDir;
+
+ @Test
+ void externalWriterReusesConversionBufferAndCleanup() throws Exception {
+ LocalFileIO fileIO = new LocalFileIO();
+ Path path = new Path(tempDir.resolve("columns.mosaic").toUri());
+ BufferAllocator allocator;
+ try (ExternalWriter writer = new
ExternalWriter(fileIO.newOutputStream(path, false))) {
+ allocator = writer.allocator();
+ HeapIntVector input = new HeapIntVector(2);
+ input.setInt(0, 11);
+ input.setInt(1, 22);
+ writer.writeColumns(new ColumnVector[] {input}, 2);
+ }
+ assertThat(allocator.getAllocatedMemory()).isZero();
+ assertThat(readValues(fileIO, path)).containsExactly(11, 22);
+ }
+
+ @Test
+ void factoryHookRetainsCompressionValidation() throws Exception {
+ int[] created = {0};
+ MosaicWriterFactory factory =
+ new MosaicWriterFactory(TYPE, CONTEXT) {
+ @Override
+ protected MosaicRecordsWriter createWriter(
+ OutputStream output,
+ RowType type,
+ FileFormatFactory.FormatContext context,
+ List<String> statsColumns,
+ Integer buckets) {
+ created[0]++;
+ assertThat(type).isEqualTo(TYPE);
+ assertThat(context).isSameAs(CONTEXT);
+ assertThat(statsColumns).isEmpty();
+ assertThat(buckets).isNull();
+ return new ExternalWriter(output);
+ }
+ };
+ assertThatThrownBy(() ->
factory.create(mock(PositionOutputStream.class), "snappy"))
+ .isInstanceOf(UnsupportedOperationException.class);
+ assertThat(created[0]).isZero();
+ LocalFileIO fileIO = new LocalFileIO();
+ Path path = new Path(tempDir.resolve("factory.mosaic").toUri());
+ try (FormatWriter writer = factory.create(fileIO.newOutputStream(path,
false), "zstd")) {
+ assertThat(writer).isInstanceOf(ExternalWriter.class);
+ assertThat(created[0]).isEqualTo(1);
+ HeapIntVector input = new HeapIntVector(1);
+ input.setInt(0, 33);
+ ((ExternalWriter) writer).writeColumns(new ColumnVector[] {input},
1);
+ }
+ assertThat(readValues(fileIO, path)).containsExactly(33);
+ }
+
+ private static List<Integer> readValues(LocalFileIO fileIO, Path path)
throws Exception {
+ List<Integer> values = new ArrayList<>();
+ MosaicFileFormat format = new MosaicFileFormat(CONTEXT);
+ try (RecordReader<InternalRow> reader =
+ format.createReaderFactory(TYPE, TYPE, Collections.emptyList())
+ .createReader(
+ new FormatReaderContext(
+ fileIO, path,
fileIO.getFileSize(path), null, null))) {
+ RecordReader.RecordIterator<InternalRow> batch;
+ while ((batch = reader.readBatch()) != null) {
+ try {
+ InternalRow row;
+ while ((row = batch.next()) != null) {
+ values.add(row.getInt(0));
+ }
+ } finally {
+ batch.releaseBatch();
+ }
+ }
+ }
+ return values;
+ }
+
+ private static class ExternalWriter extends MosaicRecordsWriter {
+
+ private ExternalWriter(OutputStream output) {
+ super(output, TYPE, CONTEXT, Collections.emptyList(), null);
+ }
+
+ private BufferAllocator allocator() {
+ return arrowWriter().getAllocator();
+ }
+
+ private void writeColumns(ColumnVector[] columns, int rows) {
+ checkNotFailed();
+ arrowWriter().write(columns, null, 0, rows);
+ flush();
+ }
+ }
+}
