[core][format] Support Arrow bundle writes

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<Field> arrowFields = vectorSchemaRoot.getSchema().getFields();
+        List<DataField> dataFields = rowType.getFields();
+        if (arrowFields.size() != dataFields.size()) {
+            return false;
+        }
+
+        Set<String> mappedNames = new HashSet<>();
+        for (int i = 0; i < arrowFields.size(); i++) {
+            String arrowName = toLowerCaseIfNeed(arrowFields.get(i).getName(), caseSensitive);
+            String dataName = toLowerCaseIfNeed(dataFields.get(i).name(), caseSensitive);
+            if (!arrowName.equals(dataName) || !mappedNames.add(arrowName)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     @Override
     public long rowCount() {
         return vectorSchemaRoot.getRowCount();

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<FieldVector> vectors = vectorSchemaRoot.getFieldVectors();
+        return !vectors.isEmpty() && allVectorsShareRootWith(vectors, allocator.getRoot());
+    }
+
     public static void serializeToIpc(VectorSchemaRoot vsr, OutputStream out) {
         try (ArrowStreamWriter writer = new ArrowStreamWriter(vsr, null, out)) {
             writer.writeBatch();
@@ -316,4 +339,15 @@ private static long zoneCastedTimestampZoneCastToEpoch(
             return instant.getEpochSecond() * 1_000_000_000 + instant.getNano();
         }
     }
+
+    private static boolean allVectorsShareRootWith(
+            List<FieldVector> vectors, BufferAllocator expectedRoot) {
+        for (FieldVector vector : vectors) {
+            if (vector.getAllocator().getRoot() != expectedRoot
+                    || !allVectorsShareRootWith(vector.getChildrenFromFields(), expectedRoot)) {
+                return false;
+            }
+        }
+        return true;
+    }
 }

paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java

@@ -18,17 +18,22 @@
 
 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;
 import org.apache.paimon.arrow.writer.ArrowFieldWriterFactoryVisitor;
 import org.apache.paimon.arrow.writer.ArrowFieldWriters;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.data.columnar.ColumnVector;
+import org.apache.paimon.types.ArrayType;
 import org.apache.paimon.types.DataField;
 import org.apache.paimon.types.DataType;
+import org.apache.paimon.types.MapType;
+import org.apache.paimon.types.MultisetType;
 import org.apache.paimon.types.RowType;
 import org.apache.paimon.types.VariantType;
+import org.apache.paimon.types.VectorType;
 import org.apache.paimon.utils.Preconditions;
 
 import org.apache.arrow.memory.BufferAllocator;
@@ -51,6 +56,7 @@ public class ArrowFormatWriter implements AutoCloseable {
 
     private final VectorSchemaRoot vectorSchemaRoot;
     private final ArrowFieldWriter[] fieldWriters;
+    private final RowType rowType;
     private final int batchSize;
     private final BufferAllocator allocator;
     @Nullable private final Long memoryUsedMaxInBytes;
@@ -171,6 +177,7 @@ 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<DataField> leftFields = ((RowType) left).getFields();
+            List<DataField> rightFields = ((RowType) right).getFields();
+            if (leftFields.size() != rightFields.size()) {
+                return false;
+            }
+            for (int i = 0; i < leftFields.size(); i++) {
+                DataField leftField = leftFields.get(i);
+                DataField rightField = rightFields.get(i);
+                if (!leftField.name().equals(rightField.name())
+                        || !hasSameLogicalLayout(leftField.type(), rightField.type())) {
+                    return false;
+                }
+            }
+            return true;
+        }
+
+        if (left instanceof ArrayType) {
+            return hasSameLogicalLayout(
+                    ((ArrayType) left).getElementType(), ((ArrayType) right).getElementType());
+        }
+
+        if (left instanceof MapType) {
+            MapType leftMap = (MapType) left;
+            MapType rightMap = (MapType) right;
+            return hasSameLogicalLayout(leftMap.getKeyType(), rightMap.getKeyType())
+                    && hasSameLogicalLayout(leftMap.getValueType(), rightMap.getValueType());
+        }
+
+        if (left instanceof MultisetType) {
+            return hasSameLogicalLayout(
+                    ((MultisetType) left).getElementType(),
+                    ((MultisetType) right).getElementType());
+        }
+
+        if (left instanceof VectorType) {
+            VectorType leftVector = (VectorType) left;
+            VectorType rightVector = (VectorType) right;
+            return leftVector.getLength() == rightVector.getLength()
+                    && hasSameLogicalLayout(
+                            leftVector.getElementType(), rightVector.getElementType());
+        }
+
+        return left.equals(right);
+    }
+
     private static RowType replaceWithShreddingType(
             RowType rowType, @Nullable RowType shreddingSchemas) {
         if (shreddingSchemas == null) {

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);

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);
+        }
+    }
 }

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) {