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 87eb04451d [core][format] Support Arrow bundle writes (#9100)
87eb04451d is described below
commit 87eb04451d51de4a86b643b3b9a6c82a2c97d0f2
Author: jianguotian <[email protected]>
AuthorDate: Thu Aug 13 14:13:21 2026 +0800
[core][format] Support Arrow bundle writes (#9100)
---
.../apache/paimon/arrow/ArrowBundleRecords.java | 33 ++
.../java/org/apache/paimon/arrow/ArrowUtils.java | 34 ++
.../paimon/arrow/vector/ArrowFormatWriter.java | 72 +++
.../paimon/arrow/writer/ArrowBundleWriter.java | 26 +-
.../org/apache/paimon/arrow/ArrowUtilsTest.java | 44 ++
.../paimon/arrow/vector/ArrowFormatWriterTest.java | 29 ++
.../paimon/arrow/writer/ArrowBundleWriterTest.java | 293 ++++++++++-
.../apache/paimon/format/BundleFormatWriter.java | 11 +-
.../shredding/InferShreddingWritePlanWriter.java | 62 +--
.../java/org/apache/paimon/io/BundleRecords.java | 6 +-
.../InferShreddingWritePlanWriterTest.java | 213 ++++++++
.../apache/paimon/io/RollingFileWriterImpl.java | 5 +-
.../io/RowDataFileSequenceNumberTracker.java | 15 +
.../org/apache/paimon/io/RowDataFileWriter.java | 9 +
.../org/apache/paimon/io/SingleFileWriter.java | 6 +-
.../paimon/io/StatsCollectingSingleFileWriter.java | 4 +
.../apache/paimon/io/RollingFileWriterTest.java | 2 +-
.../apache/paimon/io/RowDataFileWriterTest.java | 572 +++++++++++++++++++++
.../paimon/format/lance/LanceRecordsWriter.java | 17 +-
.../paimon/format/lance/jni/LanceWriter.java | 20 +-
.../format/lance/LanceRecordsWriterTest.java | 297 +++++++++++
.../paimon/format/mosaic/MosaicRecordsWriter.java | 19 +-
.../format/mosaic/MosaicReaderWriterTest.java | 58 +++
.../format/mosaic/MosaicRecordsWriterTest.java | 121 ++++-
.../paimon/format/vortex/VortexRecordsWriter.java | 89 +++-
.../format/vortex/VortexReaderWriterTest.java | 127 ++++-
26 files changed, 2066 insertions(+), 118 deletions(-)
diff --git
a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java
index 25f6603ec2..b179beafb0 100644
--- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java
+++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java
@@ -21,11 +21,18 @@ package org.apache.paimon.arrow;
import org.apache.paimon.arrow.reader.ArrowBatchReader;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.io.BundleRecords;
+import org.apache.paimon.types.DataField;
import org.apache.paimon.types.RowType;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Field;
+import java.util.HashSet;
import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
+
+import static org.apache.paimon.utils.StringUtils.toLowerCaseIfNeed;
/** Batch records for vector schema root. */
public class ArrowBundleRecords implements BundleRecords {
@@ -45,6 +52,32 @@ public class ArrowBundleRecords implements BundleRecords {
return vectorSchemaRoot;
}
+ public RowType getRowType() {
+ return rowType;
+ }
+
+ /**
+ * Returns whether row iteration reads every Arrow vector at the same
position without name
+ * remapping or synthesized null columns.
+ */
+ public boolean hasIdentityMapping() {
+ List<Field> arrowFields = vectorSchemaRoot.getSchema().getFields();
+ List<DataField> dataFields = rowType.getFields();
+ if (arrowFields.size() != dataFields.size()) {
+ return false;
+ }
+
+ Set<String> mappedNames = new HashSet<>();
+ for (int i = 0; i < arrowFields.size(); i++) {
+ String arrowName = toLowerCaseIfNeed(arrowFields.get(i).getName(),
caseSensitive);
+ String dataName = toLowerCaseIfNeed(dataFields.get(i).name(),
caseSensitive);
+ if (!arrowName.equals(dataName) || !mappedNames.add(arrowName)) {
+ return false;
+ }
+ }
+ return true;
+ }
+
@Override
public long rowCount() {
return vectorSchemaRoot.getRowCount();
diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java
index 404e6c009e..5ea4ac8aaa 100644
--- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java
+++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java
@@ -277,12 +277,35 @@ public class ArrowUtils {
return ArrowCStruct.of(array, schema);
}
+ /** Releases Arrow C Data callbacks that have not already been consumed by
native code. */
+ public static void releaseCDataIfNeeded(ArrowArray array, ArrowSchema
schema) {
+ try {
+ if (array.snapshot().release != 0) {
+ array.release();
+ }
+ } finally {
+ if (schema.snapshot().release != 0) {
+ schema.release();
+ }
+ }
+ }
+
public static byte[] serializeToIpc(VectorSchemaRoot vsr) {
ByteArrayOutputStream out = new ByteArrayOutputStream();
serializeToIpc(vsr, out);
return out.toByteArray();
}
+ /**
+ * Returns whether the schema root contains at least one vector and all
top-level and nested
+ * vectors share the root allocator of the supplied allocator.
+ */
+ public static boolean hasSameRootAllocator(
+ VectorSchemaRoot vectorSchemaRoot, BufferAllocator allocator) {
+ List<FieldVector> vectors = vectorSchemaRoot.getFieldVectors();
+ return !vectors.isEmpty() && allVectorsShareRootWith(vectors,
allocator.getRoot());
+ }
+
public static void serializeToIpc(VectorSchemaRoot vsr, OutputStream out) {
try (ArrowStreamWriter writer = new ArrowStreamWriter(vsr, null, out))
{
writer.writeBatch();
@@ -316,4 +339,15 @@ public class ArrowUtils {
return instant.getEpochSecond() * 1_000_000_000 +
instant.getNano();
}
}
+
+ private static boolean allVectorsShareRootWith(
+ List<FieldVector> vectors, BufferAllocator expectedRoot) {
+ for (FieldVector vector : vectors) {
+ if (vector.getAllocator().getRoot() != expectedRoot
+ ||
!allVectorsShareRootWith(vector.getChildrenFromFields(), expectedRoot)) {
+ return false;
+ }
+ }
+ return true;
+ }
}
diff --git
a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java
index 6133bfb550..f120b141cd 100644
---
a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java
+++
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java
@@ -18,6 +18,7 @@
package org.apache.paimon.arrow.vector;
+import org.apache.paimon.arrow.ArrowBundleRecords;
import org.apache.paimon.arrow.ArrowFieldTypeConversion;
import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.writer.ArrowFieldWriter;
@@ -25,10 +26,14 @@ import
org.apache.paimon.arrow.writer.ArrowFieldWriterFactoryVisitor;
import org.apache.paimon.arrow.writer.ArrowFieldWriters;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.data.columnar.ColumnVector;
+import org.apache.paimon.types.ArrayType;
import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataType;
+import org.apache.paimon.types.MapType;
+import org.apache.paimon.types.MultisetType;
import org.apache.paimon.types.RowType;
import org.apache.paimon.types.VariantType;
+import org.apache.paimon.types.VectorType;
import org.apache.paimon.utils.Preconditions;
import org.apache.arrow.memory.BufferAllocator;
@@ -51,6 +56,7 @@ public class ArrowFormatWriter implements AutoCloseable {
private final VectorSchemaRoot vectorSchemaRoot;
private final ArrowFieldWriter[] fieldWriters;
+ private final RowType rowType;
private final int batchSize;
private final BufferAllocator allocator;
@Nullable private final Long memoryUsedMaxInBytes;
@@ -171,6 +177,7 @@ public class ArrowFormatWriter implements AutoCloseable {
boolean closeAllocatorOnClose) {
this.allocator = allocator;
this.closeAllocatorOnClose = closeAllocatorOnClose;
+ this.rowType = rowType;
RowType outputRowType = replaceWithShreddingType(rowType,
shreddingSchemas);
vectorSchemaRoot =
@@ -303,6 +310,71 @@ public class ArrowFormatWriter implements AutoCloseable {
return allocator;
}
+ /** Returns whether direct Arrow consumption preserves this writer's row
schema. */
+ public boolean isArrowBundleSchemaCompatible(ArrowBundleRecords bundle) {
+ return !bundle.getVectorSchemaRoot().getFieldVectors().isEmpty()
+ && bundle.hasIdentityMapping()
+ && hasSameLogicalLayout(rowType, bundle.getRowType())
+ &&
vectorSchemaRoot.getSchema().equals(bundle.getVectorSchemaRoot().getSchema());
+ }
+
+ private static boolean hasSameLogicalLayout(DataType left, DataType right)
{
+ if (left == right) {
+ return true;
+ }
+ if (left == null
+ || right == null
+ || left.getClass() != right.getClass()
+ || left.isNullable() != right.isNullable()) {
+ return false;
+ }
+
+ if (left instanceof RowType) {
+ List<DataField> leftFields = ((RowType) left).getFields();
+ List<DataField> rightFields = ((RowType) right).getFields();
+ if (leftFields.size() != rightFields.size()) {
+ return false;
+ }
+ for (int i = 0; i < leftFields.size(); i++) {
+ DataField leftField = leftFields.get(i);
+ DataField rightField = rightFields.get(i);
+ if (!leftField.name().equals(rightField.name())
+ || !hasSameLogicalLayout(leftField.type(),
rightField.type())) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ if (left instanceof ArrayType) {
+ return hasSameLogicalLayout(
+ ((ArrayType) left).getElementType(), ((ArrayType)
right).getElementType());
+ }
+
+ if (left instanceof MapType) {
+ MapType leftMap = (MapType) left;
+ MapType rightMap = (MapType) right;
+ return hasSameLogicalLayout(leftMap.getKeyType(),
rightMap.getKeyType())
+ && hasSameLogicalLayout(leftMap.getValueType(),
rightMap.getValueType());
+ }
+
+ if (left instanceof MultisetType) {
+ return hasSameLogicalLayout(
+ ((MultisetType) left).getElementType(),
+ ((MultisetType) right).getElementType());
+ }
+
+ if (left instanceof VectorType) {
+ VectorType leftVector = (VectorType) left;
+ VectorType rightVector = (VectorType) right;
+ return leftVector.getLength() == rightVector.getLength()
+ && hasSameLogicalLayout(
+ leftVector.getElementType(),
rightVector.getElementType());
+ }
+
+ return left.equals(right);
+ }
+
private static RowType replaceWithShreddingType(
RowType rowType, @Nullable RowType shreddingSchemas) {
if (shreddingSchemas == null) {
diff --git
a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java
index 39a006b6d7..f8c7a22bb7 100644
---
a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java
+++
b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java
@@ -76,14 +76,22 @@ public class ArrowBundleWriter implements
BundleFormatWriter {
@Override
public void writeBundle(BundleRecords bundleRecords) throws IOException {
if (bundleRecords instanceof ArrowBundleRecords) {
- add(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot());
+ ArrowBundleRecords arrowBundle = (ArrowBundleRecords)
bundleRecords;
+ VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot();
+ if
(arrowFormatWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle)
+ && ArrowUtils.hasSameRootAllocator(root,
root.getVector(0).getAllocator())) {
+ flush();
+ add(root);
+ return;
+ }
} else if (bundleRecords instanceof VectorizedBundleRecords) {
VectorizedBundleRecords records = (VectorizedBundleRecords)
bundleRecords;
add(records.batch(), records.selected());
- } else {
- for (InternalRow row : bundleRecords) {
- addElement(row);
- }
+ return;
+ }
+
+ for (InternalRow row : bundleRecords) {
+ addElement(row);
}
}
@@ -96,9 +104,11 @@ public class ArrowBundleWriter implements
BundleFormatWriter {
ArrowUtils.serializeToCStruct(vsr, array, schema,
bufferAllocator);
long t2 = System.currentTimeMillis();
serializeCost += (t2 - t1);
- this.nativeWriter.writeIpcBytes(struct.arrayAddress(),
struct.schemaAddress());
- array.release();
- schema.release();
+ try {
+ this.nativeWriter.writeIpcBytes(struct.arrayAddress(),
struct.schemaAddress());
+ } finally {
+ ArrowUtils.releaseCDataIfNeeded(array, schema);
+ }
jniCost += (System.currentTimeMillis() - t2);
} catch (RuntimeException e) {
LOG.error("Exception happens while add vsr", e);
diff --git
a/paimon-arrow/src/test/java/org/apache/paimon/arrow/ArrowUtilsTest.java
b/paimon-arrow/src/test/java/org/apache/paimon/arrow/ArrowUtilsTest.java
index ee613a05d7..75d70c849b 100644
--- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/ArrowUtilsTest.java
+++ b/paimon-arrow/src/test/java/org/apache/paimon/arrow/ArrowUtilsTest.java
@@ -23,13 +23,19 @@ import org.apache.paimon.types.DataType;
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.FieldVector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.Test;
+import java.util.Collections;
import java.util.List;
import java.util.Random;
@@ -118,4 +124,42 @@ public class ArrowUtilsTest {
.isEqualTo(new ArrowType.FixedSizeList(4));
Assertions.assertThat(field.getChildren()).hasSize(1);
}
+
+ @Test
+ public void testSameRootAllocatorIncludesNestedVectors() {
+ try (RootAllocator allocator = new RootAllocator();
+ BufferAllocator childAllocator =
+ allocator.newChildAllocator("same-root-child", 0,
Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ nestedRoot(allocator, new IntVector("value",
childAllocator))) {
+ Assertions.assertThat(ArrowUtils.hasSameRootAllocator(root,
allocator)).isTrue();
+ }
+
+ try (RootAllocator allocator = new RootAllocator();
+ RootAllocator differentRoot = new RootAllocator();
+ VectorSchemaRoot root =
+ nestedRoot(allocator, new IntVector("value",
differentRoot))) {
+ Assertions.assertThat(ArrowUtils.hasSameRootAllocator(root,
allocator)).isFalse();
+ }
+ }
+
+ private static VectorSchemaRoot nestedRoot(BufferAllocator allocator,
FieldVector childVector) {
+ TestingStructVector structVector = new TestingStructVector("nested",
allocator);
+ structVector.putTestingChild("value", childVector);
+ return new VectorSchemaRoot(
+ Collections.singletonList(structVector.getField()),
+ Collections.singletonList(structVector),
+ 0);
+ }
+
+ private static class TestingStructVector extends StructVector {
+
+ private TestingStructVector(String name, BufferAllocator allocator) {
+ super(name, allocator,
FieldType.nullable(ArrowType.Struct.INSTANCE), null);
+ }
+
+ private void putTestingChild(String name, FieldVector childVector) {
+ putChild(name, childVector);
+ }
+ }
}
diff --git
a/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java
b/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java
index 40aa87dae9..fd257450aa 100644
---
a/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java
+++
b/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java
@@ -623,6 +623,35 @@ public class ArrowFormatWriterTest {
}
}
+ @Test
+ public void testArrowBundleSchemaCompatibilityIgnoresFieldDescription() {
+ RowType writerType = RowType.builder().field("value",
DataTypes.INT()).build();
+ RowType bundleType =
+ RowType.builder().field("value", DataTypes.INT(), "different
description").build();
+
+ try (ArrowFormatWriter writer = new ArrowFormatWriter(writerType, 1,
true)) {
+ assertThat(
+ writer.isArrowBundleSchemaCompatible(
+ new ArrowBundleRecords(
+ writer.getVectorSchemaRoot(),
bundleType, true)))
+ .isTrue();
+ }
+ }
+
+ @Test
+ public void testArrowBundleSchemaCompatibilityRequiresLogicalType() {
+ RowType writerType = RowType.builder().field("value",
DataTypes.VARCHAR(10)).build();
+ RowType bundleType = RowType.builder().field("value",
DataTypes.CHAR(10)).build();
+
+ try (ArrowFormatWriter writer = new ArrowFormatWriter(writerType, 1,
true)) {
+ assertThat(
+ writer.isArrowBundleSchemaCompatible(
+ new ArrowBundleRecords(
+ writer.getVectorSchemaRoot(),
bundleType, true)))
+ .isFalse();
+ }
+ }
+
@ParameterizedTest
@ValueSource(booleans = {false, true})
public void testWriteWithExternalAllocator(boolean allocationFailed) {
diff --git
a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java
b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java
index a43ec4e0b6..e7a82d1448 100644
---
a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java
+++
b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java
@@ -18,8 +18,11 @@
package org.apache.paimon.arrow.writer;
+import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.vector.ArrowFormatCWriter;
import org.apache.paimon.arrow.vector.ArrowFormatWriter;
+import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.columnar.ColumnVector;
import org.apache.paimon.data.columnar.VectorizedColumnBatch;
import org.apache.paimon.data.columnar.heap.HeapArrayVector;
@@ -28,10 +31,14 @@ import org.apache.paimon.data.columnar.heap.HeapLongVector;
import org.apache.paimon.data.columnar.heap.HeapMapVector;
import org.apache.paimon.data.columnar.heap.HeapRowVector;
import org.apache.paimon.fs.PositionOutputStream;
+import org.apache.paimon.io.VectorizedBundleRecords;
+import org.apache.paimon.types.DataField;
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.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.junit.jupiter.api.Test;
@@ -39,13 +46,289 @@ import org.junit.jupiter.api.Test;
import java.io.IOException;
import java.util.ArrayList;
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;
/** Tests for {@link ArrowBundleWriter}. */
public class ArrowBundleWriterTest {
+ @Test
+ public void testArrowBundleFlushesBufferedRowsBeforeDirectWrite() throws
Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024,
true);
+ List<String> events = new ArrayList<>();
+ NativeWriter nativeWriter =
+ new NativeWriter() {
+ @Override
+ public long nativeMemoryUsed() {
+ return 0;
+ }
+
+ @Override
+ public void writeIpcBytes(long arrayAddress, long
schemaAddress) {
+ events.add("rows");
+ cWriter.release();
+ }
+
+ @Override
+ public void close() {}
+ };
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter) {
+ @Override
+ public void add(VectorSchemaRoot vsr) {
+ events.add("bundle");
+ }
+ };
+
+ writer.addElement(GenericRow.of(1));
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
ArrowUtils.createVectorSchemaRoot(rowType, allocator)) {
+ setInt((IntVector) root.getVector("value"), 2);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, rowType, true));
+ }
+
+ assertThat(events).containsExactly("rows", "bundle");
+ writer.close();
+ }
+
+ @Test
+ public void testReorderedArrowBundleFallsBackToRows() throws Exception {
+ RowType writerType =
+ RowType.builder().field("a", DataTypes.INT()).field("b",
DataTypes.INT()).build();
+ RowType sourceType =
+ RowType.builder().field("b", DataTypes.INT()).field("a",
DataTypes.INT()).build();
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024,
true);
+ VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot();
+ CapturingNativeWriter nativeWriter = new
CapturingNativeWriter(writerRoot, cWriter);
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter) {
+ @Override
+ public void add(VectorSchemaRoot vsr) {
+ throw new AssertionError("Reordered Arrow bundle must
use row fallback.");
+ }
+ };
+
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
ArrowUtils.createVectorSchemaRoot(sourceType, allocator)) {
+ setInt((IntVector) root.getVector("b"), 20);
+ setInt((IntVector) root.getVector("a"), 10);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, writerType, true));
+ }
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(1);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20);
+ }
+
+ @Test
+ public void testMissingColumnArrowBundleFallsBackToRows() throws Exception
{
+ RowType writerType =
+ RowType.builder().field("a", DataTypes.INT()).field("b",
DataTypes.INT()).build();
+ RowType sourceType = RowType.builder().field("a",
DataTypes.INT()).build();
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024,
true);
+ VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot();
+ CapturingNativeWriter nativeWriter = new
CapturingNativeWriter(writerRoot, cWriter);
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter) {
+ @Override
+ public void add(VectorSchemaRoot vsr) {
+ throw new AssertionError(
+ "Arrow bundle with a missing column must use
row fallback.");
+ }
+ };
+
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
ArrowUtils.createVectorSchemaRoot(sourceType, allocator)) {
+ setInt((IntVector) root.getVector("a"), 10);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, writerType, true));
+ }
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(1);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10);
+ assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1))
+ .containsExactly((Object) null);
+ }
+
+ @Test
+ public void testDescriptionOnlyRowTypeDifferenceUsesDirectWrite() throws
Exception {
+ RowType writerType = RowType.builder().field("value",
DataTypes.INT()).build();
+ RowType bundleType =
+ new RowType(
+ Collections.singletonList(
+ new DataField(
+ 0, "value", DataTypes.INT(),
"different description")));
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024,
true);
+ VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot();
+ CapturingNativeWriter nativeWriter = new
CapturingNativeWriter(writerRoot, cWriter);
+ int[] directWrites = {0};
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter) {
+ @Override
+ public void add(VectorSchemaRoot vsr) {
+ directWrites[0]++;
+ }
+ };
+
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
ArrowUtils.createVectorSchemaRoot(writerType, allocator)) {
+ setInt((IntVector) root.getVector("value"), 10);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, bundleType, true));
+ }
+ writer.close();
+
+ assertThat(directWrites[0]).isEqualTo(1);
+ assertThat(nativeWriter.snapshots).isEmpty();
+ }
+
+ @Test
+ public void testNonIdentityNameMappingFallsBackToRows() throws Exception {
+ RowType writerType =
+ RowType.builder().field("A", DataTypes.INT()).field("a",
DataTypes.INT()).build();
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024,
true);
+ VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot();
+ CapturingNativeWriter nativeWriter = new
CapturingNativeWriter(writerRoot, cWriter);
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter) {
+ @Override
+ public void add(VectorSchemaRoot vsr) {
+ throw new AssertionError(
+ "Non-identity Arrow name mapping must use row
fallback.");
+ }
+ };
+
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
ArrowUtils.createVectorSchemaRoot(writerType, allocator)) {
+ setInt((IntVector) root.getVector("A"), 10);
+ setInt((IntVector) root.getVector("a"), 20);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, writerType,
false));
+ }
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(1);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(20);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20);
+ }
+
+ @Test
+ public void testMixedAllocatorRootsFallBackToRows() throws Exception {
+ RowType rowType =
+ RowType.builder().field("a", DataTypes.INT()).field("b",
DataTypes.INT()).build();
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024,
true);
+ VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot();
+ CapturingNativeWriter nativeWriter = new
CapturingNativeWriter(writerRoot, cWriter);
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter) {
+ @Override
+ public void add(VectorSchemaRoot vsr) {
+ throw new AssertionError("Mixed-root Arrow bundle must
use row fallback.");
+ }
+ };
+
+ try (RootAllocator firstAllocator = new RootAllocator();
+ RootAllocator secondAllocator = new RootAllocator()) {
+ FieldVector firstVector =
+
writerRoot.getSchema().getFields().get(0).createVector(firstAllocator);
+ FieldVector secondVector =
+
writerRoot.getSchema().getFields().get(1).createVector(secondAllocator);
+ try (VectorSchemaRoot root =
+ new VectorSchemaRoot(
+ writerRoot.getSchema(), Arrays.asList(firstVector,
secondVector), 1)) {
+ setInt((IntVector) firstVector, 10);
+ setInt((IntVector) secondVector, 20);
+
+ writer.writeBundle(new ArrowBundleRecords(root, rowType,
true));
+ }
+ }
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(1);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20);
+ }
+
+ @Test
+ public void testDirectWriteFailureReleasesExportedReferences() throws
Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024,
true);
+ RuntimeException failure = new RuntimeException("Expected native write
failure.");
+ NativeWriter nativeWriter =
+ new NativeWriter() {
+ @Override
+ public long nativeMemoryUsed() {
+ return 0;
+ }
+
+ @Override
+ public void writeIpcBytes(long arrayAddress, long
schemaAddress) {
+ throw failure;
+ }
+
+ @Override
+ public void close() {}
+ };
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter);
+
+ try {
+ try (RootAllocator allocator = new RootAllocator()) {
+ try (VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType, allocator))
{
+ setInt((IntVector) root.getVector("value"), 10);
+ root.setRowCount(1);
+
+ assertThatThrownBy(
+ () ->
+ writer.writeBundle(
+ new
ArrowBundleRecords(root, rowType, true)))
+ .isSameAs(failure);
+ }
+
+ assertThat(allocator.getAllocatedMemory()).isZero();
+ }
+ } finally {
+ writer.close();
+ }
+ }
+
+ @Test
+ public void testVectorizedBundleFlushesBufferedRowsBeforeWrite() throws
Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024,
true);
+ CapturingNativeWriter nativeWriter =
+ new CapturingNativeWriter(cWriter.getVectorSchemaRoot(),
cWriter);
+ ArrowBundleWriter writer =
+ new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter,
nativeWriter);
+
+ writer.addElement(GenericRow.of(1));
+ HeapIntVector vector = new HeapIntVector(2);
+ vector.setInt(0, 2);
+ vector.setInt(1, 3);
+ VectorizedColumnBatch batch = new VectorizedColumnBatch(new
ColumnVector[] {vector});
+ batch.setNumRows(2);
+ writer.writeBundle(new VectorizedBundleRecords(batch, null));
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(2);
+
assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(1);
+
assertThat(nativeWriter.snapshots.get(1).objectColumns.get(0)).containsExactly(2,
3);
+ }
+
@Test
public void testAddBatchWithoutDeletionVector() throws IOException {
RowType rowType = RowType.of(DataTypes.INT(), DataTypes.BIGINT());
@@ -578,12 +861,12 @@ public class ArrowBundleWriterTest {
if (fv instanceof IntVector) {
intValues = new int[rowCount];
for (int i = 0; i < rowCount; i++) {
- intValues[i] = ((IntVector) fv).get(i);
+ intValues[i] = fv.isNull(i) ? 0 : ((IntVector)
fv).get(i);
}
} else if (fv instanceof BigIntVector) {
longValues = new long[rowCount];
for (int i = 0; i < rowCount; i++) {
- longValues[i] = ((BigIntVector) fv).get(i);
+ longValues[i] = fv.isNull(i) ? 0L : ((BigIntVector)
fv).get(i);
}
}
List<Object> colObjects = new ArrayList<>();
@@ -618,6 +901,12 @@ public class ArrowBundleWriterTest {
}
}
+ private static void setInt(IntVector vector, int value) {
+ vector.allocateNew(1);
+ vector.setSafe(0, value);
+ vector.setValueCount(1);
+ }
+
private static class NoOpPositionOutputStream extends PositionOutputStream
{
private long pos = 0;
diff --git
a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java
b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java
index e0ba09b612..95be4e9339 100644
---
a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java
+++
b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java
@@ -22,11 +22,18 @@ import org.apache.paimon.io.BundleRecords;
import java.io.IOException;
-/** Format write with bundle interface. */
+/**
+ * Format writer with a row-equivalent bundle interface.
+ *
+ * <p>Implementations may consume a compatible bundle natively, convert or
copy it, or fall back to
+ * row-by-row writes. {@link #writeBundle} must preserve the same record
values and order as
+ * invoking {@link #addElement} for every record. It must not retain borrowed
buffers after the
+ * method returns unless it has copied them or acquired independent ownership.
+ */
public interface BundleFormatWriter extends FormatWriter {
/**
- * Write a bundle of records directly.
+ * Writes a bundle with semantics equivalent to invoking {@link
#addElement} for every record.
*
* @param bundle the records to be written
* @throws IOException if exception happens
diff --git
a/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java
b/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java
index 22f102d814..37b6674e01 100644
---
a/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java
+++
b/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java
@@ -26,12 +26,10 @@ import org.apache.paimon.fs.PositionOutputStream;
import org.apache.paimon.io.BundleRecords;
import org.apache.paimon.utils.InternalRowUtils;
-import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.Iterator;
import java.util.List;
/** Buffers initial rows, infers a per-file shredding write plan, and writes
physical rows. */
@@ -43,7 +41,6 @@ public class InferShreddingWritePlanWriter implements
BundleFormatWriter {
private final String compression;
private final List<InternalRow> bufferedRows;
- private final List<BundleRecords> bufferedBundles;
@Nullable private FormatWriter actualWriter;
private boolean planFinalized = false;
@@ -59,7 +56,6 @@ public class InferShreddingWritePlanWriter implements
BundleFormatWriter {
this.out = out;
this.compression = compression;
this.bufferedRows = new ArrayList<>();
- this.bufferedBundles = new ArrayList<>();
}
@Override
@@ -80,14 +76,8 @@ public class InferShreddingWritePlanWriter implements
BundleFormatWriter {
@Override
public void writeBundle(BundleRecords bundle) throws IOException {
if (!planFinalized) {
- final List<InternalRow> rows = new ArrayList<>();
for (InternalRow row : bundle) {
- rows.add(InternalRowUtils.copyInternalRow(row,
writePlanFactory.logicalRowType()));
- }
- bufferedBundles.add(new CopiedBundleRecords(rows));
- totalBufferedRowCount += bundle.rowCount();
- if (totalBufferedRowCount >=
writePlanFactory.inferBufferRowCount()) {
- finalizePlanAndFlush();
+ addElement(row);
}
return;
}
@@ -123,57 +113,15 @@ public class InferShreddingWritePlanWriter implements
BundleFormatWriter {
}
private void finalizePlanAndFlush() throws IOException {
- ShreddingWritePlan writePlan =
writePlanFactory.createWritePlan(collectAllRows());
+ ShreddingWritePlan writePlan =
writePlanFactory.createWritePlan(bufferedRows);
actualWriter =
ShreddingWritePlanWriterFactory.createWriterWithPlan(
writerFactory, writePlanFactory, out, compression,
writePlan);
planFinalized = true;
- if (!bufferedBundles.isEmpty()) {
- BundleFormatWriter bundleWriter = (BundleFormatWriter)
actualWriter;
- for (BundleRecords bundle : bufferedBundles) {
- bundleWriter.writeBundle(bundle);
- }
- bufferedBundles.clear();
- } else {
- for (InternalRow row : bufferedRows) {
- actualWriter.addElement(row);
- }
- bufferedRows.clear();
- }
- }
-
- private List<InternalRow> collectAllRows() {
- if (bufferedBundles.isEmpty()) {
- return bufferedRows;
- }
-
- List<InternalRow> allRows = new ArrayList<>();
- for (BundleRecords bundle : bufferedBundles) {
- for (InternalRow row : bundle) {
- allRows.add(row);
- }
- }
- return allRows;
- }
-
- private static class CopiedBundleRecords implements BundleRecords {
-
- private final List<InternalRow> rows;
-
- private CopiedBundleRecords(List<InternalRow> rows) {
- this.rows = rows;
- }
-
- @Override
- @Nonnull
- public Iterator<InternalRow> iterator() {
- return rows.iterator();
- }
-
- @Override
- public long rowCount() {
- return rows.size();
+ for (InternalRow row : bufferedRows) {
+ actualWriter.addElement(row);
}
+ bufferedRows.clear();
}
}
diff --git
a/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java
b/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java
index fad92112b1..89758b6157 100644
--- a/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java
+++ b/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java
@@ -32,9 +32,11 @@ import org.apache.paimon.data.InternalRow;
public interface BundleRecords extends Iterable<InternalRow> {
/**
- * The total row count of this batch.
+ * The stable, non-negative row count of this batch.
*
- * @return the number of row count.
+ * <p>The count must equal the number of records exposed by {@link
#iterator()}.
+ *
+ * @return the number of records in this batch.
*/
long rowCount();
}
diff --git
a/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java
b/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java
new file mode 100644
index 0000000000..be7a8b197e
--- /dev/null
+++
b/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java
@@ -0,0 +1,213 @@
+/*
+ * 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.shredding;
+
+import org.apache.paimon.data.GenericRow;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.shredding.ShreddingWritePlan;
+import org.apache.paimon.format.BundleFormatWriter;
+import org.apache.paimon.format.FormatWriter;
+import org.apache.paimon.fs.PositionOutputStream;
+import org.apache.paimon.io.BundleRecords;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link InferShreddingWritePlanWriter}. */
+class InferShreddingWritePlanWriterTest {
+
+ private static final RowType ROW_TYPE =
+ RowType.builder().field("value", DataTypes.INT()).build();
+
+ @Test
+ void testMixedRowsAndBundlesPreserveOrderAndInferenceBoundary() throws
Exception {
+ TestingWriterFactory writerFactory = new TestingWriterFactory();
+ TestingWritePlanFactory writePlanFactory = new
TestingWritePlanFactory(3);
+ InferShreddingWritePlanWriter writer =
+ new InferShreddingWritePlanWriter(
+ writerFactory, writePlanFactory, new
NoOpPositionOutputStream(), "none");
+
+ writer.addElement(GenericRow.of(1));
+ writer.writeBundle(bundle(GenericRow.of(2), GenericRow.of(3),
GenericRow.of(4)));
+ writer.addElement(GenericRow.of(5));
+ writer.writeBundle(bundle(GenericRow.of(6), GenericRow.of(7)));
+ writer.close();
+
+ assertThat(writePlanFactory.sampleValues).containsExactly(1, 2, 3);
+ assertThat(writerFactory.writer.values).containsExactly(101, 102, 103,
104, 105, 106, 107);
+ assertThat(writerFactory.writer.bundleWriteCount).isEqualTo(1);
+ }
+
+ @Test
+ void testCloseFinalizesBufferedRowsAndBundle() throws Exception {
+ TestingWriterFactory writerFactory = new TestingWriterFactory();
+ TestingWritePlanFactory writePlanFactory = new
TestingWritePlanFactory(10);
+ InferShreddingWritePlanWriter writer =
+ new InferShreddingWritePlanWriter(
+ writerFactory, writePlanFactory, new
NoOpPositionOutputStream(), "none");
+
+ writer.addElement(GenericRow.of(1));
+ writer.writeBundle(bundle(GenericRow.of(2), GenericRow.of(3)));
+ writer.close();
+
+ assertThat(writePlanFactory.sampleValues).containsExactly(1, 2, 3);
+ assertThat(writerFactory.writer.values).containsExactly(101, 102, 103);
+ }
+
+ private static BundleRecords bundle(InternalRow... rows) {
+ List<InternalRow> records = Arrays.asList(rows);
+ return new BundleRecords() {
+ @Override
+ public Iterator<InternalRow> iterator() {
+ return records.iterator();
+ }
+
+ @Override
+ public long rowCount() {
+ return records.size();
+ }
+ };
+ }
+
+ private static class TestingWriterFactory implements
SupportsShreddingWritePlan {
+
+ private final TestingFormatWriter writer = new TestingFormatWriter();
+
+ @Override
+ public FormatWriter createWithShreddingWritePlan(
+ PositionOutputStream out, String compression,
ShreddingWritePlan writePlan) {
+ return writer;
+ }
+ }
+
+ private static class TestingFormatWriter implements BundleFormatWriter {
+
+ private final List<Integer> values = new ArrayList<>();
+ private int bundleWriteCount;
+
+ @Override
+ public void addElement(InternalRow element) {
+ values.add(element.getInt(0));
+ }
+
+ @Override
+ public void writeBundle(BundleRecords bundle) throws IOException {
+ bundleWriteCount++;
+ for (InternalRow row : bundle) {
+ addElement(row);
+ }
+ }
+
+ @Override
+ public boolean reachTargetSize(boolean suggestedCheck, long
targetSize) {
+ return false;
+ }
+
+ @Override
+ public void close() {}
+ }
+
+ private static class TestingWritePlanFactory implements
ShreddingWritePlanFactory {
+
+ private final int inferBufferRowCount;
+ private final List<Integer> sampleValues = new ArrayList<>();
+
+ private TestingWritePlanFactory(int inferBufferRowCount) {
+ this.inferBufferRowCount = inferBufferRowCount;
+ }
+
+ @Override
+ public RowType logicalRowType() {
+ return ROW_TYPE;
+ }
+
+ @Override
+ public boolean shouldCreateWritePlan() {
+ return true;
+ }
+
+ @Override
+ public boolean shouldInferWritePlan() {
+ return true;
+ }
+
+ @Override
+ public int inferBufferRowCount() {
+ return inferBufferRowCount;
+ }
+
+ @Override
+ public ShreddingWritePlan createWritePlan(List<InternalRow>
sampleRows) {
+ for (InternalRow row : sampleRows) {
+ sampleValues.add(row.getInt(0));
+ }
+ return new TestingWritePlan();
+ }
+ }
+
+ private static class TestingWritePlan implements ShreddingWritePlan {
+
+ @Override
+ public RowType logicalRowType() {
+ return ROW_TYPE;
+ }
+
+ @Override
+ public RowType physicalRowType() {
+ return ROW_TYPE;
+ }
+
+ @Override
+ public InternalRow toPhysicalRow(InternalRow row) {
+ return GenericRow.of(row.getInt(0) + 100);
+ }
+ }
+
+ private static class NoOpPositionOutputStream extends PositionOutputStream
{
+
+ @Override
+ public long getPos() {
+ return 0;
+ }
+
+ @Override
+ public void write(int b) {}
+
+ @Override
+ public void write(byte[] b) {}
+
+ @Override
+ public void write(byte[] b, int off, int len) {}
+
+ @Override
+ public void flush() {}
+
+ @Override
+ public void close() {}
+ }
+}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java
b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java
index 0ebdc5feb7..fe35d51fdd 100644
--- a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java
+++ b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java
@@ -110,9 +110,10 @@ public class RollingFileWriterImpl<T, R> implements
RollingFileWriter<T, R> {
openCurrentWriter();
}
+ long rowCount = bundle.rowCount();
currentWriter.writeBundle(bundle);
- recordCount += bundle.rowCount();
- currentFileRecordCount += bundle.rowCount();
+ recordCount += rowCount;
+ currentFileRecordCount += rowCount;
if (rollingFile(true)) {
closeCurrentWriter();
diff --git
a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java
b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java
index 0372c7f885..73104ab946 100644
---
a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java
+++
b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java
@@ -26,6 +26,9 @@ import org.apache.paimon.utils.LongCounter;
import java.util.function.Supplier;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+import static org.apache.paimon.utils.Preconditions.checkState;
+
/**
* Tracks sequence number range for rows written to a data file.
*
@@ -99,4 +102,16 @@ public class RowDataFileSequenceNumberTracker {
hasNullSeqNumber = true;
}
}
+
+ boolean supportsRowCountUpdate() {
+ return seqNumberFieldIndex == -1;
+ }
+
+ void updateByRowCount(long rowCount) {
+ checkArgument(rowCount >= 0, "Row count must not be negative.");
+ checkState(
+ supportsRowCountUpdate(),
+ "Cannot update sequence numbers by row count when row tracking
is enabled.");
+ seqNumCounter.add(rowCount);
+ }
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java
b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java
index 361a90facf..8d5fb5ede3 100644
--- a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java
+++ b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java
@@ -116,6 +116,15 @@ public class RowDataFileWriter extends
StatsCollectingSingleFileWriter<InternalR
@Override
public void writeBundle(BundleRecords bundle) throws IOException {
+ if (auxiliaryFileWriters.isEmpty()
+ && sequenceNumberTracker.supportsRowCountUpdate()
+ && !requiresPerRecordStats()) {
+ long rowCount = bundle.rowCount();
+ super.writeBundle(bundle);
+ sequenceNumberTracker.updateByRowCount(rowCount);
+ return;
+ }
+
for (InternalRow row : bundle) {
write(row);
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java
b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java
index 8ae8f4bdc0..96f273091d 100644
--- a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java
+++ b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java
@@ -40,6 +40,8 @@ import java.io.UncheckedIOException;
import java.util.Optional;
import java.util.function.Function;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+
/**
* A {@link FileWriter} to produce a single file.
*
@@ -145,6 +147,8 @@ public abstract class SingleFileWriter<T, R> implements
FileWriter<T, R> {
}
try {
+ long rowCount = bundle.rowCount();
+ checkArgument(rowCount >= 0, "Row count must not be negative.");
if (writer instanceof BundleFormatWriter) {
((BundleFormatWriter) writer).writeBundle(bundle);
} else {
@@ -152,7 +156,7 @@ public abstract class SingleFileWriter<T, R> implements
FileWriter<T, R> {
writer.addElement(row);
}
}
- recordCount += bundle.rowCount();
+ recordCount += rowCount;
} catch (Throwable e) {
LOG.warn("Exception occurs when writing file {}. Cleaning up.",
path, e);
abort();
diff --git
a/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java
b/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java
index 773ca08a81..7f660a37f4 100644
---
a/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java
+++
b/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java
@@ -75,6 +75,10 @@ public abstract class StatsCollectingSingleFileWriter<T, R>
extends SingleFileWr
super.writeBundle(bundle);
}
+ protected final boolean requiresPerRecordStats() {
+ return statsRequirePerRecord;
+ }
+
public SimpleColStats[] fieldStats(long fileSize) throws IOException {
Preconditions.checkState(closed, "Cannot access metric unless the
writer is closed.");
if (isStatsDisabled) {
diff --git
a/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java
b/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java
index ac1c9442e2..5eb6fd8a1d 100644
--- a/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java
+++ b/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java
@@ -171,7 +171,7 @@ public class RollingFileWriterTest {
assertThat(files.get(2).rowCount()).isEqualTo(30);
}
- private static BundleRecords bundle(int rowCount) {
+ private static SingleUseBundleRecords bundle(int rowCount) {
List<InternalRow> rows = new ArrayList<>();
for (int i = 0; i < rowCount; i++) {
rows.add(GenericRow.of(i));
diff --git
a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java
b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java
new file mode 100644
index 0000000000..192ba65e69
--- /dev/null
+++ b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java
@@ -0,0 +1,572 @@
+/*
+ * 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.io;
+
+import org.apache.paimon.CoreOptions;
+import org.apache.paimon.data.GenericRow;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.fileindex.FileIndexOptions;
+import org.apache.paimon.format.BundleFormatWriter;
+import org.apache.paimon.format.FileFormat;
+import org.apache.paimon.format.FormatWriter;
+import org.apache.paimon.format.FormatWriterFactory;
+import org.apache.paimon.format.SimpleColStats;
+import org.apache.paimon.format.SupportsDirectWrite;
+import org.apache.paimon.fs.FileIO;
+import org.apache.paimon.fs.Path;
+import org.apache.paimon.fs.PositionOutputStream;
+import org.apache.paimon.manifest.FileSource;
+import org.apache.paimon.options.Options;
+import org.apache.paimon.table.SpecialFields;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+import org.apache.paimon.utils.LongCounter;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Iterator;
+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;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/** Test for {@link RowDataFileWriter}. */
+class RowDataFileWriterTest {
+
+ private static final Path PATH = new Path("file:/tmp/data-file");
+ private static final RowType ROW_TYPE = RowType.builder().field("id",
DataTypes.INT()).build();
+
+ @Test
+ void testEligibleBundleIsForwardedWithoutIteration() throws Exception {
+ FileIO fileIO = fileIO();
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ LongCounter sequenceCounter = new LongCounter(5);
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO,
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ sequenceCounter,
+ new FileIndexOptions());
+ TrackingBundleRecords bundle =
+ trackingRows(GenericRow.of(1), GenericRow.of(2),
GenericRow.of(3));
+
+ writer.writeBundle(bundle);
+
+ assertThat(bundle.iteratorCalls).isZero();
+ assertThat(formatWriter.writtenBundle).isSameAs(bundle);
+ assertThat(formatWriter.bundleWrites).isEqualTo(1);
+ assertThat(formatWriter.rowWrites).isZero();
+ assertThat(writer.recordCount()).isEqualTo(3);
+ assertThat(sequenceCounter.getValue()).isEqualTo(8);
+
+ writer.close();
+ DataFileMeta result = writer.result();
+ assertThat(result.rowCount()).isEqualTo(3);
+ assertThat(result.minSequenceNumber()).isEqualTo(5);
+ assertThat(result.maxSequenceNumber()).isEqualTo(7);
+ }
+
+ @Test
+ void testExtractorStatsAllowBundleForwarding() throws Exception {
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ TestingExtractStatsProducer statsProducer = new
TestingExtractStatsProducer();
+ LongCounter sequenceCounter = new LongCounter(5);
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO(),
+ ROW_TYPE,
+ formatWriter,
+ statsProducer,
+ sequenceCounter,
+ new FileIndexOptions());
+ TrackingBundleRecords bundle =
+ trackingRows(GenericRow.of(1), GenericRow.of(2),
GenericRow.of(3));
+
+ writer.writeBundle(bundle);
+
+ assertThat(bundle.iteratorCalls).isZero();
+ assertThat(formatWriter.writtenBundle).isSameAs(bundle);
+ assertThat(formatWriter.bundleWrites).isEqualTo(1);
+ assertThat(formatWriter.rowWrites).isZero();
+ assertThat(writer.recordCount()).isEqualTo(3);
+ assertThat(sequenceCounter.getValue()).isEqualTo(8);
+
+ writer.close();
+ DataFileMeta result = writer.result();
+ assertThat(statsProducer.extractCalls).isEqualTo(1);
+ assertThat(result.valueStats().minValues().getInt(0)).isEqualTo(1);
+ assertThat(result.valueStats().maxValues().getInt(0)).isEqualTo(3);
+ assertThat(result.valueStats().nullCounts().getLong(0)).isZero();
+ }
+
+ @Test
+ void testPlainFormatWriterFallsBackToRows() throws Exception {
+ TestingFormatWriter formatWriter = new TestingFormatWriter();
+ LongCounter sequenceCounter = new LongCounter();
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO(),
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ sequenceCounter,
+ new FileIndexOptions());
+
+ writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2)));
+
+ assertThat(formatWriter.rowWrites).isEqualTo(2);
+ assertThat(writer.recordCount()).isEqualTo(2);
+ assertThat(sequenceCounter.getValue()).isEqualTo(2);
+ }
+
+ @Test
+ void testBundleFormatWriterCanChooseRowFallback() throws Exception {
+ TestingFallbackBundleFormatWriter formatWriter = new
TestingFallbackBundleFormatWriter();
+ LongCounter sequenceCounter = new LongCounter();
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO(),
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ sequenceCounter,
+ new FileIndexOptions());
+
+ writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2)));
+
+ assertThat(formatWriter.bundleWrites).isEqualTo(1);
+ assertThat(formatWriter.rowWrites).isEqualTo(2);
+ assertThat(writer.recordCount()).isEqualTo(2);
+ assertThat(sequenceCounter.getValue()).isEqualTo(2);
+ }
+
+ @Test
+ void testNegativeBundleRowCountIsRejectedBeforeWriting() throws Exception {
+ FileIO fileIO = fileIO();
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ LongCounter sequenceCounter = new LongCounter();
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO,
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ sequenceCounter,
+ new FileIndexOptions());
+
+ assertThatThrownBy(() -> writer.writeBundle(new
InvalidRowCountBundleRecords(-1)))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Row count must not be negative.");
+
+ assertThat(formatWriter.bundleWrites).isZero();
+ assertThat(formatWriter.closeCalls).isEqualTo(1);
+ assertThat(writer.recordCount()).isZero();
+ assertThat(sequenceCounter.getValue()).isZero();
+ verify(fileIO).deleteQuietly(PATH);
+ }
+
+ @Test
+ void testBundleWriteFailureCleansUpWithoutAdvancingMetadata() throws
Exception {
+ FileIO fileIO = fileIO();
+ IOException failure = new IOException("bundle write failed");
+ TestingThrowingBundleFormatWriter formatWriter =
+ new TestingThrowingBundleFormatWriter(failure);
+ LongCounter sequenceCounter = new LongCounter(5);
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO,
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ sequenceCounter,
+ new FileIndexOptions());
+
+ assertThatThrownBy(() ->
writer.writeBundle(rows(GenericRow.of(1)))).isSameAs(failure);
+
+ assertThat(formatWriter.closeCalls).isEqualTo(1);
+ assertThat(writer.recordCount()).isZero();
+ assertThat(sequenceCounter.getValue()).isEqualTo(5);
+ verify(fileIO).deleteQuietly(PATH);
+ }
+
+ @Test
+ void testPerRecordStatsFallBackToRows() throws Exception {
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ TestingStatsProducer statsProducer = new TestingStatsProducer();
+ LongCounter sequenceCounter = new LongCounter();
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO(),
+ ROW_TYPE,
+ formatWriter,
+ statsProducer,
+ sequenceCounter,
+ new FileIndexOptions());
+
+ writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2)));
+
+ assertThat(formatWriter.bundleWrites).isZero();
+ assertThat(formatWriter.rowWrites).isEqualTo(2);
+ assertThat(statsProducer.collectedRows).isEqualTo(2);
+ assertThat(sequenceCounter.getValue()).isEqualTo(2);
+ }
+
+ @Test
+ void testRowTrackingFallsBackToRows() throws Exception {
+ RowType rowTrackingType =
+ RowType.builder()
+ .field("id", DataTypes.INT())
+ .field(SpecialFields.SEQUENCE_NUMBER.name(),
DataTypes.BIGINT())
+ .build();
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ LongCounter sequenceCounter = new LongCounter(20);
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO(),
+ rowTrackingType,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ sequenceCounter,
+ new FileIndexOptions());
+
+ writer.writeBundle(rows(GenericRow.of(1, 7L), GenericRow.of(2, 11L)));
+
+ assertThat(formatWriter.bundleWrites).isZero();
+ assertThat(formatWriter.rowWrites).isEqualTo(2);
+ assertThat(sequenceCounter.getValue()).isEqualTo(22);
+
+ writer.close();
+ DataFileMeta result = writer.result();
+ assertThat(result.minSequenceNumber()).isEqualTo(7);
+ assertThat(result.maxSequenceNumber()).isEqualTo(11);
+ }
+
+ @Test
+ void testFileIndexFallsBackToRows() throws Exception {
+ Options options = new Options();
+ options.set("file-index.bitmap.columns", "id");
+ FileIndexOptions fileIndexOptions = new FileIndexOptions(new
CoreOptions(options));
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO(),
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ new LongCounter(),
+ fileIndexOptions);
+
+ writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2)));
+
+ assertThat(formatWriter.bundleWrites).isZero();
+ assertThat(formatWriter.rowWrites).isEqualTo(2);
+
+ writer.close();
+ assertThat(writer.result().embeddedIndex()).isNotNull();
+ }
+
+ @Test
+ void testRowSidecarFallsBackToRows() throws Exception {
+ FileIO fileIO = fileIO();
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ TestingFormatWriter sidecarWriter = new TestingFormatWriter();
+ FileFormat rowSidecarFormat = mock(FileFormat.class);
+ when(rowSidecarFormat.createWriterFactory(ROW_TYPE))
+ .thenReturn(new TestingFormatWriterFactory(sidecarWriter));
+ Path rowSidecarPath = new Path("file:/tmp/data-file.row");
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO,
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ new LongCounter(),
+ new FileIndexOptions(),
+ rowSidecarFormat,
+ rowSidecarPath);
+
+ writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2)));
+
+ assertThat(formatWriter.bundleWrites).isZero();
+ assertThat(formatWriter.rowWrites).isEqualTo(2);
+ assertThat(sidecarWriter.rowWrites).isEqualTo(2);
+
+ writer.close();
+
assertThat(writer.result().extraFiles()).containsExactly(rowSidecarPath.getName());
+ }
+
+ private static FileIO fileIO() throws IOException {
+ FileIO fileIO = mock(FileIO.class);
+ when(fileIO.getFileSize(PATH)).thenReturn(123L);
+ return fileIO;
+ }
+
+ private static RowDataFileWriter createWriter(
+ FileIO fileIO,
+ RowType rowType,
+ FormatWriter formatWriter,
+ SimpleStatsProducer statsProducer,
+ LongCounter sequenceCounter,
+ FileIndexOptions fileIndexOptions) {
+ return createWriter(
+ fileIO,
+ rowType,
+ formatWriter,
+ statsProducer,
+ sequenceCounter,
+ fileIndexOptions,
+ null,
+ null);
+ }
+
+ private static RowDataFileWriter createWriter(
+ FileIO fileIO,
+ RowType rowType,
+ FormatWriter formatWriter,
+ SimpleStatsProducer statsProducer,
+ LongCounter sequenceCounter,
+ FileIndexOptions fileIndexOptions,
+ FileFormat rowSidecarFormat,
+ Path rowSidecarPath) {
+ return new RowDataFileWriter(
+ fileIO,
+ new FileWriterContext(
+ new TestingFormatWriterFactory(formatWriter),
statsProducer, "none"),
+ PATH,
+ rowType,
+ 1L,
+ () -> sequenceCounter,
+ fileIndexOptions,
+ FileSource.APPEND,
+ false,
+ false,
+ false,
+ null,
+ rowSidecarFormat,
+ rowSidecarPath);
+ }
+
+ private static BundleRecords rows(InternalRow... rows) {
+ return new ListBundleRecords(Arrays.asList(rows));
+ }
+
+ private static TrackingBundleRecords trackingRows(InternalRow... rows) {
+ return new TrackingBundleRecords(Arrays.asList(rows));
+ }
+
+ private static class TestingFormatWriterFactory
+ implements FormatWriterFactory, SupportsDirectWrite {
+
+ private final FormatWriter writer;
+
+ private TestingFormatWriterFactory(FormatWriter writer) {
+ this.writer = writer;
+ }
+
+ @Override
+ public FormatWriter create(PositionOutputStream out, String
compression) {
+ return writer;
+ }
+
+ @Override
+ public FormatWriter create(FileIO fileIO, Path path, String
compression) {
+ return writer;
+ }
+ }
+
+ private static class TestingFormatWriter implements FormatWriter {
+
+ int rowWrites;
+ int closeCalls;
+
+ @Override
+ public void addElement(InternalRow element) {
+ rowWrites++;
+ }
+
+ @Override
+ public boolean reachTargetSize(boolean suggestedCheck, long
targetSize) {
+ return false;
+ }
+
+ @Override
+ public void close() {
+ closeCalls++;
+ }
+ }
+
+ private static class TestingBundleFormatWriter extends TestingFormatWriter
+ implements BundleFormatWriter {
+
+ private int bundleWrites;
+ private BundleRecords writtenBundle;
+
+ @Override
+ public void writeBundle(BundleRecords bundle) {
+ bundleWrites++;
+ writtenBundle = bundle;
+ }
+ }
+
+ private static class TestingFallbackBundleFormatWriter extends
TestingFormatWriter
+ implements BundleFormatWriter {
+
+ private int bundleWrites;
+
+ @Override
+ public void writeBundle(BundleRecords bundle) {
+ bundleWrites++;
+ for (InternalRow row : bundle) {
+ addElement(row);
+ }
+ }
+ }
+
+ private static class TestingThrowingBundleFormatWriter extends
TestingFormatWriter
+ implements BundleFormatWriter {
+
+ private final IOException failure;
+
+ private TestingThrowingBundleFormatWriter(IOException failure) {
+ this.failure = failure;
+ }
+
+ @Override
+ public void writeBundle(BundleRecords bundle) throws IOException {
+ throw failure;
+ }
+ }
+
+ private static class TestingStatsProducer implements SimpleStatsProducer {
+
+ private int collectedRows;
+
+ @Override
+ public boolean isStatsDisabled() {
+ return false;
+ }
+
+ @Override
+ public boolean requirePerRecord() {
+ return true;
+ }
+
+ @Override
+ public void collect(InternalRow row) {
+ collectedRows++;
+ }
+
+ @Override
+ public SimpleColStats[] extract(FileIO fileIO, Path path, long length)
{
+ return new SimpleColStats[] {SimpleColStats.NONE};
+ }
+ }
+
+ private static class TestingExtractStatsProducer implements
SimpleStatsProducer {
+
+ private int extractCalls;
+
+ @Override
+ public boolean isStatsDisabled() {
+ return false;
+ }
+
+ @Override
+ public boolean requirePerRecord() {
+ return false;
+ }
+
+ @Override
+ public void collect(InternalRow row) {
+ throw new AssertionError("Extractor-backed statistics must not
collect rows.");
+ }
+
+ @Override
+ public SimpleColStats[] extract(FileIO fileIO, Path path, long length)
{
+ extractCalls++;
+ return new SimpleColStats[] {new SimpleColStats(1, 3, 0L)};
+ }
+ }
+
+ private static class InvalidRowCountBundleRecords implements BundleRecords
{
+
+ private final long rowCount;
+
+ private InvalidRowCountBundleRecords(long rowCount) {
+ this.rowCount = rowCount;
+ }
+
+ @Override
+ public Iterator<InternalRow> iterator() {
+ throw new AssertionError("Invalid row count must be rejected
before row iteration.");
+ }
+
+ @Override
+ public long rowCount() {
+ return rowCount;
+ }
+ }
+
+ private static class TrackingBundleRecords implements BundleRecords {
+
+ private final List<InternalRow> rows;
+ private int iteratorCalls;
+
+ private TrackingBundleRecords(List<InternalRow> rows) {
+ this.rows = rows;
+ }
+
+ @Override
+ public Iterator<InternalRow> iterator() {
+ iteratorCalls++;
+ return rows.iterator();
+ }
+
+ @Override
+ public long rowCount() {
+ return rows.size();
+ }
+ }
+
+ private static class ListBundleRecords implements BundleRecords {
+
+ private final List<InternalRow> rows;
+
+ private ListBundleRecords(List<InternalRow> rows) {
+ this.rows = rows;
+ }
+
+ @Override
+ public Iterator<InternalRow> iterator() {
+ return rows.iterator();
+ }
+
+ @Override
+ public long rowCount() {
+ return rows.size();
+ }
+ }
+}
diff --git
a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java
b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java
index ee1f572462..a23d95489f 100644
---
a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java
+++
b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java
@@ -19,6 +19,7 @@
package org.apache.paimon.format.lance;
import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.vector.ArrowFormatWriter;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.format.BundleFormatWriter;
@@ -65,12 +66,20 @@ public class LanceRecordsWriter implements
BundleFormatWriter {
@Override
public void writeBundle(BundleRecords bundleRecords) throws IOException {
if (bundleRecords instanceof ArrowBundleRecords) {
- add(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot());
- } else {
- for (InternalRow row : bundleRecords) {
- addElement(row);
+ ArrowBundleRecords arrowBundle = (ArrowBundleRecords)
bundleRecords;
+ VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot();
+ if (arrowFormatWriter.isArrowBundleSchemaCompatible(arrowBundle)
+ && ArrowUtils.hasSameRootAllocator(root,
arrowFormatWriter.getAllocator())) {
+ flush();
+
nativeWriter.ensureInitialized(arrowFormatWriter.getAllocator());
+ add(root);
+ return;
}
}
+
+ for (InternalRow row : bundleRecords) {
+ addElement(row);
+ }
}
public void add(VectorSchemaRoot vsr) throws IOException {
diff --git
a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java
b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java
index f9329907f9..b3e29eb0a0 100644
---
a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java
+++
b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java
@@ -18,6 +18,8 @@
package org.apache.paimon.format.lance.jni;
+import org.apache.paimon.arrow.ArrowUtils;
+
import com.lancedb.lance.file.LanceFileWriter;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.FieldVector;
@@ -32,6 +34,7 @@ public class LanceWriter {
private final String path;
private final Map<String, String> storageOptions;
private LanceFileWriter writer;
+ private BufferAllocator allocator;
private long bytesWritten = 0;
public LanceWriter(String path, Map<String, String> storageOptions) {
@@ -44,12 +47,25 @@ public class LanceWriter {
}
public void writeVsr(VectorSchemaRoot vsr) throws IOException {
- initWriteLazy(vsr.getVector(0).getAllocator());
+ BufferAllocator sourceAllocator = vsr.getVector(0).getAllocator();
+ initWriteLazy(sourceAllocator);
+ if (!ArrowUtils.hasSameRootAllocator(vsr, allocator)) {
+ throw new IllegalArgumentException(
+ "Lance writer cannot consume Arrow buffers from a
different allocator root.");
+ }
this.bytesWritten +=
vsr.getFieldVectors().stream().mapToLong(FieldVector::getBufferSize).sum();
this.writer.write(vsr);
}
+ /**
+ * Initializes the native writer with an allocator whose lifetime is owned
by the surrounding
+ * format writer.
+ */
+ public void ensureInitialized(BufferAllocator bufferAllocator) throws
IOException {
+ initWriteLazy(bufferAllocator);
+ }
+
public void close() throws IOException {
if (writer != null) {
try {
@@ -58,6 +74,7 @@ public class LanceWriter {
throw new IOException(e);
}
this.writer = null;
+ this.allocator = null;
}
}
@@ -68,6 +85,7 @@ public class LanceWriter {
private void initWriteLazy(BufferAllocator bufferAllocator) throws
IOException {
if (writer == null) {
writer = LanceFileWriter.open(path, bufferAllocator, null,
storageOptions);
+ allocator = bufferAllocator;
}
}
}
diff --git
a/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java
b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java
new file mode 100644
index 0000000000..7d5161d5dc
--- /dev/null
+++
b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java
@@ -0,0 +1,297 @@
+/*
+ * 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.lance;
+
+import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
+import org.apache.paimon.arrow.vector.ArrowFormatWriter;
+import org.apache.paimon.data.GenericRow;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.serializer.InternalRowSerializer;
+import org.apache.paimon.format.lance.jni.LanceWriter;
+import org.apache.paimon.fs.Path;
+import org.apache.paimon.reader.RecordReaderIterator;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link LanceRecordsWriter}. */
+class LanceRecordsWriterTest {
+
+ @Test
+ void testArrowBundlePreservesRowBundleRowOrder() throws Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024,
true);
+ BufferAllocator writerAllocator = arrowWriter.getAllocator();
+ CapturingLanceWriter nativeWriter = new CapturingLanceWriter();
+ LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L,
arrowWriter, nativeWriter);
+
+ writer.addElement(GenericRow.of(1));
+ try (BufferAllocator sourceAllocator =
+ arrowWriter
+ .getAllocator()
+ .newChildAllocator("lance-bundle-test", 0,
Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("value"), 2);
+ root.setRowCount(1);
+ writer.writeBundle(new ArrowBundleRecords(root, rowType, true));
+ }
+ writer.addElement(GenericRow.of(3));
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(3);
+
assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(1);
+
assertThat(nativeWriter.snapshots.get(1).values.get(0)).containsExactly(2);
+
assertThat(nativeWriter.snapshots.get(2).values.get(0)).containsExactly(3);
+
assertThat(nativeWriter.initializedAllocator).isSameAs(writerAllocator);
+ }
+
+ @Test
+ void testNativeWriterUsesDirectSameRootAndFallbackCrossRoot(@TempDir
java.nio.file.Path tempDir)
+ throws Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ String file = tempDir.resolve("bundle_" +
UUID.randomUUID()).toString();
+ ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024,
true);
+ TrackingLanceWriter nativeWriter = new TrackingLanceWriter(file);
+ try (LanceRecordsWriter writer =
+ new LanceRecordsWriter(
+ nativeWriter::getWrittenPosition, arrowWriter,
nativeWriter)) {
+ writer.addElement(GenericRow.of(1));
+ try (BufferAllocator sourceAllocator =
+ arrowWriter
+ .getAllocator()
+ .newChildAllocator(
+ "lance-direct-native-test", 0,
Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType,
sourceAllocator)) {
+ nativeWriter.sameRootBundle = root;
+ setInt((IntVector) root.getVector("value"), 2);
+ root.setRowCount(1);
+ writer.writeBundle(new ArrowBundleRecords(root, rowType,
true));
+ }
+
+ try (BufferAllocator sourceAllocator = new
org.apache.arrow.memory.RootAllocator();
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType,
sourceAllocator)) {
+ nativeWriter.crossRootBundle = root;
+ setInt((IntVector) root.getVector("value"), 3);
+ root.setRowCount(1);
+ writer.writeBundle(new ArrowBundleRecords(root, rowType,
true));
+ }
+ writer.addElement(GenericRow.of(4));
+ }
+
+ assertThat(nativeWriter.sameRootBundleWrites).isEqualTo(1);
+ assertThat(nativeWriter.crossRootBundleWrites).isZero();
+
+ InternalRowSerializer serializer = new InternalRowSerializer(rowType);
+ List<Integer> actual = new ArrayList<>();
+ try (RecordReaderIterator<InternalRow> iterator =
+ new RecordReaderIterator<>(
+ new LanceRecordsReader(
+ new Path(file), null, rowType, 1024,
Collections.emptyMap()))) {
+ while (iterator.hasNext()) {
+ actual.add(serializer.copy(iterator.next()).getInt(0));
+ }
+ }
+ assertThat(actual).containsExactly(1, 2, 3, 4);
+ }
+
+ @Test
+ void testReorderedArrowBundleFallsBackToRows() throws Exception {
+ RowType writerType =
+ RowType.builder().field("a", DataTypes.INT()).field("b",
DataTypes.INT()).build();
+ RowType sourceType =
+ RowType.builder().field("b", DataTypes.INT()).field("a",
DataTypes.INT()).build();
+ ArrowFormatWriter arrowWriter = new ArrowFormatWriter(writerType,
1024, true);
+ CapturingLanceWriter nativeWriter = new CapturingLanceWriter();
+ LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L,
arrowWriter, nativeWriter);
+
+ try (BufferAllocator sourceAllocator =
+ arrowWriter
+ .getAllocator()
+ .newChildAllocator("lance-schema-test", 0,
Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(sourceType,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("b"), 20);
+ setInt((IntVector) root.getVector("a"), 10);
+ root.setRowCount(1);
+ writer.writeBundle(new ArrowBundleRecords(root, writerType, true));
+ }
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(1);
+ Snapshot snapshot = nativeWriter.snapshots.get(0);
+ assertThat(snapshot.fieldNames).containsExactly("a", "b");
+ assertThat(snapshot.values.get(0)).containsExactly(10);
+ assertThat(snapshot.values.get(1)).containsExactly(20);
+ }
+
+ @Test
+ void testMissingColumnArrowBundleFallsBackToRows() throws Exception {
+ RowType writerType =
+ RowType.builder().field("a", DataTypes.INT()).field("b",
DataTypes.INT()).build();
+ RowType sourceType = RowType.builder().field("a",
DataTypes.INT()).build();
+ ArrowFormatWriter arrowWriter = new ArrowFormatWriter(writerType,
1024, true);
+ CapturingLanceWriter nativeWriter = new CapturingLanceWriter();
+ LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L,
arrowWriter, nativeWriter);
+
+ try (BufferAllocator sourceAllocator =
+ arrowWriter
+ .getAllocator()
+
.newChildAllocator("lance-missing-column-test", 0, Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(sourceType,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("a"), 10);
+ root.setRowCount(1);
+ writer.writeBundle(new ArrowBundleRecords(root, writerType, true));
+ }
+ writer.close();
+
+ assertThat(nativeWriter.snapshots).hasSize(1);
+ Snapshot snapshot = nativeWriter.snapshots.get(0);
+ assertThat(snapshot.fieldNames).containsExactly("a", "b");
+ assertThat(snapshot.values.get(0)).containsExactly(10);
+ assertThat(snapshot.values.get(1)).containsExactly((Integer) null);
+ }
+
+ @Test
+ void testDifferentAllocatorRootFallsBackToRows() throws Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024,
true);
+ CapturingLanceWriter nativeWriter = new CapturingLanceWriter();
+ LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L,
arrowWriter, nativeWriter);
+
+ try (BufferAllocator sourceAllocator = new
org.apache.arrow.memory.RootAllocator();
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType,
sourceAllocator)) {
+ nativeWriter.disallowedRoot = root;
+ setInt((IntVector) root.getVector("value"), 10);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, rowType, true));
+ }
+ writer.close();
+
+ assertThat(nativeWriter.disallowedRootWrites).isZero();
+ assertThat(nativeWriter.snapshots).hasSize(1);
+
assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(10);
+ }
+
+ private static void setInt(IntVector vector, int value) {
+ vector.allocateNew(1);
+ vector.setSafe(0, value);
+ vector.setValueCount(1);
+ }
+
+ private static class CapturingLanceWriter extends LanceWriter {
+
+ private final List<Snapshot> snapshots = new ArrayList<>();
+ private BufferAllocator initializedAllocator;
+ private VectorSchemaRoot disallowedRoot;
+ private int disallowedRootWrites;
+
+ private CapturingLanceWriter() {
+ super("unused", Collections.emptyMap());
+ }
+
+ @Override
+ public void ensureInitialized(BufferAllocator bufferAllocator) {
+ initializedAllocator = bufferAllocator;
+ }
+
+ @Override
+ public void writeVsr(VectorSchemaRoot root) {
+ if (root == disallowedRoot) {
+ disallowedRootWrites++;
+ }
+ List<String> fieldNames =
+ root.getSchema().getFields().stream()
+ .map(field -> field.getName())
+ .collect(Collectors.toList());
+ List<List<Integer>> values = new ArrayList<>();
+ for (int column = 0; column < root.getFieldVectors().size();
column++) {
+ IntVector vector = (IntVector) root.getVector(column);
+ List<Integer> columnValues = new ArrayList<>();
+ for (int row = 0; row < root.getRowCount(); row++) {
+ columnValues.add(vector.isNull(row) ? null :
vector.get(row));
+ }
+ values.add(columnValues);
+ }
+ snapshots.add(new Snapshot(fieldNames, values));
+ }
+
+ @Override
+ public void close() throws IOException {}
+
+ @Override
+ public String path() {
+ return "unused";
+ }
+ }
+
+ private static class TrackingLanceWriter extends LanceWriter {
+
+ private VectorSchemaRoot sameRootBundle;
+ private VectorSchemaRoot crossRootBundle;
+ private int sameRootBundleWrites;
+ private int crossRootBundleWrites;
+
+ private TrackingLanceWriter(String path) {
+ super(path, Collections.emptyMap());
+ }
+
+ @Override
+ public void writeVsr(VectorSchemaRoot root) throws IOException {
+ if (root == sameRootBundle) {
+ sameRootBundleWrites++;
+ }
+ if (root == crossRootBundle) {
+ crossRootBundleWrites++;
+ }
+ super.writeVsr(root);
+ }
+ }
+
+ private static class Snapshot {
+
+ private final List<String> fieldNames;
+ private final List<List<Integer>> values;
+
+ private Snapshot(List<String> fieldNames, List<List<Integer>> values) {
+ this.fieldNames = fieldNames;
+ this.values = values;
+ }
+ }
+}
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 944f78c4b1..1a3e4d4149 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
@@ -19,6 +19,7 @@
package org.apache.paimon.format.mosaic;
import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.vector.ArrowFormatWriter;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.format.BundleFormatWriter;
@@ -125,13 +126,21 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
@Override
public void writeBundle(BundleRecords bundleRecords) {
if (bundleRecords instanceof ArrowBundleRecords) {
- flush();
- nativeWriter.write(((ArrowBundleRecords)
bundleRecords).getVectorSchemaRoot());
- } else {
- for (InternalRow row : bundleRecords) {
- addElement(row);
+ 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);
+ return;
}
}
+
+ for (InternalRow row : bundleRecords) {
+ addElement(row);
+ }
}
@Override
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 60efceed08..931f4d7033 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
@@ -18,10 +18,13 @@
package org.apache.paimon.format.mosaic;
+import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.data.serializer.InternalRowSerializer;
+import org.apache.paimon.format.BundleFormatWriter;
import org.apache.paimon.format.FileFormatFactory;
import org.apache.paimon.format.FormatReaderContext;
import org.apache.paimon.format.FormatReaderFactory;
@@ -37,6 +40,9 @@ import org.apache.paimon.reader.RecordReader;
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.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
@@ -44,6 +50,7 @@ import org.junit.jupiter.api.io.TempDir;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
+import java.util.Iterator;
import java.util.List;
import java.util.UUID;
@@ -80,6 +87,38 @@ class MosaicReaderWriterTest {
assertThat(result.get(1).getString(1).toString()).isEqualTo("world");
}
+ @Test
+ void testPublicWriterFactoryFallsBackForCrossRootArrowBundle() throws
IOException {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ Path path = newPath();
+ MosaicFileFormat format = createFormat();
+ FormatWriterFactory writerFactory =
format.createWriterFactory(rowType);
+ LocalFileIO fileIO = new LocalFileIO();
+
+ try (RootAllocator sourceAllocator = new RootAllocator();
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType,
sourceAllocator);
+ BundleFormatWriter writer =
+ (BundleFormatWriter)
+
writerFactory.create(fileIO.newOutputStream(path, false), "zstd")) {
+ IntVector vector = (IntVector) root.getVector("value");
+ vector.allocateNew(2);
+ vector.setSafe(0, 7);
+ vector.setSafe(1, 9);
+ vector.setValueCount(2);
+ root.setRowCount(2);
+
+ CountingArrowBundleRecords bundle = new
CountingArrowBundleRecords(root, rowType);
+ writer.writeBundle(bundle);
+ assertThat(bundle.iteratorCalls()).isEqualTo(1);
+ }
+
+ List<InternalRow> result = readAll(rowType, rowType, path, null);
+ assertThat(result).hasSize(2);
+ assertThat(result.get(0).getInt(0)).isEqualTo(7);
+ assertThat(result.get(1).getInt(0)).isEqualTo(9);
+ }
+
@Test
void testNullValues() throws IOException {
RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING());
@@ -358,4 +397,23 @@ class MosaicReaderWriterTest {
return false;
}
}
+
+ private static class CountingArrowBundleRecords extends ArrowBundleRecords
{
+
+ private int iteratorCalls;
+
+ private CountingArrowBundleRecords(VectorSchemaRoot root, RowType
rowType) {
+ super(root, rowType, true);
+ }
+
+ @Override
+ public Iterator<InternalRow> iterator() {
+ iteratorCalls++;
+ return super.iterator();
+ }
+
+ private int iteratorCalls() {
+ return iteratorCalls;
+ }
+ }
}
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 58a8dc5252..344d5a0022 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
@@ -18,12 +18,18 @@
package org.apache.paimon.format.mosaic;
+import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.format.FileFormatFactory;
+import org.apache.paimon.mosaic.MosaicWriter;
import org.apache.paimon.options.Options;
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.junit.jupiter.api.Test;
import java.io.ByteArrayOutputStream;
@@ -31,15 +37,22 @@ import java.util.Collections;
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.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.same;
+import static org.mockito.Mockito.verify;
/** Test for {@link MosaicRecordsWriter}. */
class MosaicRecordsWriterTest {
+ private static final FileFormatFactory.FormatContext FORMAT_CONTEXT =
+ new FileFormatFactory.FormatContext(new Options(), 1024, 1024);
+
@Test
void testConstructorFailureClosesCreatedResources() {
RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING());
- FileFormatFactory.FormatContext formatContext =
- new FileFormatFactory.FormatContext(new Options(), 1024, 1024);
CloseCountingRootAllocator allocator = new
CloseCountingRootAllocator();
RuntimeException failure = new RuntimeException("native writer
failed");
@@ -48,7 +61,7 @@ class MosaicRecordsWriterTest {
new MosaicRecordsWriter(
new ByteArrayOutputStream(),
rowType,
- formatContext,
+ FORMAT_CONTEXT,
Collections.emptyList(),
null,
allocator,
@@ -60,6 +73,108 @@ class MosaicRecordsWriterTest {
assertThat(allocator.closeCount()).isEqualTo(1);
}
+ @Test
+ void testCompatibleSameRootArrowBundleUsesDirectWrite() throws Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ RootAllocator writerAllocator = new RootAllocator();
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ MosaicRecordsWriter writer = createWriter(rowType, writerAllocator,
nativeWriter);
+
+ try (BufferAllocator sourceAllocator =
+
writerAllocator.newChildAllocator("mosaic-direct-test", 0, Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("value"), 1);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, rowType, true));
+
+ verify(nativeWriter).write(same(root));
+ }
+ writer.close();
+ }
+
+ @Test
+ void testDifferentAllocatorRootFallsBackToRows() throws Exception {
+ RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
+ RootAllocator writerAllocator = new RootAllocator();
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ MosaicRecordsWriter writer = createWriter(rowType, writerAllocator,
nativeWriter);
+
+ try (RootAllocator sourceAllocator = new RootAllocator();
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(rowType,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("value"), 1);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, rowType, true));
+
+ verify(nativeWriter, never()).write(same(root));
+ }
+ writer.close();
+
+ verify(nativeWriter).write(any(VectorSchemaRoot.class));
+ }
+
+ @Test
+ void testReorderedArrowBundleFallsBackToRows() throws Exception {
+ RowType writerType =
+ RowType.builder().field("a", DataTypes.INT()).field("b",
DataTypes.INT()).build();
+ RowType sourceType =
+ RowType.builder().field("b", DataTypes.INT()).field("a",
DataTypes.INT()).build();
+ RootAllocator writerAllocator = new RootAllocator();
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ MosaicRecordsWriter writer = createWriter(writerType, writerAllocator,
nativeWriter);
+ doAnswer(
+ invocation -> {
+ VectorSchemaRoot written =
invocation.getArgument(0);
+
assertThat(written.getSchema().getFields().get(0).getName())
+ .isEqualTo("a");
+
assertThat(written.getSchema().getFields().get(1).getName())
+ .isEqualTo("b");
+ assertThat(((IntVector)
written.getVector("a")).get(0)).isEqualTo(10);
+ assertThat(((IntVector)
written.getVector("b")).get(0)).isEqualTo(20);
+ return null;
+ })
+ .when(nativeWriter)
+ .write(any(VectorSchemaRoot.class));
+
+ try (BufferAllocator sourceAllocator =
+ writerAllocator.newChildAllocator(
+ "mosaic-reordered-test", 0, Long.MAX_VALUE);
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(sourceType,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("b"), 20);
+ setInt((IntVector) root.getVector("a"), 10);
+ root.setRowCount(1);
+
+ writer.writeBundle(new ArrowBundleRecords(root, writerType, true));
+
+ verify(nativeWriter, never()).write(same(root));
+ }
+ writer.close();
+
+ verify(nativeWriter).write(any(VectorSchemaRoot.class));
+ }
+
+ private static MosaicRecordsWriter createWriter(
+ RowType rowType, RootAllocator allocator, MosaicWriter
nativeWriter) {
+ return new MosaicRecordsWriter(
+ new ByteArrayOutputStream(),
+ rowType,
+ FORMAT_CONTEXT,
+ Collections.emptyList(),
+ null,
+ allocator,
+ (outputStream, arrowSchema, options, bufferAllocator) ->
nativeWriter);
+ }
+
+ private static void setInt(IntVector vector, int value) {
+ vector.allocateNew(1);
+ vector.setSafe(0, value);
+ vector.setValueCount(1);
+ }
+
private static class CloseCountingRootAllocator extends RootAllocator {
private int closeCount;
diff --git
a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java
b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java
index 7798fa09d4..502b4185e8 100644
---
a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java
+++
b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java
@@ -19,6 +19,7 @@
package org.apache.paimon.format.vortex;
import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.vector.ArrowCStruct;
import org.apache.paimon.arrow.vector.ArrowFormatCWriter;
import org.apache.paimon.data.InternalRow;
@@ -38,6 +39,9 @@ import org.apache.arrow.vector.types.pojo.Schema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import javax.annotation.Nullable;
+
+import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
@@ -65,6 +69,11 @@ public class VortexRecordsWriter implements
BundleFormatWriter {
private final List<AutoCloseable> retainedResources;
private ArrowFormatCWriter currentWriter;
+ // Direct bundles are copied into this writer-owned allocator. Vortex
synchronously imports
+ // independent Arrow buffer references before writeBatch returns, while
nativeWriter.close()
+ // completes outstanding asynchronous writes before this allocator is
closed.
+ @Nullable private RootAllocator bundleAllocator;
+
private long jniCost = 0;
private long ffiBytes = 0;
@@ -102,13 +111,17 @@ public class VortexRecordsWriter implements
BundleFormatWriter {
@Override
public void writeBundle(BundleRecords bundleRecords) throws IOException {
if (bundleRecords instanceof ArrowBundleRecords) {
- flush();
- writeBundleVsr(((ArrowBundleRecords)
bundleRecords).getVectorSchemaRoot());
- } else {
- for (InternalRow row : bundleRecords) {
- addElement(row);
+ ArrowBundleRecords arrowBundle = (ArrowBundleRecords)
bundleRecords;
+ if
(currentWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle)) {
+ flush();
+ writeBundleVsr(arrowBundle.getVectorSchemaRoot());
+ return;
}
}
+
+ for (InternalRow row : bundleRecords) {
+ addElement(row);
+ }
}
@Override
@@ -136,10 +149,26 @@ public class VortexRecordsWriter implements
BundleFormatWriter {
// Release all retained resources now that async writes are done.
for (AutoCloseable res : retainedResources) {
- closeQuietly(res);
+ try {
+ closeQuietly(res);
+ } catch (Throwable t) {
+ throwable = addSuppressed(throwable, t);
+ }
}
retainedResources.clear();
- closeQuietly(currentWriter);
+ try {
+ closeQuietly(currentWriter);
+ } catch (Throwable t) {
+ throwable = addSuppressed(throwable, t);
+ }
+
+ if (bundleAllocator != null) {
+ try {
+ bundleAllocator.close();
+ } catch (Throwable t) {
+ throwable = addSuppressed(throwable, t);
+ }
+ }
try {
session.close();
@@ -170,34 +199,38 @@ public class VortexRecordsWriter implements
BundleFormatWriter {
}
}
- /** Write an external VSR (from writeBundle) via IPC copy into an
independent allocator. */
+ /** Write an external VSR (from writeBundle) via IPC copy into the
writer-owned allocator. */
private void writeBundleVsr(VectorSchemaRoot vsr) throws IOException {
ffiBytes += bufferBytes(vsr);
- byte[] ipc = org.apache.paimon.arrow.ArrowUtils.serializeToIpc(vsr);
- RootAllocator bundleAllocator = new RootAllocator(Long.MAX_VALUE);
- try {
- ArrowStreamReader ipcReader =
- new ArrowStreamReader(new
java.io.ByteArrayInputStream(ipc), bundleAllocator);
- ipcReader.loadNextBatch();
- VectorSchemaRoot copy = ipcReader.getVectorSchemaRoot();
-
- ArrowArray arrowArray = ArrowArray.allocateNew(bundleAllocator);
- ArrowSchema arrowSchema = ArrowSchema.allocateNew(bundleAllocator);
- Data.exportVectorSchemaRoot(bundleAllocator, copy, null,
arrowArray, arrowSchema);
+ byte[] ipc = ArrowUtils.serializeToIpc(vsr);
+ RootAllocator allocator = bundleAllocator();
+ try (ArrowStreamReader reader =
+ new ArrowStreamReader(new ByteArrayInputStream(ipc),
allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator);
+ ArrowSchema schema = ArrowSchema.allocateNew(allocator)) {
+ if (!reader.loadNextBatch()) {
+ throw new IOException("Arrow IPC copy did not contain a record
batch.");
+ }
+ Data.exportVectorSchemaRoot(
+ allocator, reader.getVectorSchemaRoot(), null, array,
schema);
+
long t1 = System.currentTimeMillis();
- nativeWriter.writeBatch(arrowArray.memoryAddress(),
arrowSchema.memoryAddress());
+ try {
+ nativeWriter.writeBatch(array.memoryAddress(),
schema.memoryAddress());
+ } finally {
+ ArrowUtils.releaseCDataIfNeeded(array, schema);
+ }
jniCost += (System.currentTimeMillis() - t1);
- // Retain all resources that own the exported C Data buffers and
release
- // callbacks. Rust holds async zero-copy references via
Arc<FFI_ArrowArray>.
- // Order matters: close ipcReader (owns VectorSchemaRoot) before
allocator.
- retainedResources.add(ipcReader);
- retainedResources.add(bundleAllocator);
- } catch (Exception e) {
- closeQuietly(bundleAllocator);
- throw e instanceof IOException ? (IOException) e : new
IOException(e);
}
}
+ private RootAllocator bundleAllocator() {
+ if (bundleAllocator == null) {
+ bundleAllocator = new RootAllocator(Long.MAX_VALUE);
+ }
+ return bundleAllocator;
+ }
+
private static long bufferBytes(VectorSchemaRoot vsr) {
long bytes = 0;
for (int i = 0; i < vsr.getFieldVectors().size(); i++) {
diff --git
a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java
b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java
index e97e3cf80f..81dab2774d 100644
---
a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java
+++
b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java
@@ -19,6 +19,7 @@
package org.apache.paimon.format.vortex;
import org.apache.paimon.arrow.ArrowBundleRecords;
+import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.vector.ArrowFormatWriter;
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.GenericRow;
@@ -45,6 +46,9 @@ import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.RoaringBitmap32;
import dev.vortex.jni.NativeRuntime;
+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.api.io.TempDir;
@@ -264,7 +268,7 @@ public class VortexReaderWriterTest {
}
@Test
- public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers(
+ public void
testArrowBundleRecordsWriteDoesNotBorrowCallerBuffersAcrossBatches(
@TempDir java.nio.file.Path tempDir) throws Exception {
RowType rowType =
RowType.builder()
@@ -289,6 +293,16 @@ public class VortexReaderWriterTest {
arrowWriter.write(GenericRow.of(2,
BinaryString.fromString("world")));
arrowWriter.flush();
+ ((BundleFormatWriter) writer)
+ .writeBundle(
+ new ArrowBundleRecords(
+ arrowWriter.getVectorSchemaRoot(),
rowType, true));
+
+ arrowWriter.reset();
+ arrowWriter.write(GenericRow.of(3,
BinaryString.fromString("second")));
+ arrowWriter.write(GenericRow.of(4,
BinaryString.fromString("batch")));
+ arrowWriter.flush();
+
((BundleFormatWriter) writer)
.writeBundle(
new ArrowBundleRecords(
@@ -312,11 +326,114 @@ public class VortexReaderWriterTest {
actualRows.add(serializer.copy(iterator.next()));
}
- assertEquals(2, actualRows.size());
+ assertEquals(4, actualRows.size());
assertEquals(1, actualRows.get(0).getInt(0));
assertEquals(BinaryString.fromString("hello"),
actualRows.get(0).getString(1));
assertEquals(2, actualRows.get(1).getInt(0));
assertEquals(BinaryString.fromString("world"),
actualRows.get(1).getString(1));
+ assertEquals(3, actualRows.get(2).getInt(0));
+ assertEquals(BinaryString.fromString("second"),
actualRows.get(2).getString(1));
+ assertEquals(4, actualRows.get(3).getInt(0));
+ assertEquals(BinaryString.fromString("batch"),
actualRows.get(3).getString(1));
+ }
+ }
+
+ @Test
+ public void testArrowBundlePreservesRowBundleRowOrder(@TempDir
java.nio.file.Path tempDir)
+ throws Exception {
+ RowType rowType =
+ RowType.builder()
+ .field("f_int", DataTypes.INT())
+ .field("f_string", DataTypes.STRING())
+ .build();
+ VortexFileFormat format =
+ new VortexFileFormatFactory()
+ .create(new FileFormatFactory.FormatContext(new
Options(), 1024, 1024));
+ FileIO fileIO = new LocalFileIO();
+ Path testFile =
+ new Path(new Path(tempDir.toUri()), "test_bundle_order_" +
UUID.randomUUID());
+
+ try (FormatWriter writer =
+ ((SupportsDirectWrite)
format.createWriterFactory(rowType))
+ .create(fileIO, testFile, "");
+ ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType,
1024, true)) {
+ writer.addElement(GenericRow.of(1,
BinaryString.fromString("row")));
+
+ arrowWriter.write(GenericRow.of(2,
BinaryString.fromString("bundle")));
+ arrowWriter.flush();
+ ((BundleFormatWriter) writer)
+ .writeBundle(
+ new ArrowBundleRecords(
+ arrowWriter.getVectorSchemaRoot(),
rowType, true));
+
+ writer.addElement(GenericRow.of(3,
BinaryString.fromString("row")));
+ }
+
+ InternalRowSerializer serializer = new InternalRowSerializer(rowType);
+ FormatReaderFactory readerFactory =
format.createReaderFactory(rowType, rowType, null);
+ try (RecordReader<InternalRow> reader =
+ readerFactory.createReader(
+ new FormatReaderContext(
+ fileIO, testFile,
fileIO.getFileSize(testFile), null));
+ RecordReaderIterator<InternalRow> iterator = new
RecordReaderIterator<>(reader)) {
+ List<InternalRow> actualRows = new ArrayList<>();
+ while (iterator.hasNext()) {
+ actualRows.add(serializer.copy(iterator.next()));
+ }
+
+ assertEquals(3, actualRows.size());
+ assertEquals(1, actualRows.get(0).getInt(0));
+ assertEquals(2, actualRows.get(1).getInt(0));
+ assertEquals(3, actualRows.get(2).getInt(0));
+ }
+ }
+
+ @Test
+ public void testReorderedArrowBundleFallsBackToRows(@TempDir
java.nio.file.Path tempDir)
+ throws Exception {
+ RowType writerType =
+ RowType.builder().field("a", DataTypes.INT()).field("b",
DataTypes.INT()).build();
+ RowType sourceType =
+ RowType.builder().field("b", DataTypes.INT()).field("a",
DataTypes.INT()).build();
+
+ Options options = new Options();
+ VortexFileFormat format =
+ new VortexFileFormatFactory()
+ .create(new FileFormatFactory.FormatContext(options,
1024, 1024));
+ FileIO fileIO = new LocalFileIO();
+ Path testFile =
+ new Path(new Path(tempDir.toUri()), "test_reordered_bundle_" +
UUID.randomUUID());
+
+ try (FormatWriter writer =
+ ((SupportsDirectWrite)
format.createWriterFactory(writerType))
+ .create(fileIO, testFile, "");
+ RootAllocator sourceAllocator = new RootAllocator();
+ VectorSchemaRoot root =
+ ArrowUtils.createVectorSchemaRoot(sourceType,
sourceAllocator)) {
+ setInt((IntVector) root.getVector("b"), 20);
+ setInt((IntVector) root.getVector("a"), 10);
+ root.setRowCount(1);
+
+ ((BundleFormatWriter) writer)
+ .writeBundle(new ArrowBundleRecords(root, writerType,
true));
+ }
+
+ InternalRowSerializer serializer = new
InternalRowSerializer(writerType);
+ FormatReaderFactory readerFactory =
+ format.createReaderFactory(writerType, writerType, null);
+ try (RecordReader<InternalRow> reader =
+ readerFactory.createReader(
+ new FormatReaderContext(
+ fileIO, testFile,
fileIO.getFileSize(testFile), null));
+ RecordReaderIterator<InternalRow> iterator = new
RecordReaderIterator<>(reader)) {
+ List<InternalRow> actualRows = new ArrayList<>();
+ while (iterator.hasNext()) {
+ actualRows.add(serializer.copy(iterator.next()));
+ }
+
+ assertEquals(1, actualRows.size());
+ assertEquals(10, actualRows.get(0).getInt(0));
+ assertEquals(20, actualRows.get(0).getInt(1));
}
}
@@ -367,6 +484,12 @@ public class VortexReaderWriterTest {
}
}
+ private static void setInt(IntVector vector, int value) {
+ vector.allocateNew(1);
+ vector.setSafe(0, value);
+ vector.setValueCount(1);
+ }
+
@Test
public void testReadWithVirtualRowTrackingField(@TempDir
java.nio.file.Path tempDir)
throws Exception {
