[mosaic] Expose column-batch reads and writer extension hooks

paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java

@@ -41,7 +41,7 @@
 /** Reader from a {@link VectorSchemaRoot} to paimon rows. */
 public class ArrowBatchReader {
 
-    private final VectorizedColumnBatch batch;
+    private final VectorizedColumnBatch reusableBatch;
     private final Arrow2PaimonVectorConverter[] convertors;
     private final RowType projectedRowType;
     private final boolean caseSensitive;
@@ -57,19 +57,53 @@ public ArrowBatchReader(
             RowType rowType,
             boolean caseSensitive,
             Arrow2PaimonVectorConverter.Arrow2PaimonVectorConvertorVisitor visitor) {
-        ColumnVector[] columnVectors = new ColumnVector[rowType.getFieldCount()];
+        this.reusableBatch = new VectorizedColumnBatch(new ColumnVector[rowType.getFieldCount()]);
         this.convertors = new Arrow2PaimonVectorConverter[rowType.getFieldCount()];
-        this.batch = new VectorizedColumnBatch(columnVectors);
         this.projectedRowType = rowType;
 
-        for (int i = 0; i < columnVectors.length; i++) {
+        for (int i = 0; i < convertors.length; i++) {
             this.convertors[i] =
                     Arrow2PaimonVectorConverter.construct(visitor, rowType.getTypeAt(i));
         }
         this.caseSensitive = caseSensitive;
     }
 
     public Iterable<InternalRow> readBatch(VectorSchemaRoot vsr) {
+        populateBatch(vsr, reusableBatch);
+        int rowCount = reusableBatch.getNumRows();
+        final ColumnarRow columnarRow = new ColumnarRow(reusableBatch);
+        return () ->
+                new Iterator<InternalRow>() {
+                    private int position = 0;
+
+                    @Override
+                    public boolean hasNext() {
+                        return position < rowCount;
+                    }
+
+                    @Override
+                    public InternalRow next() {
+                        columnarRow.setRowId(position);
+                        position++;
+                        return columnarRow;
+                    }
+                };
+    }
+
+    /**
+     * Wraps an Arrow batch as Paimon column vectors without materializing rows.
+     *
+     * <p>The returned batch container is not reused, but its columns borrow vectors owned by {@code
+     * vsr} and must not be used after the root is released.
+     */
+    public VectorizedColumnBatch readVectorizedBatch(VectorSchemaRoot vsr) {
+        VectorizedColumnBatch resultBatch =
+                new VectorizedColumnBatch(new ColumnVector[projectedRowType.getFieldCount()]);
+        populateBatch(vsr, resultBatch);
+        return resultBatch;
+    }
+
+    private void populateBatch(VectorSchemaRoot vsr, VectorizedColumnBatch targetBatch) {
         int[] mapping = new int[projectedRowType.getFieldCount()];
         Schema arrowSchema = vsr.getSchema();
         Map<String, Integer> arrowFieldIndex = new HashMap<>();
@@ -83,32 +117,14 @@ public Iterable<InternalRow> readBatch(VectorSchemaRoot vsr) {
             mapping[i] = arrowFieldIndex.getOrDefault(fieldName, -1);
         }
 
-        for (int i = 0; i < batch.columns.length; i++) {
+        for (int i = 0; i < targetBatch.columns.length; i++) {
             if (mapping[i] >= 0) {
-                batch.columns[i] = convertors[i].convertVector(vsr.getVector(mapping[i]));
+                targetBatch.columns[i] = convertors[i].convertVector(vsr.getVector(mapping[i]));
             } else {
-                batch.columns[i] = AllNullColumnVector.INSTANCE;
+                targetBatch.columns[i] = AllNullColumnVector.INSTANCE;
             }
         }
 
-        int rowCount = vsr.getRowCount();
-        batch.setNumRows(vsr.getRowCount());
-        final ColumnarRow columnarRow = new ColumnarRow(batch);
-        return () ->
-                new Iterator<InternalRow>() {
-                    private int position = 0;
-
-                    @Override
-                    public boolean hasNext() {
-                        return position < rowCount;
-                    }
-
-                    @Override
-                    public InternalRow next() {
-                        columnarRow.setRowId(position);
-                        position++;
-                        return columnarRow;
-                    }
-                };
+        targetBatch.setNumRows(vsr.getRowCount());
     }
 }

paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java

@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.arrow.reader;
+
+import org.apache.paimon.arrow.ArrowUtils;
+import org.apache.paimon.data.columnar.ColumnarRow;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link ArrowBatchReader}. */
+class ArrowBatchReaderTest {
+
+    @Test
+    void testReadBatchWrapperIsReused() {
+        RowType rowType = RowType.builder().field("id", DataTypes.INT()).build();
+        try (RootAllocator allocator = new RootAllocator();
+                VectorSchemaRoot firstRoot = intRoot(rowType, allocator, 11);
+                VectorSchemaRoot secondRoot = intRoot(rowType, allocator, 22)) {
+            ArrowBatchReader reader = new ArrowBatchReader(rowType, true);
+
+            ColumnarRow first = (ColumnarRow) reader.readBatch(firstRoot).iterator().next();
+            VectorizedColumnBatch reusableBatch = first.batch();
+            assertThat(first.getInt(0)).isEqualTo(11);
+
+            ColumnarRow second = (ColumnarRow) reader.readBatch(secondRoot).iterator().next();
+
+            assertThat(second.batch()).isSameAs(reusableBatch);
+            assertThat(second.getInt(0)).isEqualTo(22);
+        }
+    }
+
+    @Test
+    void testVectorizedBatchWrappersAreNotReused() {
+        RowType rowType = RowType.builder().field("id", DataTypes.INT()).build();
+        try (RootAllocator allocator = new RootAllocator();
+                VectorSchemaRoot firstRoot = intRoot(rowType, allocator, 11);
+                VectorSchemaRoot secondRoot = intRoot(rowType, allocator, 22, 33)) {
+            ArrowBatchReader reader = new ArrowBatchReader(rowType, true);
+
+            VectorizedColumnBatch first = reader.readVectorizedBatch(firstRoot);
+            VectorizedColumnBatch second = reader.readVectorizedBatch(secondRoot);
+
+            assertThat(first).isNotSameAs(second);
+            assertThat(first.columns).isNotSameAs(second.columns);
+            assertThat(first.getNumRows()).isEqualTo(1);
+            assertThat(second.getNumRows()).isEqualTo(2);
+            assertThat(first.getInt(0, 0)).isEqualTo(11);
+            assertThat(second.getInt(0, 0)).isEqualTo(22);
+            assertThat(second.getInt(1, 0)).isEqualTo(33);
+        }
+    }
+
+    private static VectorSchemaRoot intRoot(
+            RowType rowType, RootAllocator allocator, int... values) {
+        VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(rowType, allocator);
+        IntVector vector = (IntVector) root.getVector(0);
+        vector.allocateNew(values.length);
+        for (int i = 0; i < values.length; i++) {
+            vector.setSafe(i, values[i]);
+        }
+        vector.setValueCount(values.length);
+        root.setRowCount(values.length);
+        return root;
+    }
+}