[core] Enhance ArrowBundleWriter to more batch scenarios

Author: JingsongLi

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

@@ -20,6 +20,7 @@
 
 import org.apache.paimon.arrow.ArrowUtils;
 import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.columnar.ColumnVector;
 import org.apache.paimon.types.RowType;
 
 import org.apache.arrow.c.ArrowArray;
@@ -70,10 +71,19 @@ public ArrowFormatCWriter(ArrowFormatWriter arrowFormatWriter) {
         schema = ArrowSchema.allocateNew(allocator);
     }
 
+    public ArrowFormatWriter formatWriter() {
+        return realWriter;
+    }
+
     public boolean write(InternalRow currentRow) {
         return realWriter.write(currentRow);
     }
 
+    public void write(
+            ColumnVector[] columns, @Nullable int[] pickedInColumn, int startIndex, int batchRows) {
+        realWriter.write(columns, pickedInColumn, startIndex, batchRows);
+    }
+
     public ArrowCStruct toCStruct() {
         VectorSchemaRoot vectorSchemaRoot = realWriter.getVectorSchemaRoot();
         return ArrowUtils.serializeToCStruct(

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

@@ -24,10 +24,12 @@
 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.DataField;
 import org.apache.paimon.types.DataType;
 import org.apache.paimon.types.RowType;
 import org.apache.paimon.types.VariantType;
+import org.apache.paimon.utils.Preconditions;
 
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
@@ -187,6 +189,15 @@ public boolean write(InternalRow currentRow) {
         return true;
     }
 
+    public void write(
+            ColumnVector[] columns, @Nullable int[] pickedInColumn, int startIndex, int batchRows) {
+        Preconditions.checkState(rowId == 0, "rowId must be 0 before writing columns.");
+        for (int i = 0; i < columns.length; i++) {
+            fieldWriters[i].write(columns[i], pickedInColumn, startIndex, batchRows);
+        }
+        rowId = batchRows;
+    }
+
     public long memoryUsed() {
         vectorSchemaRoot.setRowCount(rowId);
         long memoryUsed = 0;
@@ -213,6 +224,14 @@ public void close() {
         allocator.close();
     }
 
+    public int getBatchSize() {
+        return batchSize;
+    }
+
+    public ArrowFieldWriter[] getFieldWriters() {
+        return fieldWriters;
+    }
+
     public VectorSchemaRoot getVectorSchemaRoot() {
         return vectorSchemaRoot;
     }

paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java

@@ -23,9 +23,18 @@
 import org.apache.paimon.arrow.vector.ArrowCStruct;
 import org.apache.paimon.arrow.vector.ArrowFormatCWriter;
 import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.columnar.ColumnVector;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
+import org.apache.paimon.deletionvectors.DeletionFileRecordIterator;
+import org.apache.paimon.deletionvectors.DeletionVectorJudger;
 import org.apache.paimon.format.BundleFormatWriter;
 import org.apache.paimon.fs.PositionOutputStream;
 import org.apache.paimon.io.BundleRecords;
+import org.apache.paimon.io.DeletionFileIteratorRecords;
+import org.apache.paimon.io.VectorizedIteratorRecords;
+import org.apache.paimon.reader.FileRecordIterator;
+import org.apache.paimon.reader.VectorizedRecordIterator;
+import org.apache.paimon.utils.IntArrayList;
 
 import org.apache.arrow.c.ArrowArray;
 import org.apache.arrow.c.ArrowSchema;
@@ -34,6 +43,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 
 /** Arrow bundle writer. */
@@ -72,11 +83,23 @@ public void addElement(InternalRow internalRow) {
     public void writeBundle(BundleRecords bundleRecords) throws IOException {
         if (bundleRecords instanceof ArrowBundleRecords) {
             add(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot());
-        } else {
-            for (InternalRow row : bundleRecords) {
-                addElement(row);
+            return;
+        } else if (bundleRecords instanceof VectorizedIteratorRecords) {
+            add(((VectorizedIteratorRecords) bundleRecords).vectorizedIterator().batch(), null);
+            return;
+        } else if (bundleRecords instanceof DeletionFileIteratorRecords) {
+            DeletionFileRecordIterator iterator =
+                    ((DeletionFileIteratorRecords) bundleRecords).deletionFileIterator();
+            FileRecordIterator<InternalRow> recordIterator = iterator.iterator();
+            if (recordIterator instanceof VectorizedRecordIterator) {
+                add(((VectorizedRecordIterator) recordIterator).batch(), iterator.deletionVector());
+                return;
             }
         }
+
+        for (InternalRow row : bundleRecords) {
+            addElement(row);
+        }
     }
 
     public void add(VectorSchemaRoot vsr) {
@@ -98,6 +121,47 @@ public void add(VectorSchemaRoot vsr) {
         }
     }
 
+    public void add(
+            VectorizedColumnBatch batch, @Nullable DeletionVectorJudger deletionVectorJudger) {
+        if (!arrowFormatWriter.empty()) {
+            flush();
+        }
+
+        int batchSize = arrowFormatWriter.formatWriter().getBatchSize();
+        ColumnVector[] columns = batch.columns;
+
+        int[] pickedInColumn = null;
+        int totalNumRows;
+
+        if (deletionVectorJudger != null) {
+            int originNumRows = batch.getNumRows();
+            IntArrayList picked = new IntArrayList(originNumRows);
+            for (int i = 0; i < originNumRows; i++) {
+                if (!deletionVectorJudger.isDeleted(i)) {
+                    picked.add(i);
+                }
+            }
+            if (picked.size() == originNumRows) {
+                totalNumRows = originNumRows;
+            } else {
+                pickedInColumn = picked.toArray();
+                totalNumRows = pickedInColumn.length;
+            }
+        } else {
+            totalNumRows = batch.getNumRows();
+        }
+
+        int startIndex = 0;
+        while (startIndex < totalNumRows) {
+            int batchRows = Math.min(batchSize, totalNumRows - startIndex);
+            arrowFormatWriter.write(columns, pickedInColumn, startIndex, batchRows);
+            startIndex += batchRows;
+            if (startIndex < totalNumRows) {
+                flush();
+            }
+        }
+    }
+
     @Override
     public boolean reachTargetSize(boolean suggestedCheck, long targetSize) throws IOException {
         return suggestedCheck && (underlyingStream.getPos() > targetSize);

paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowFieldWriters.java

@@ -105,7 +105,7 @@ protected void doWrite(
             for (int i = 0; i < batchRows; i++) {
                 int row = getRowNumber(startIndex, i, pickedInColumn);
                 if (columnVector.isNullAt(row)) {
-                    varCharVector.setNull(row);
+                    varCharVector.setNull(i);
                 } else {
                     byte[] value = ((BytesColumnVector) columnVector).getBytes(row).getBytes();
                     varCharVector.setSafe(i, value);
@@ -154,7 +154,7 @@ protected void doWrite(
             for (int i = 0; i < batchRows; i++) {
                 int row = getRowNumber(startIndex, i, pickedInColumn);
                 if (columnVector.isNullAt(row)) {
-                    bitVector.setNull(row);
+                    bitVector.setNull(i);
                 } else {
                     int value = ((BooleanColumnVector) columnVector).getBoolean(row) ? 1 : 0;
                     bitVector.setSafe(i, value);
@@ -473,7 +473,7 @@ protected void doWrite(
                 if (columnVector.isNullAt(row)) {
                     timeMilliVector.setNull(i);
                 } else {
-                    int value = ((IntColumnVector) columnVector).getInt(i);
+                    int value = ((IntColumnVector) columnVector).getInt(row);
                     timeMilliVector.setSafe(i, value);
                 }
             }