This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 6d3e2e6898 [mosaic] Size the writer's Arrow vectors by
write.batch-size (#9741)
6d3e2e6898 is described below
commit 6d3e2e6898dda197dd356b65c99324119c3a4c9a
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Fri Sep 11 17:47:25 2026 +0800
[mosaic] Size the writer's Arrow vectors by write.batch-size (#9741)
---
.../paimon/format/mosaic/MosaicRecordsWriter.java | 19 +++
.../format/mosaic/MosaicRecordsWriterTest.java | 136 +++++++++++++++++++++
2 files changed, 155 insertions(+)
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 1a3e4d4149..f4a96937ce 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
@@ -33,7 +33,10 @@ import org.apache.paimon.types.RowType;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BaseValueVector;
+import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.BaseRepeatedValueVector;
import org.apache.arrow.vector.types.pojo.Schema;
import javax.annotation.Nullable;
@@ -101,6 +104,13 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
createdArrowWriter =
ArrowFormatWriter.forBorrowedAllocator(
rowType, writeBatchSize, true, allocator,
writeBatchMemory);
+ // Only batches smaller than Arrow's default allocation are sized
by the batch.
+ if (writeBatchSize < BaseValueVector.INITIAL_VALUE_ALLOCATION) {
+ for (FieldVector vector :
+
createdArrowWriter.getVectorSchemaRoot().getFieldVectors()) {
+ setInitialCapacity(vector, writeBatchSize);
+ }
+ }
Schema arrowSchema =
createdArrowWriter.getVectorSchemaRoot().getSchema();
createdNativeWriter =
nativeWriterFactory.create(outputStream, arrowSchema,
options, allocator);
@@ -113,6 +123,15 @@ public class MosaicRecordsWriter implements
BundleFormatWriter {
this.nativeWriter = createdNativeWriter;
}
+ private static void setInitialCapacity(FieldVector vector, int capacity) {
+ if (vector instanceof BaseRepeatedValueVector) {
+ // The plain overload would size the element vector for 5 elements
per row.
+ ((BaseRepeatedValueVector) vector).setInitialCapacity(capacity,
1.0);
+ } else {
+ vector.setInitialCapacity(capacity);
+ }
+ }
+
@Override
public void addElement(InternalRow internalRow) {
if (!arrowFormatWriter.write(internalRow)) {
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 344d5a0022..be69fb6207 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
@@ -20,8 +20,13 @@ 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.BinaryString;
+import org.apache.paimon.data.GenericArray;
+import org.apache.paimon.data.GenericRow;
import org.apache.paimon.format.FileFormatFactory;
import org.apache.paimon.mosaic.MosaicWriter;
+import org.apache.paimon.options.MemorySize;
import org.apache.paimon.options.Options;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
@@ -157,6 +162,137 @@ class MosaicRecordsWriterTest {
verify(nativeWriter).write(any(VectorSchemaRoot.class));
}
+ @Test
+ void testVectorsAreSizedByWriteBatchSize() throws Exception {
+ // 2,000 columns: Arrow's default per-vector allocation would exceed
60 MB here.
+ RowType.Builder builder = RowType.builder();
+ for (int i = 0; i < 2000; i++) {
+ builder.field("c" + i, DataTypes.DOUBLE());
+ }
+ RowType wideType = builder.build();
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ try (RootAllocator allocator = new RootAllocator()) {
+ MosaicRecordsWriter writer =
+ new MosaicRecordsWriter(
+ new ByteArrayOutputStream(),
+ wideType,
+ new FileFormatFactory.FormatContext(new Options(),
1024, 4),
+ Collections.emptyList(),
+ null,
+ allocator,
+ (outputStream, arrowSchema, options,
bufferAllocator) -> nativeWriter);
+ GenericRow row = new GenericRow(wideType.getFieldCount());
+ row.setField(0, 1.0d);
+ writer.addElement(row);
+ assertThat(allocator.getAllocatedMemory()).isLessThan(8L * 1024 *
1024);
+ writer.close();
+ }
+ }
+
+ @Test
+ void testLargeBatchSizeWithSmallMemoryBudgetKeepsArrowDefaultAllocation()
throws Exception {
+ // write.batch-size=65536 with write.batch-memory=128 KiB: the memory
budget flushes long
+ // before the row limit, so the first row must not allocate for the
whole batch size.
+ RowType rowType = mixedRowType(100);
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ try (RootAllocator allocator = new RootAllocator(16L * 1024 * 1024)) {
+ MosaicRecordsWriter writer =
+ createWriter(
+ rowType, allocator, nativeWriter, 65536,
MemorySize.ofKibiBytes(128));
+ writer.addElement(firstRow(rowType));
+ assertThat(allocator.getAllocatedMemory())
+ .isEqualTo(baselineFirstRowAllocation(rowType, 65536));
+ writer.close();
+ }
+ }
+
+ @Test
+ void testInitialCapacityNeverExceedsArrowDefaultAllocation() throws
Exception {
+ // DOUBLE, STRING and ARRAY<DOUBLE> columns cover the fixed-width,
variable-width and
+ // repeated vector families, whose default sizing differs.
+ RowType rowType = mixedRowType(60);
+ for (int batchSize : new int[] {4, 1024, 3969, 3970, 65536}) {
+ long baseline = baselineFirstRowAllocation(rowType, batchSize);
+ MosaicWriter nativeWriter = mock(MosaicWriter.class);
+ try (RootAllocator allocator = new RootAllocator()) {
+ MosaicRecordsWriter writer =
+ createWriter(
+ rowType,
+ allocator,
+ nativeWriter,
+ batchSize,
+ MemorySize.VALUE_128_MB);
+ try {
+ writer.addElement(firstRow(rowType));
+ assertThat(allocator.getAllocatedMemory())
+ .as("batch size %d", batchSize)
+ .isLessThanOrEqualTo(baseline);
+ // Arrow rounds buffers to powers of two, so only clearly
smaller batches
+ // allocate less.
+ if (batchSize <= 1024) {
+ assertThat(allocator.getAllocatedMemory())
+ .as("batch size %d", batchSize)
+ .isLessThan(baseline);
+ }
+ } finally {
+ writer.close();
+ }
+ }
+ }
+ }
+
+ private static RowType mixedRowType(int columnsPerType) {
+ RowType.Builder builder = RowType.builder();
+ for (int i = 0; i < columnsPerType; i++) {
+ builder.field("d" + i, DataTypes.DOUBLE());
+ builder.field("s" + i, DataTypes.STRING());
+ builder.field("a" + i, DataTypes.ARRAY(DataTypes.DOUBLE()));
+ }
+ return builder.build();
+ }
+
+ private static GenericRow firstRow(RowType rowType) {
+ GenericRow row = new GenericRow(rowType.getFieldCount());
+ row.setField(0, 1.0d);
+ row.setField(1, BinaryString.fromString("one"));
+ row.setField(2, new GenericArray(new Object[] {1.0d}));
+ return row;
+ }
+
+ /** First-row allocation of the same writer without the capacity loop. */
+ private static long baselineFirstRowAllocation(RowType rowType, int
batchSize) {
+ try (RootAllocator allocator = new RootAllocator()) {
+ ArrowFormatWriter writer =
+ ArrowFormatWriter.forBorrowedAllocator(
+ rowType,
+ batchSize,
+ true,
+ allocator,
+ MemorySize.VALUE_128_MB.getBytes());
+ assertThat(writer.write(firstRow(rowType))).isTrue();
+ long allocated = allocator.getAllocatedMemory();
+ writer.close();
+ return allocated;
+ }
+ }
+
+ private static MosaicRecordsWriter createWriter(
+ RowType rowType,
+ RootAllocator allocator,
+ MosaicWriter nativeWriter,
+ int writeBatchSize,
+ MemorySize writeBatchMemory) {
+ return new MosaicRecordsWriter(
+ new ByteArrayOutputStream(),
+ rowType,
+ new FileFormatFactory.FormatContext(
+ new Options(), 1024, writeBatchSize, writeBatchMemory),
+ Collections.emptyList(),
+ null,
+ allocator,
+ (outputStream, arrowSchema, options, bufferAllocator) ->
nativeWriter);
+ }
+
private static MosaicRecordsWriter createWriter(
RowType rowType, RootAllocator allocator, MosaicWriter
nativeWriter) {
return new MosaicRecordsWriter(
