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 25f6603ec22e..b179beafb0cd 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 @@ 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 VectorSchemaRoot getVectorSchemaRoot() { 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 arrowFields = vectorSchemaRoot.getSchema().getFields(); + List dataFields = rowType.getFields(); + if (arrowFields.size() != dataFields.size()) { + return false; + } + + Set 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 404e6c009e27..5ea4ac8aaac9 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 static ArrowCStruct serializeToCStruct( 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 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 @@ private static long zoneCastedTimestampZoneCastToEpoch( return instant.getEpochSecond() * 1_000_000_000 + instant.getNano(); } } + + private static boolean allVectorsShareRootWith( + List 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 6133bfb55095..f120b141cd84 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.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 @@ private ArrowFormatWriter( boolean closeAllocatorOnClose) { this.allocator = allocator; this.closeAllocatorOnClose = closeAllocatorOnClose; + this.rowType = rowType; RowType outputRowType = replaceWithShreddingType(rowType, shreddingSchemas); vectorSchemaRoot = @@ -303,6 +310,71 @@ public BufferAllocator getAllocator() { 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 leftFields = ((RowType) left).getFields(); + List 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 39a006b6d7ed..f8c7a22bb74c 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 void addElement(InternalRow internalRow) { @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 void add(VectorSchemaRoot vsr) { 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 ee613a05d706..75d70c849b9c 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.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 void testVectorType() { .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 40aa87dae970..fd257450aaa9 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 void testWriterWithBorrowedAllocatorDoesNotCloseAllocator() { } } + @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 a43ec4e0b65e..e7a82d14484c 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.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 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 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 void writeIpcBytes(long arrayAddress, long schemaAddress) { 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 colObjects = new ArrayList<>(); @@ -618,6 +901,12 @@ static class Snapshot { } } + 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 e0ba09b61243..95be4e9339b9 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 java.io.IOException; -/** Format write with bundle interface. */ +/** + * Format writer with a row-equivalent bundle interface. + * + *

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 22f102d814b5..37b6674e0107 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.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 bufferedRows; - private final List bufferedBundles; @Nullable private FormatWriter actualWriter; private boolean planFinalized = false; @@ -59,7 +56,6 @@ public InferShreddingWritePlanWriter( this.out = out; this.compression = compression; this.bufferedRows = new ArrayList<>(); - this.bufferedBundles = new ArrayList<>(); } @Override @@ -80,14 +76,8 @@ public void addElement(InternalRow row) throws IOException { @Override public void writeBundle(BundleRecords bundle) throws IOException { if (!planFinalized) { - final List 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 void close() throws IOException { } 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 collectAllRows() { - if (bufferedBundles.isEmpty()) { - return bufferedRows; - } - - List 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 rows; - - private CopiedBundleRecords(List rows) { - this.rows = rows; - } - - @Override - @Nonnull - public Iterator 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 fad92112b11d..89758b6157e1 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 @@ public interface BundleRecords extends Iterable { /** - * The total row count of this batch. + * The stable, non-negative row count of this batch. * - * @return the number of row count. + *

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 000000000000..be7a8b197ea0 --- /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 records = Arrays.asList(rows); + return new BundleRecords() { + @Override + public Iterator 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 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 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 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 0ebdc5feb790..fe35d51fdd5f 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 void writeBundle(BundleRecords bundle) throws IOException { 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 0372c7f885b4..73104ab946d6 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 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 void update(InternalRow row) { 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 361a90facf9a..8d5fb5ede353 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 void write(InternalRow row) throws IOException { @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 8ae8f4bdc075..96f273091dd7 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.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 void writeBundle(BundleRecords bundle) throws IOException { } 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 void writeBundle(BundleRecords bundle) throws IOException { 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 773ca08a815f..7f660a37f453 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 void writeBundle(BundleRecords bundle) throws IOException { 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 ac1c9442e2a9..5eb6fd8a1d31 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 void testRollingByRowsWithBundle() throws IOException { assertThat(files.get(2).rowCount()).isEqualTo(30); } - private static BundleRecords bundle(int rowCount) { + private static SingleUseBundleRecords bundle(int rowCount) { List 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 000000000000..192ba65e6977 --- /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 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 rows; + private int iteratorCalls; + + private TrackingBundleRecords(List rows) { + this.rows = rows; + } + + @Override + public Iterator iterator() { + iteratorCalls++; + return rows.iterator(); + } + + @Override + public long rowCount() { + return rows.size(); + } + } + + private static class ListBundleRecords implements BundleRecords { + + private final List rows; + + private ListBundleRecords(List rows) { + this.rows = rows; + } + + @Override + public Iterator 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 ee1f572462d7..a23d95489f9e 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 void addElement(InternalRow internalRow) throws IOException { @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 f9329907f951..b3e29eb0a0b3 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 storageOptions; private LanceFileWriter writer; + private BufferAllocator allocator; private long bytesWritten = 0; public LanceWriter(String path, Map storageOptions) { @@ -44,12 +47,25 @@ public Long getWrittenPosition() { } 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 void close() throws IOException { throw new IOException(e); } this.writer = null; + this.allocator = null; } } @@ -68,6 +85,7 @@ public String path() { 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 000000000000..7d5161d5dca2 --- /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 actual = new ArrayList<>(); + try (RecordReaderIterator 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 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 fieldNames = + root.getSchema().getFields().stream() + .map(field -> field.getName()) + .collect(Collectors.toList()); + List> values = new ArrayList<>(); + for (int column = 0; column < root.getFieldVectors().size(); column++) { + IntVector vector = (IntVector) root.getVector(column); + List 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 fieldNames; + private final List> values; + + private Snapshot(List fieldNames, List> 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 944f78c4b147..1a3e4d4149b6 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 void addElement(InternalRow internalRow) { @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 60efceed08e3..931f4d703318 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.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 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 @@ void testWriteAndRead() throws IOException { 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 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 @@ private static boolean isNativeAvailable() { return false; } } + + private static class CountingArrowBundleRecords extends ArrowBundleRecords { + + private int iteratorCalls; + + private CountingArrowBundleRecords(VectorSchemaRoot root, RowType rowType) { + super(root, rowType, true); + } + + @Override + public Iterator 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 58a8dc5252a2..344d5a0022b6 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 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 @@ void testConstructorFailureClosesCreatedResources() { new MosaicRecordsWriter( new ByteArrayOutputStream(), rowType, - formatContext, + FORMAT_CONTEXT, Collections.emptyList(), null, allocator, @@ -60,6 +73,108 @@ void testConstructorFailureClosesCreatedResources() { 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 7798fa09d419..502b4185e870 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.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 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 void addElement(InternalRow internalRow) throws IOException { @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 void close() throws IOException { // 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 @@ private void flush() throws IOException { } } - /** 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. - // 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 e97e3cf80fcc..81dab2774d19 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.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 void testWriteDoesNotLeakArrowMemoryOnClose(@TempDir java.nio.file.Path t } @Test - public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( + public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffersAcrossBatches( @TempDir java.nio.file.Path tempDir) throws Exception { RowType rowType = RowType.builder() @@ -289,6 +293,16 @@ public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( 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 void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( 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 reader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, testFile, fileIO.getFileSize(testFile), null)); + RecordReaderIterator iterator = new RecordReaderIterator<>(reader)) { + List 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 reader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, testFile, fileIO.getFileSize(testFile), null)); + RecordReaderIterator iterator = new RecordReaderIterator<>(reader)) { + List 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 void testReadWithSelection(@TempDir java.nio.file.Path tempDir) throws Ex } } + 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 {