This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 34db147b8cefb1dc562edf42eccb1df443dbedbf 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(
