From 0a43c9f87c6db4b220ab0342a69377c15dc10d22 Mon Sep 17 00:00:00 2001 From: mingfeng Date: Fri, 7 Aug 2026 05:16:43 -0700 Subject: [PATCH 1/6] [core] Add direct write path for BundleRecords --- .../paimon/arrow/ArrowBundleRecords.java | 4 + .../paimon/arrow/reader/ArrowBatchReader.java | 68 +-- .../reader/ArrowVectorizedRecordIterator.java | 34 ++ .../arrow/reader/ArrowBatchReaderTest.java | 89 ++++ .../org/apache/paimon/io/BundleRecords.java | 10 + .../io/RowDataFileSequenceNumberTracker.java | 15 + .../apache/paimon/io/RowDataFileWriter.java | 9 + .../io/StatsCollectingSingleFileWriter.java | 4 + .../paimon/io/RowDataFileWriterTest.java | 401 ++++++++++++++++++ 9 files changed, 608 insertions(+), 26 deletions(-) create mode 100644 paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java create mode 100644 paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java index 25f6603ec22e..828649010c5d 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java @@ -45,6 +45,10 @@ public VectorSchemaRoot getVectorSchemaRoot() { return vectorSchemaRoot; } + public RowType getRowType() { + return rowType; + } + @Override public long rowCount() { return vectorSchemaRoot.getRowCount(); diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java index ab939d70b4e7..5548dfb2fc65 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java +++ b/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,12 +57,11 @@ 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)); } @@ -70,6 +69,41 @@ public ArrowBatchReader( } public Iterable readBatch(VectorSchemaRoot vsr) { + populateBatch(vsr, reusableBatch); + int rowCount = reusableBatch.getNumRows(); + final ColumnarRow columnarRow = new ColumnarRow(reusableBatch); + return () -> + new Iterator() { + 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. + * + *

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 arrowFieldIndex = new HashMap<>(); @@ -83,32 +117,14 @@ public Iterable 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() { - 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()); } } diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java new file mode 100644 index 000000000000..de2f38e51a15 --- /dev/null +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java @@ -0,0 +1,34 @@ +/* + * 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.ArrowBundleRecords; +import org.apache.paimon.reader.VectorizedRecordIterator; + +/** A {@link VectorizedRecordIterator} which can expose its Arrow batch for direct bundle writes. */ +public interface ArrowVectorizedRecordIterator extends VectorizedRecordIterator { + + /** + * Returns a borrowed view of the Arrow vectors backing {@link #batch()}. + * + *

The caller does not own the batch and must not retain or close it. Its row order and count + * correspond to {@link #batch()}, and it is valid only until {@link #releaseBatch()}. + */ + ArrowBundleRecords arrowBundle(); +} diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java new file mode 100644 index 000000000000..c70f21ad446e --- /dev/null +++ b/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; + } +} diff --git a/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java b/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java index fad92112b11d..faeff1a5d776 100644 --- a/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java +++ b/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java @@ -31,6 +31,16 @@ @Public public interface BundleRecords extends Iterable { + /** + * Whether this bundle can be passed directly to a matching format writer. + * + *

The writer must consume the bundle synchronously. The producer retains ownership of any + * borrowed native buffers and may release them as soon as {@code writeBundle} returns. + */ + default boolean isDirectWriteBundle() { + return false; + } + /** * The total row count of this batch. * diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java index 0372c7f885b4..73104ab946d6 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileSequenceNumberTracker.java @@ -26,6 +26,9 @@ import java.util.function.Supplier; +import static org.apache.paimon.utils.Preconditions.checkArgument; +import static org.apache.paimon.utils.Preconditions.checkState; + /** * Tracks sequence number range for rows written to a data file. * @@ -99,4 +102,16 @@ public void update(InternalRow row) { hasNullSeqNumber = true; } } + + boolean supportsRowCountUpdate() { + return seqNumberFieldIndex == -1; + } + + void updateByRowCount(long rowCount) { + checkArgument(rowCount >= 0, "Row count must not be negative."); + checkState( + supportsRowCountUpdate(), + "Cannot update sequence numbers by row count when row tracking is enabled."); + seqNumCounter.add(rowCount); + } } diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java index 02dff0fd8233..712c029abc9b 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java @@ -146,6 +146,15 @@ public void write(InternalRow row) throws IOException { @Override public void writeBundle(BundleRecords bundle) throws IOException { + if (bundle.isDirectWriteBundle() + && auxiliaryFileWriters.isEmpty() + && sequenceNumberTracker.supportsRowCountUpdate() + && !requiresPerRecordStats()) { + super.writeBundle(bundle); + sequenceNumberTracker.updateByRowCount(bundle.rowCount()); + return; + } + for (InternalRow row : bundle) { write(row); } diff --git a/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java index 773ca08a815f..7f660a37f453 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java @@ -75,6 +75,10 @@ public void writeBundle(BundleRecords bundle) throws IOException { super.writeBundle(bundle); } + protected final boolean requiresPerRecordStats() { + return statsRequirePerRecord; + } + public SimpleColStats[] fieldStats(long fileSize) throws IOException { Preconditions.checkState(closed, "Cannot access metric unless the writer is closed."); if (isStatsDisabled) { diff --git a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java new file mode 100644 index 000000000000..eab450de59dc --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java @@ -0,0 +1,401 @@ +/* + * 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.io; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.fileindex.FileIndexOptions; +import org.apache.paimon.format.BundleFormatWriter; +import org.apache.paimon.format.FormatWriter; +import org.apache.paimon.format.FormatWriterFactory; +import org.apache.paimon.format.SimpleColStats; +import org.apache.paimon.format.SupportsDirectWrite; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.PositionOutputStream; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.options.Options; +import org.apache.paimon.table.SpecialFields; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.LongCounter; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** Test for {@link RowDataFileWriter}. */ +class RowDataFileWriterTest { + + private static final Path PATH = new Path("file:/tmp/data-file"); + private static final RowType ROW_TYPE = RowType.builder().field("id", DataTypes.INT()).build(); + + @Test + void testBundleFastPathDoesNotIterateRows() throws Exception { + FileIO fileIO = fileIO(); + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + LongCounter sequenceCounter = new LongCounter(5); + RowDataFileWriter writer = + createWriter( + fileIO, + ROW_TYPE, + formatWriter, + SimpleStatsProducer.disabledProducer(), + sequenceCounter, + new FileIndexOptions()); + BundleRecords bundle = new DirectNonIterableBundleRecords(3); + + writer.writeBundle(bundle); + + assertThat(formatWriter.writtenBundle).isSameAs(bundle); + assertThat(formatWriter.bundleWrites).isEqualTo(1); + assertThat(formatWriter.rowWrites).isZero(); + assertThat(writer.recordCount()).isEqualTo(3); + assertThat(sequenceCounter.getValue()).isEqualTo(8); + + writer.close(); + DataFileMeta result = writer.result(); + assertThat(result.rowCount()).isEqualTo(3); + assertThat(result.minSequenceNumber()).isEqualTo(5); + assertThat(result.maxSequenceNumber()).isEqualTo(7); + } + + @Test + void testExtractorStatsAllowBundleFastPath() throws Exception { + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + TestingExtractStatsProducer statsProducer = new TestingExtractStatsProducer(); + LongCounter sequenceCounter = new LongCounter(5); + RowDataFileWriter writer = + createWriter( + fileIO(), + ROW_TYPE, + formatWriter, + statsProducer, + sequenceCounter, + new FileIndexOptions()); + BundleRecords bundle = new DirectNonIterableBundleRecords(3); + + writer.writeBundle(bundle); + + assertThat(formatWriter.writtenBundle).isSameAs(bundle); + assertThat(formatWriter.bundleWrites).isEqualTo(1); + assertThat(formatWriter.rowWrites).isZero(); + assertThat(writer.recordCount()).isEqualTo(3); + assertThat(sequenceCounter.getValue()).isEqualTo(8); + + writer.close(); + DataFileMeta result = writer.result(); + assertThat(statsProducer.extractCalls).isEqualTo(1); + assertThat(result.valueStats().minValues().getInt(0)).isEqualTo(1); + assertThat(result.valueStats().maxValues().getInt(0)).isEqualTo(3); + assertThat(result.valueStats().nullCounts().getLong(0)).isZero(); + } + + @Test + void testUnmarkedBundleFallsBackToRows() throws Exception { + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + RowDataFileWriter writer = + createWriter( + fileIO(), + ROW_TYPE, + formatWriter, + SimpleStatsProducer.disabledProducer(), + new LongCounter(), + new FileIndexOptions()); + + writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); + + assertThat(formatWriter.bundleWrites).isZero(); + assertThat(formatWriter.rowWrites).isEqualTo(2); + assertThat(writer.recordCount()).isEqualTo(2); + } + + @Test + void testPerRecordStatsFallBackToRows() throws Exception { + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + TestingStatsProducer statsProducer = new TestingStatsProducer(); + LongCounter sequenceCounter = new LongCounter(); + RowDataFileWriter writer = + createWriter( + fileIO(), + ROW_TYPE, + formatWriter, + statsProducer, + sequenceCounter, + new FileIndexOptions()); + + writer.writeBundle(directRows(GenericRow.of(1), GenericRow.of(2))); + + assertThat(formatWriter.bundleWrites).isZero(); + assertThat(formatWriter.rowWrites).isEqualTo(2); + assertThat(statsProducer.collectedRows).isEqualTo(2); + assertThat(sequenceCounter.getValue()).isEqualTo(2); + } + + @Test + void testRowTrackingFallsBackToRows() throws Exception { + RowType rowTrackingType = + RowType.builder() + .field("id", DataTypes.INT()) + .field(SpecialFields.SEQUENCE_NUMBER.name(), DataTypes.BIGINT()) + .build(); + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + LongCounter sequenceCounter = new LongCounter(20); + RowDataFileWriter writer = + createWriter( + fileIO(), + rowTrackingType, + formatWriter, + SimpleStatsProducer.disabledProducer(), + sequenceCounter, + new FileIndexOptions()); + + writer.writeBundle(directRows(GenericRow.of(1, 7L), GenericRow.of(2, 11L))); + + assertThat(formatWriter.bundleWrites).isZero(); + assertThat(formatWriter.rowWrites).isEqualTo(2); + assertThat(sequenceCounter.getValue()).isEqualTo(22); + + writer.close(); + DataFileMeta result = writer.result(); + assertThat(result.minSequenceNumber()).isEqualTo(7); + assertThat(result.maxSequenceNumber()).isEqualTo(11); + } + + @Test + void testFileIndexFallsBackToRows() throws Exception { + Options options = new Options(); + options.set("file-index.bitmap.columns", "id"); + FileIndexOptions fileIndexOptions = new FileIndexOptions(new CoreOptions(options)); + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + RowDataFileWriter writer = + createWriter( + fileIO(), + ROW_TYPE, + formatWriter, + SimpleStatsProducer.disabledProducer(), + new LongCounter(), + fileIndexOptions); + + writer.writeBundle(directRows(GenericRow.of(1), GenericRow.of(2))); + + assertThat(formatWriter.bundleWrites).isZero(); + assertThat(formatWriter.rowWrites).isEqualTo(2); + + writer.close(); + assertThat(writer.result().embeddedIndex()).isNotNull(); + } + + private static FileIO fileIO() throws IOException { + FileIO fileIO = mock(FileIO.class); + when(fileIO.getFileSize(PATH)).thenReturn(123L); + return fileIO; + } + + private static RowDataFileWriter createWriter( + FileIO fileIO, + RowType rowType, + TestingBundleFormatWriter formatWriter, + SimpleStatsProducer statsProducer, + LongCounter sequenceCounter, + FileIndexOptions fileIndexOptions) { + return new RowDataFileWriter( + fileIO, + new FileWriterContext( + new TestingFormatWriterFactory(formatWriter), statsProducer, "none"), + PATH, + rowType, + 1L, + () -> sequenceCounter, + fileIndexOptions, + FileSource.APPEND, + false, + false, + false, + null); + } + + private static BundleRecords rows(InternalRow... rows) { + return new ListBundleRecords(Arrays.asList(rows)); + } + + private static BundleRecords directRows(InternalRow... rows) { + return new DirectListBundleRecords(Arrays.asList(rows)); + } + + private static class TestingFormatWriterFactory + implements FormatWriterFactory, SupportsDirectWrite { + + private final TestingBundleFormatWriter writer; + + private TestingFormatWriterFactory(TestingBundleFormatWriter writer) { + this.writer = writer; + } + + @Override + public FormatWriter create(PositionOutputStream out, String compression) { + return writer; + } + + @Override + public FormatWriter create(FileIO fileIO, Path path, String compression) { + return writer; + } + } + + private static class TestingBundleFormatWriter implements BundleFormatWriter { + + private int rowWrites; + private int bundleWrites; + private BundleRecords writtenBundle; + + @Override + public void addElement(InternalRow element) { + rowWrites++; + } + + @Override + public void writeBundle(BundleRecords bundle) { + bundleWrites++; + writtenBundle = bundle; + } + + @Override + public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { + return false; + } + + @Override + public void close() {} + } + + private static class TestingStatsProducer implements SimpleStatsProducer { + + private int collectedRows; + + @Override + public boolean isStatsDisabled() { + return false; + } + + @Override + public boolean requirePerRecord() { + return true; + } + + @Override + public void collect(InternalRow row) { + collectedRows++; + } + + @Override + public SimpleColStats[] extract(FileIO fileIO, Path path, long length) { + return new SimpleColStats[] {SimpleColStats.NONE}; + } + } + + private static class TestingExtractStatsProducer implements SimpleStatsProducer { + + private int extractCalls; + + @Override + public boolean isStatsDisabled() { + return false; + } + + @Override + public boolean requirePerRecord() { + return false; + } + + @Override + public void collect(InternalRow row) { + throw new AssertionError("Extractor-backed statistics must not collect rows."); + } + + @Override + public SimpleColStats[] extract(FileIO fileIO, Path path, long length) { + extractCalls++; + return new SimpleColStats[] {new SimpleColStats(1, 3, 0L)}; + } + } + + private static class DirectNonIterableBundleRecords implements BundleRecords { + + private final long rowCount; + + private DirectNonIterableBundleRecords(long rowCount) { + this.rowCount = rowCount; + } + + @Override + public boolean isDirectWriteBundle() { + return true; + } + + @Override + public Iterator iterator() { + throw new AssertionError("Direct bundle write must not iterate rows."); + } + + @Override + public long rowCount() { + return rowCount; + } + } + + private static class ListBundleRecords implements BundleRecords { + + private final List rows; + + private ListBundleRecords(List rows) { + this.rows = rows; + } + + @Override + public Iterator iterator() { + return rows.iterator(); + } + + @Override + public long rowCount() { + return rows.size(); + } + } + + private static class DirectListBundleRecords extends ListBundleRecords { + + private DirectListBundleRecords(List rows) { + super(rows); + } + + @Override + public boolean isDirectWriteBundle() { + return true; + } + } +} From de0c74ee7235c7c466baa12f6e1ee29a407bdda6 Mon Sep 17 00:00:00 2001 From: mingfeng Date: Sat, 8 Aug 2026 02:48:14 -0700 Subject: [PATCH 2/6] [core] Let format writers safely handle bundle writes - Remove the producer-side direct-write flag. - Require row-equivalent writes and safe borrowed-buffer handling. - Add schema and allocator fallback checks. - Fix row-count and shredding consistency. - Add regression tests for Arrow, Lance, Mosaic, and Vortex. --- .../paimon/arrow/ArrowBundleRecords.java | 29 ++ .../org/apache/paimon/arrow/ArrowUtils.java | 37 ++ .../reader/ArrowVectorizedRecordIterator.java | 5 +- .../arrow/vector/ArrowFormatWriter.java | 11 + .../arrow/writer/ArrowBundleWriter.java | 16 +- .../paimon/arrow/writer/NativeWriter.java | 6 + .../apache/paimon/arrow/ArrowUtilsTest.java | 44 +++ .../arrow/writer/ArrowBundleWriterTest.java | 188 ++++++++++ .../paimon/format/BundleFormatWriter.java | 13 +- .../InferShreddingWritePlanWriter.java | 61 +--- .../shredding/ShreddingFormatWriter.java | 4 +- .../org/apache/paimon/io/BundleRecords.java | 14 +- .../InferShreddingWritePlanWriterTest.java | 201 +++++++++++ .../paimon/io/RollingFileWriterImpl.java | 6 +- .../apache/paimon/io/RowDataFileWriter.java | 6 +- .../apache/paimon/io/SingleFileWriter.java | 6 +- .../paimon/io/RollingFileWriterTest.java | 19 +- .../paimon/io/RowDataFileWriterTest.java | 188 ++++++++-- .../format/lance/LanceRecordsWriter.java | 15 +- .../paimon/format/lance/jni/LanceWriter.java | 20 +- .../format/lance/LanceRecordsWriterTest.java | 189 ++++++++++ .../format/mosaic/MosaicRecordsWriter.java | 17 +- .../MosaicBundleWriteIntegrationTest.java | 336 ++++++++++++++++++ .../mosaic/MosaicRecordsWriterTest.java | 120 ++++++- .../format/vortex/VortexRecordsWriter.java | 12 +- .../format/vortex/VortexReaderWriterTest.java | 59 +++ 26 files changed, 1481 insertions(+), 141 deletions(-) create mode 100644 paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java create mode 100644 paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java create mode 100644 paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java index 828649010c5d..b179beafb0cd 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowBundleRecords.java +++ b/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 { @@ -49,6 +56,28 @@ 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 arrowFields = vectorSchemaRoot.getSchema().getFields(); + List dataFields = rowType.getFields(); + if (arrowFields.size() != dataFields.size()) { + return false; + } + + Set 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(); diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java index 404e6c009e27..5176f4b4ab32 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java @@ -283,6 +283,22 @@ public static byte[] serializeToIpc(VectorSchemaRoot vsr) { return out.toByteArray(); } + /** Returns whether every vector in the root shares the allocator's root allocator. */ + public static boolean hasSameRootAllocator( + VectorSchemaRoot vectorSchemaRoot, BufferAllocator allocator) { + if (vectorSchemaRoot.getFieldVectors().isEmpty()) { + return false; + } + + BufferAllocator expectedRoot = rootAllocator(allocator); + for (FieldVector vector : vectorSchemaRoot.getFieldVectors()) { + if (!hasSameRootAllocator(vector, expectedRoot)) { + return false; + } + } + return true; + } + public static void serializeToIpc(VectorSchemaRoot vsr, OutputStream out) { try (ArrowStreamWriter writer = new ArrowStreamWriter(vsr, null, out)) { writer.writeBatch(); @@ -316,4 +332,25 @@ private static long zoneCastedTimestampZoneCastToEpoch( return instant.getEpochSecond() * 1_000_000_000 + instant.getNano(); } } + + private static BufferAllocator rootAllocator(BufferAllocator allocator) { + BufferAllocator current = allocator; + while (current.getParentAllocator() != null) { + current = current.getParentAllocator(); + } + return current; + } + + private static boolean hasSameRootAllocator(FieldVector vector, BufferAllocator expectedRoot) { + if (rootAllocator(vector.getAllocator()) != expectedRoot) { + return false; + } + + for (FieldVector child : vector.getChildrenFromFields()) { + if (!hasSameRootAllocator(child, expectedRoot)) { + return false; + } + } + return true; + } } diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java index de2f38e51a15..f96f6195d98f 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java @@ -27,8 +27,9 @@ public interface ArrowVectorizedRecordIterator extends VectorizedRecordIterator /** * Returns a borrowed view of the Arrow vectors backing {@link #batch()}. * - *

The caller does not own the batch and must not retain or close it. Its row order and count - * correspond to {@link #batch()}, and it is valid only until {@link #releaseBatch()}. + *

The caller does not own the batch and must not retain or close it. Its row order, count, + * and values correspond to {@link #batch()}, and it is valid only until {@link + * #releaseBatch()}. */ ArrowBundleRecords arrowBundle(); } diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java index 6133bfb55095..dfcbd72ae6eb 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java @@ -18,6 +18,7 @@ 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; @@ -51,6 +52,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 +173,7 @@ private ArrowFormatWriter( boolean closeAllocatorOnClose) { this.allocator = allocator; this.closeAllocatorOnClose = closeAllocatorOnClose; + this.rowType = rowType; RowType outputRowType = replaceWithShreddingType(rowType, shreddingSchemas); vectorSchemaRoot = @@ -303,6 +306,14 @@ 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() + && rowType.equals(bundle.getRowType()) + && vectorSchemaRoot.getSchema().equals(bundle.getVectorSchemaRoot().getSchema()); + } + private static RowType replaceWithShreddingType( RowType rowType, @Nullable RowType shreddingSchemas) { if (shreddingSchemas == null) { diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java index 39a006b6d7ed..d0ebbc7ddbf8 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java @@ -76,15 +76,21 @@ 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; } + + BundleFormatWriter.super.writeBundle(bundleRecords); } public void add(VectorSchemaRoot vsr) { diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java index 2fee3bcdc361..3a02e6effeca 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java @@ -25,6 +25,12 @@ public abstract class NativeWriter { public abstract long nativeMemoryUsed(); + /** + * Writes an Arrow batch represented by C Data Interface addresses. + * + *

The implementation must consume the batch synchronously or acquire independent ownership + * before returning. Both addresses become invalid immediately after this method returns. + */ public abstract void writeIpcBytes(long arrayAddress, long schemaAddress); public abstract void close(); diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/ArrowUtilsTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/ArrowUtilsTest.java index ee613a05d706..75d70c849b9c 100644 --- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/ArrowUtilsTest.java +++ b/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); + } + } } diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java index a43ec4e0b65e..ba932bc3e1f6 100644 --- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java +++ b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java @@ -18,8 +18,11 @@ package org.apache.paimon.arrow.writer; +import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.arrow.vector.ArrowFormatCWriter; import org.apache.paimon.arrow.vector.ArrowFormatWriter; +import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.columnar.ColumnVector; import org.apache.paimon.data.columnar.VectorizedColumnBatch; import org.apache.paimon.data.columnar.heap.HeapArrayVector; @@ -28,10 +31,13 @@ import org.apache.paimon.data.columnar.heap.HeapMapVector; import org.apache.paimon.data.columnar.heap.HeapRowVector; import org.apache.paimon.fs.PositionOutputStream; +import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.Test; @@ -39,6 +45,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -46,6 +53,181 @@ /** Tests for {@link ArrowBundleWriter}. */ public class ArrowBundleWriterTest { + @Test + public void testArrowBundleFlushesBufferedRowsBeforeDirectWrite() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); + List events = new ArrayList<>(); + NativeWriter nativeWriter = + new NativeWriter() { + @Override + public long nativeMemoryUsed() { + return 0; + } + + @Override + public void writeIpcBytes(long arrayAddress, long schemaAddress) { + events.add("rows"); + cWriter.release(); + } + + @Override + public void close() {} + }; + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + events.add("bundle"); + } + }; + + writer.addElement(GenericRow.of(1)); + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(rowType, allocator)) { + setInt((IntVector) root.getVector("value"), 2); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + + assertThat(events).containsExactly("rows", "bundle"); + writer.close(); + } + + @Test + public void testReorderedArrowBundleFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError("Reordered Arrow bundle must use row fallback."); + } + }; + + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(sourceType, allocator)) { + setInt((IntVector) root.getVector("b"), 20); + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); + } + + @Test + public void testLogicalRowTypeMismatchFallsBackToRows() throws Exception { + RowType writerType = RowType.builder().field("value", DataTypes.INT()).build(); + RowType bundleType = + new RowType( + Collections.singletonList( + new DataField( + 0, "value", DataTypes.INT(), "different description"))); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError( + "Logically incompatible Arrow bundle must use row fallback."); + } + }; + + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(writerType, allocator)) { + setInt((IntVector) root.getVector("value"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, bundleType, true)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); + } + + @Test + public void testNonIdentityNameMappingFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("A", DataTypes.INT()).field("a", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError( + "Non-identity Arrow name mapping must use row fallback."); + } + }; + + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(writerType, allocator)) { + setInt((IntVector) root.getVector("A"), 10); + setInt((IntVector) root.getVector("a"), 20); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, writerType, false)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(20); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); + } + + @Test + public void testMixedAllocatorRootsFallBackToRows() throws Exception { + RowType rowType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError("Mixed-root Arrow bundle must use row fallback."); + } + }; + + try (RootAllocator firstAllocator = new RootAllocator(); + RootAllocator secondAllocator = new RootAllocator()) { + FieldVector firstVector = + writerRoot.getSchema().getFields().get(0).createVector(firstAllocator); + FieldVector secondVector = + writerRoot.getSchema().getFields().get(1).createVector(secondAllocator); + try (VectorSchemaRoot root = + new VectorSchemaRoot( + writerRoot.getSchema(), Arrays.asList(firstVector, secondVector), 1)) { + setInt((IntVector) firstVector, 10); + setInt((IntVector) secondVector, 20); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); + } + @Test public void testAddBatchWithoutDeletionVector() throws IOException { RowType rowType = RowType.of(DataTypes.INT(), DataTypes.BIGINT()); @@ -618,6 +800,12 @@ static class Snapshot { } } + private static void setInt(IntVector vector, int value) { + vector.allocateNew(1); + vector.setSafe(0, value); + vector.setValueCount(1); + } + private static class NoOpPositionOutputStream extends PositionOutputStream { private long pos = 0; diff --git a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java index e0ba09b61243..1ffabd45977d 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java @@ -18,6 +18,7 @@ package org.apache.paimon.format; +import org.apache.paimon.data.InternalRow; import org.apache.paimon.io.BundleRecords; import java.io.IOException; @@ -26,10 +27,18 @@ public interface BundleFormatWriter extends FormatWriter { /** - * Write a bundle of records directly. + * Writes a bundle with semantics equivalent to invoking {@link #addElement} for every record. + * + *

The implementation may consume the bundle natively, convert or copy it, or fall back to + * row-by-row writes. It must not retain borrowed buffers after this method returns unless it + * has copied them or acquired independent ownership. * * @param bundle the records to be written * @throws IOException if exception happens */ - void writeBundle(BundleRecords bundle) throws IOException; + default void writeBundle(BundleRecords bundle) throws IOException { + for (InternalRow row : bundle) { + addElement(row); + } + } } diff --git a/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java b/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java index 22f102d814b5..dacc6f4687f7 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java @@ -26,12 +26,10 @@ import org.apache.paimon.io.BundleRecords; import org.apache.paimon.utils.InternalRowUtils; -import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.io.IOException; import java.util.ArrayList; -import java.util.Iterator; import java.util.List; /** Buffers initial rows, infers a per-file shredding write plan, and writes physical rows. */ @@ -43,7 +41,6 @@ public class InferShreddingWritePlanWriter implements BundleFormatWriter { private final String compression; private final List bufferedRows; - private final List bufferedBundles; @Nullable private FormatWriter actualWriter; private boolean planFinalized = false; @@ -59,7 +56,6 @@ public InferShreddingWritePlanWriter( this.out = out; this.compression = compression; this.bufferedRows = new ArrayList<>(); - this.bufferedBundles = new ArrayList<>(); } @Override @@ -80,12 +76,11 @@ public void addElement(InternalRow row) throws IOException { @Override public void writeBundle(BundleRecords bundle) throws IOException { if (!planFinalized) { - final List rows = new ArrayList<>(); for (InternalRow row : bundle) { - rows.add(InternalRowUtils.copyInternalRow(row, writePlanFactory.logicalRowType())); + bufferedRows.add( + InternalRowUtils.copyInternalRow(row, writePlanFactory.logicalRowType())); + totalBufferedRowCount++; } - bufferedBundles.add(new CopiedBundleRecords(rows)); - totalBufferedRowCount += bundle.rowCount(); if (totalBufferedRowCount >= writePlanFactory.inferBufferRowCount()) { finalizePlanAndFlush(); } @@ -123,57 +118,15 @@ public void close() throws IOException { } private void finalizePlanAndFlush() throws IOException { - ShreddingWritePlan writePlan = writePlanFactory.createWritePlan(collectAllRows()); + ShreddingWritePlan writePlan = writePlanFactory.createWritePlan(bufferedRows); actualWriter = ShreddingWritePlanWriterFactory.createWriterWithPlan( writerFactory, writePlanFactory, out, compression, writePlan); planFinalized = true; - if (!bufferedBundles.isEmpty()) { - BundleFormatWriter bundleWriter = (BundleFormatWriter) actualWriter; - for (BundleRecords bundle : bufferedBundles) { - bundleWriter.writeBundle(bundle); - } - bufferedBundles.clear(); - } else { - for (InternalRow row : bufferedRows) { - actualWriter.addElement(row); - } - bufferedRows.clear(); - } - } - - private List collectAllRows() { - if (bufferedBundles.isEmpty()) { - return bufferedRows; - } - - List allRows = new ArrayList<>(); - for (BundleRecords bundle : bufferedBundles) { - for (InternalRow row : bundle) { - allRows.add(row); - } - } - return allRows; - } - - private static class CopiedBundleRecords implements BundleRecords { - - private final List rows; - - private CopiedBundleRecords(List rows) { - this.rows = rows; - } - - @Override - @Nonnull - public Iterator iterator() { - return rows.iterator(); - } - - @Override - public long rowCount() { - return rows.size(); + for (InternalRow row : bufferedRows) { + actualWriter.addElement(row); } + bufferedRows.clear(); } } diff --git a/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java b/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java index 97a9d4d2da9d..316825b29866 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java @@ -64,9 +64,7 @@ public void writeBundle(BundleRecords bundle) throws IOException { return; } - for (InternalRow row : bundle) { - addElement(row); - } + BundleFormatWriter.super.writeBundle(bundle); } @Override diff --git a/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java b/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java index faeff1a5d776..89758b6157e1 100644 --- a/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java +++ b/paimon-common/src/main/java/org/apache/paimon/io/BundleRecords.java @@ -32,19 +32,11 @@ public interface BundleRecords extends Iterable { /** - * Whether this bundle can be passed directly to a matching format writer. + * The stable, non-negative row count of this batch. * - *

The writer must consume the bundle synchronously. The producer retains ownership of any - * borrowed native buffers and may release them as soon as {@code writeBundle} returns. - */ - default boolean isDirectWriteBundle() { - return false; - } - - /** - * The total row count of this batch. + *

The count must equal the number of records exposed by {@link #iterator()}. * - * @return the number of row count. + * @return the number of records in this batch. */ long rowCount(); } diff --git a/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java b/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java new file mode 100644 index 000000000000..bcf10caf456a --- /dev/null +++ b/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java @@ -0,0 +1,201 @@ +/* + * 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.format.shredding; + +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.data.shredding.ShreddingWritePlan; +import org.apache.paimon.format.FormatWriter; +import org.apache.paimon.fs.PositionOutputStream; +import org.apache.paimon.io.BundleRecords; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link InferShreddingWritePlanWriter}. */ +class InferShreddingWritePlanWriterTest { + + private static final RowType ROW_TYPE = + RowType.builder().field("value", DataTypes.INT()).build(); + + @Test + void testMixedRowsAndBundlesPreserveOrderAndInferenceBoundary() throws Exception { + TestingWriterFactory writerFactory = new TestingWriterFactory(); + TestingWritePlanFactory writePlanFactory = new TestingWritePlanFactory(3); + InferShreddingWritePlanWriter writer = + new InferShreddingWritePlanWriter( + writerFactory, writePlanFactory, new NoOpPositionOutputStream(), "none"); + + writer.addElement(GenericRow.of(1)); + writer.writeBundle(bundle(GenericRow.of(2), GenericRow.of(3), GenericRow.of(4))); + writer.addElement(GenericRow.of(5)); + writer.writeBundle(bundle(GenericRow.of(6), GenericRow.of(7))); + writer.close(); + + assertThat(writePlanFactory.sampleValues).containsExactly(1, 2, 3, 4); + assertThat(writerFactory.writer.values).containsExactly(101, 102, 103, 104, 105, 106, 107); + } + + @Test + void testCloseFinalizesBufferedRowsAndBundle() throws Exception { + TestingWriterFactory writerFactory = new TestingWriterFactory(); + TestingWritePlanFactory writePlanFactory = new TestingWritePlanFactory(10); + InferShreddingWritePlanWriter writer = + new InferShreddingWritePlanWriter( + writerFactory, writePlanFactory, new NoOpPositionOutputStream(), "none"); + + writer.addElement(GenericRow.of(1)); + writer.writeBundle(bundle(GenericRow.of(2), GenericRow.of(3))); + writer.close(); + + assertThat(writePlanFactory.sampleValues).containsExactly(1, 2, 3); + assertThat(writerFactory.writer.values).containsExactly(101, 102, 103); + } + + private static BundleRecords bundle(InternalRow... rows) { + List records = Arrays.asList(rows); + return new BundleRecords() { + @Override + public Iterator iterator() { + return records.iterator(); + } + + @Override + public long rowCount() { + return records.size(); + } + }; + } + + private static class TestingWriterFactory implements SupportsShreddingWritePlan { + + private final TestingFormatWriter writer = new TestingFormatWriter(); + + @Override + public FormatWriter createWithShreddingWritePlan( + PositionOutputStream out, String compression, ShreddingWritePlan writePlan) { + return writer; + } + } + + private static class TestingFormatWriter implements FormatWriter { + + private final List values = new ArrayList<>(); + + @Override + public void addElement(InternalRow element) { + values.add(element.getInt(0)); + } + + @Override + public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { + return false; + } + + @Override + public void close() {} + } + + private static class TestingWritePlanFactory implements ShreddingWritePlanFactory { + + private final int inferBufferRowCount; + private final List sampleValues = new ArrayList<>(); + + private TestingWritePlanFactory(int inferBufferRowCount) { + this.inferBufferRowCount = inferBufferRowCount; + } + + @Override + public RowType logicalRowType() { + return ROW_TYPE; + } + + @Override + public boolean shouldCreateWritePlan() { + return true; + } + + @Override + public boolean shouldInferWritePlan() { + return true; + } + + @Override + public int inferBufferRowCount() { + return inferBufferRowCount; + } + + @Override + public ShreddingWritePlan createWritePlan(List sampleRows) { + for (InternalRow row : sampleRows) { + sampleValues.add(row.getInt(0)); + } + return new TestingWritePlan(); + } + } + + private static class TestingWritePlan implements ShreddingWritePlan { + + @Override + public RowType logicalRowType() { + return ROW_TYPE; + } + + @Override + public RowType physicalRowType() { + return ROW_TYPE; + } + + @Override + public InternalRow toPhysicalRow(InternalRow row) { + return GenericRow.of(row.getInt(0) + 100); + } + } + + private static class NoOpPositionOutputStream extends PositionOutputStream { + + @Override + public long getPos() { + return 0; + } + + @Override + public void write(int b) {} + + @Override + public void write(byte[] b) {} + + @Override + public void write(byte[] b, int off, int len) {} + + @Override + public void flush() {} + + @Override + public void close() {} + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java index 0ebdc5feb790..618733b99718 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java @@ -110,9 +110,11 @@ public void writeBundle(BundleRecords bundle) throws IOException { openCurrentWriter(); } + long previousRecordCount = currentWriter.recordCount(); currentWriter.writeBundle(bundle); - recordCount += bundle.rowCount(); - currentFileRecordCount += bundle.rowCount(); + long writtenRecordCount = currentWriter.recordCount() - previousRecordCount; + recordCount += writtenRecordCount; + currentFileRecordCount += writtenRecordCount; if (rollingFile(true)) { closeCurrentWriter(); diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java index 712c029abc9b..e0e28713e09c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java @@ -146,12 +146,12 @@ public void write(InternalRow row) throws IOException { @Override public void writeBundle(BundleRecords bundle) throws IOException { - if (bundle.isDirectWriteBundle() - && auxiliaryFileWriters.isEmpty() + if (auxiliaryFileWriters.isEmpty() && sequenceNumberTracker.supportsRowCountUpdate() && !requiresPerRecordStats()) { + long previousRecordCount = recordCount(); super.writeBundle(bundle); - sequenceNumberTracker.updateByRowCount(bundle.rowCount()); + sequenceNumberTracker.updateByRowCount(recordCount() - previousRecordCount); return; } diff --git a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java index 8ae8f4bdc075..96f273091dd7 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java @@ -40,6 +40,8 @@ import java.util.Optional; import java.util.function.Function; +import static org.apache.paimon.utils.Preconditions.checkArgument; + /** * A {@link FileWriter} to produce a single file. * @@ -145,6 +147,8 @@ public void writeBundle(BundleRecords bundle) throws IOException { } try { + long rowCount = bundle.rowCount(); + checkArgument(rowCount >= 0, "Row count must not be negative."); if (writer instanceof BundleFormatWriter) { ((BundleFormatWriter) writer).writeBundle(bundle); } else { @@ -152,7 +156,7 @@ public void writeBundle(BundleRecords bundle) throws IOException { writer.addElement(row); } } - recordCount += bundle.rowCount(); + recordCount += rowCount; } catch (Throwable e) { LOG.warn("Exception occurs when writing file {}. Cleaning up.", path, e); abort(); diff --git a/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java index bbfb221fbee3..94ded0e49ced 100644 --- a/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java @@ -169,7 +169,22 @@ public void testRollingByRowsWithBundle() throws IOException { assertThat(files.get(2).rowCount()).isEqualTo(30); } - private static BundleRecords bundle(int rowCount) { + @Test + public void testRollingWriterDoesNotReadBundleRowCountAgain() throws IOException { + initialize("parquet", false, 1024L * 1024 * 1024, 100L); + SingleUseBundleRecords bundle = bundle(150); + + rollingFileWriter.writeBundle(bundle); + rollingFileWriter.close(); + + assertThat(bundle.rowCountCalls).isEqualTo(1); + assertThat(rollingFileWriter.result()) + .singleElement() + .extracting(DataFileMeta::rowCount) + .isEqualTo(150L); + } + + private static SingleUseBundleRecords bundle(int rowCount) { List rows = new ArrayList<>(); for (int i = 0; i < rowCount; i++) { rows.add(GenericRow.of(i)); @@ -314,6 +329,7 @@ private static class SingleUseBundleRecords implements BundleRecords { private final List rows; private boolean iterated; + private int rowCountCalls; private SingleUseBundleRecords(List rows) { this.rows = rows; @@ -330,6 +346,7 @@ public Iterator iterator() { @Override public long rowCount() { + rowCountCalls++; return rows.size(); } } diff --git a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java index eab450de59dc..6ed4cb77ddad 100644 --- a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java @@ -45,7 +45,9 @@ import java.util.List; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; /** Test for {@link RowDataFileWriter}. */ @@ -55,7 +57,7 @@ class RowDataFileWriterTest { private static final RowType ROW_TYPE = RowType.builder().field("id", DataTypes.INT()).build(); @Test - void testBundleFastPathDoesNotIterateRows() throws Exception { + void testEligibleBundleIsForwardedWithoutIteration() throws Exception { FileIO fileIO = fileIO(); TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); LongCounter sequenceCounter = new LongCounter(5); @@ -67,10 +69,13 @@ void testBundleFastPathDoesNotIterateRows() throws Exception { SimpleStatsProducer.disabledProducer(), sequenceCounter, new FileIndexOptions()); - BundleRecords bundle = new DirectNonIterableBundleRecords(3); + TrackingBundleRecords bundle = + trackingRows(GenericRow.of(1), GenericRow.of(2), GenericRow.of(3)); writer.writeBundle(bundle); + assertThat(bundle.rowCountCalls).isEqualTo(1); + assertThat(bundle.iteratorCalls).isZero(); assertThat(formatWriter.writtenBundle).isSameAs(bundle); assertThat(formatWriter.bundleWrites).isEqualTo(1); assertThat(formatWriter.rowWrites).isZero(); @@ -85,7 +90,7 @@ void testBundleFastPathDoesNotIterateRows() throws Exception { } @Test - void testExtractorStatsAllowBundleFastPath() throws Exception { + void testExtractorStatsAllowBundleForwarding() throws Exception { TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); TestingExtractStatsProducer statsProducer = new TestingExtractStatsProducer(); LongCounter sequenceCounter = new LongCounter(5); @@ -97,10 +102,13 @@ void testExtractorStatsAllowBundleFastPath() throws Exception { statsProducer, sequenceCounter, new FileIndexOptions()); - BundleRecords bundle = new DirectNonIterableBundleRecords(3); + TrackingBundleRecords bundle = + trackingRows(GenericRow.of(1), GenericRow.of(2), GenericRow.of(3)); writer.writeBundle(bundle); + assertThat(bundle.rowCountCalls).isEqualTo(1); + assertThat(bundle.iteratorCalls).isZero(); assertThat(formatWriter.writtenBundle).isSameAs(bundle); assertThat(formatWriter.bundleWrites).isEqualTo(1); assertThat(formatWriter.rowWrites).isZero(); @@ -116,22 +124,92 @@ void testExtractorStatsAllowBundleFastPath() throws Exception { } @Test - void testUnmarkedBundleFallsBackToRows() throws Exception { - TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + void testPlainFormatWriterFallsBackToRows() throws Exception { + TestingFormatWriter formatWriter = new TestingFormatWriter(); + LongCounter sequenceCounter = new LongCounter(); RowDataFileWriter writer = createWriter( fileIO(), ROW_TYPE, formatWriter, SimpleStatsProducer.disabledProducer(), - new LongCounter(), + sequenceCounter, new FileIndexOptions()); writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); - assertThat(formatWriter.bundleWrites).isZero(); assertThat(formatWriter.rowWrites).isEqualTo(2); assertThat(writer.recordCount()).isEqualTo(2); + assertThat(sequenceCounter.getValue()).isEqualTo(2); + } + + @Test + void testBundleFormatWriterCanChooseRowFallback() throws Exception { + TestingFallbackBundleFormatWriter formatWriter = new TestingFallbackBundleFormatWriter(); + LongCounter sequenceCounter = new LongCounter(); + RowDataFileWriter writer = + createWriter( + fileIO(), + ROW_TYPE, + formatWriter, + SimpleStatsProducer.disabledProducer(), + sequenceCounter, + new FileIndexOptions()); + + writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); + + assertThat(formatWriter.rowWrites).isEqualTo(2); + assertThat(writer.recordCount()).isEqualTo(2); + assertThat(sequenceCounter.getValue()).isEqualTo(2); + } + + @Test + void testNegativeBundleRowCountIsRejectedBeforeWriting() throws Exception { + FileIO fileIO = fileIO(); + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + LongCounter sequenceCounter = new LongCounter(); + RowDataFileWriter writer = + createWriter( + fileIO, + ROW_TYPE, + formatWriter, + SimpleStatsProducer.disabledProducer(), + sequenceCounter, + new FileIndexOptions()); + + assertThatThrownBy(() -> writer.writeBundle(new InvalidRowCountBundleRecords(-1))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Row count must not be negative."); + + assertThat(formatWriter.bundleWrites).isZero(); + assertThat(formatWriter.closeCalls).isEqualTo(1); + assertThat(writer.recordCount()).isZero(); + assertThat(sequenceCounter.getValue()).isZero(); + verify(fileIO).deleteQuietly(PATH); + } + + @Test + void testBundleWriteFailureCleansUpWithoutAdvancingMetadata() throws Exception { + FileIO fileIO = fileIO(); + IOException failure = new IOException("bundle write failed"); + TestingThrowingBundleFormatWriter formatWriter = + new TestingThrowingBundleFormatWriter(failure); + LongCounter sequenceCounter = new LongCounter(5); + RowDataFileWriter writer = + createWriter( + fileIO, + ROW_TYPE, + formatWriter, + SimpleStatsProducer.disabledProducer(), + sequenceCounter, + new FileIndexOptions()); + + assertThatThrownBy(() -> writer.writeBundle(rows(GenericRow.of(1)))).isSameAs(failure); + + assertThat(formatWriter.closeCalls).isEqualTo(1); + assertThat(writer.recordCount()).isZero(); + assertThat(sequenceCounter.getValue()).isEqualTo(5); + verify(fileIO).deleteQuietly(PATH); } @Test @@ -148,7 +226,7 @@ void testPerRecordStatsFallBackToRows() throws Exception { sequenceCounter, new FileIndexOptions()); - writer.writeBundle(directRows(GenericRow.of(1), GenericRow.of(2))); + writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); assertThat(formatWriter.bundleWrites).isZero(); assertThat(formatWriter.rowWrites).isEqualTo(2); @@ -174,7 +252,7 @@ void testRowTrackingFallsBackToRows() throws Exception { sequenceCounter, new FileIndexOptions()); - writer.writeBundle(directRows(GenericRow.of(1, 7L), GenericRow.of(2, 11L))); + writer.writeBundle(rows(GenericRow.of(1, 7L), GenericRow.of(2, 11L))); assertThat(formatWriter.bundleWrites).isZero(); assertThat(formatWriter.rowWrites).isEqualTo(2); @@ -201,7 +279,7 @@ void testFileIndexFallsBackToRows() throws Exception { new LongCounter(), fileIndexOptions); - writer.writeBundle(directRows(GenericRow.of(1), GenericRow.of(2))); + writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); assertThat(formatWriter.bundleWrites).isZero(); assertThat(formatWriter.rowWrites).isEqualTo(2); @@ -219,7 +297,7 @@ private static FileIO fileIO() throws IOException { private static RowDataFileWriter createWriter( FileIO fileIO, RowType rowType, - TestingBundleFormatWriter formatWriter, + FormatWriter formatWriter, SimpleStatsProducer statsProducer, LongCounter sequenceCounter, FileIndexOptions fileIndexOptions) { @@ -243,16 +321,16 @@ private static BundleRecords rows(InternalRow... rows) { return new ListBundleRecords(Arrays.asList(rows)); } - private static BundleRecords directRows(InternalRow... rows) { - return new DirectListBundleRecords(Arrays.asList(rows)); + private static TrackingBundleRecords trackingRows(InternalRow... rows) { + return new TrackingBundleRecords(Arrays.asList(rows)); } private static class TestingFormatWriterFactory implements FormatWriterFactory, SupportsDirectWrite { - private final TestingBundleFormatWriter writer; + private final FormatWriter writer; - private TestingFormatWriterFactory(TestingBundleFormatWriter writer) { + private TestingFormatWriterFactory(FormatWriter writer) { this.writer = writer; } @@ -267,30 +345,56 @@ public FormatWriter create(FileIO fileIO, Path path, String compression) { } } - private static class TestingBundleFormatWriter implements BundleFormatWriter { + private static class TestingFormatWriter implements FormatWriter { - private int rowWrites; - private int bundleWrites; - private BundleRecords writtenBundle; + int rowWrites; + int closeCalls; @Override public void addElement(InternalRow element) { rowWrites++; } + @Override + public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { + return false; + } + + @Override + public void close() { + closeCalls++; + } + } + + private static class TestingBundleFormatWriter extends TestingFormatWriter + implements BundleFormatWriter { + + private int bundleWrites; + private BundleRecords writtenBundle; + @Override public void writeBundle(BundleRecords bundle) { bundleWrites++; writtenBundle = bundle; } + } - @Override - public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { - return false; + private static class TestingFallbackBundleFormatWriter extends TestingFormatWriter + implements BundleFormatWriter {} + + private static class TestingThrowingBundleFormatWriter extends TestingFormatWriter + implements BundleFormatWriter { + + private final IOException failure; + + private TestingThrowingBundleFormatWriter(IOException failure) { + this.failure = failure; } @Override - public void close() {} + public void writeBundle(BundleRecords bundle) throws IOException { + throw failure; + } } private static class TestingStatsProducer implements SimpleStatsProducer { @@ -344,22 +448,17 @@ public SimpleColStats[] extract(FileIO fileIO, Path path, long length) { } } - private static class DirectNonIterableBundleRecords implements BundleRecords { + private static class InvalidRowCountBundleRecords implements BundleRecords { private final long rowCount; - private DirectNonIterableBundleRecords(long rowCount) { + private InvalidRowCountBundleRecords(long rowCount) { this.rowCount = rowCount; } - @Override - public boolean isDirectWriteBundle() { - return true; - } - @Override public Iterator iterator() { - throw new AssertionError("Direct bundle write must not iterate rows."); + throw new AssertionError("Invalid row count must be rejected before row iteration."); } @Override @@ -368,34 +467,45 @@ public long rowCount() { } } - private static class ListBundleRecords implements BundleRecords { + private static class TrackingBundleRecords implements BundleRecords { private final List rows; + private int iteratorCalls; + private int rowCountCalls; - private ListBundleRecords(List rows) { + private TrackingBundleRecords(List rows) { this.rows = rows; } @Override public Iterator iterator() { + iteratorCalls++; return rows.iterator(); } @Override public long rowCount() { + rowCountCalls++; return rows.size(); } } - private static class DirectListBundleRecords extends ListBundleRecords { + private static class ListBundleRecords implements BundleRecords { - private DirectListBundleRecords(List rows) { - super(rows); + private final List rows; + + private ListBundleRecords(List rows) { + this.rows = rows; } @Override - public boolean isDirectWriteBundle() { - return true; + public Iterator iterator() { + return rows.iterator(); + } + + @Override + public long rowCount() { + return rows.size(); } } } diff --git a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java index ee1f572462d7..d315f44633ec 100644 --- a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java +++ b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java @@ -19,6 +19,7 @@ package org.apache.paimon.format.lance; import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.arrow.vector.ArrowFormatWriter; import org.apache.paimon.data.InternalRow; import org.apache.paimon.format.BundleFormatWriter; @@ -65,12 +66,18 @@ public void addElement(InternalRow internalRow) throws IOException { @Override public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { - add(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); - } else { - for (InternalRow row : bundleRecords) { - addElement(row); + ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; + VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot(); + if (arrowFormatWriter.isArrowBundleSchemaCompatible(arrowBundle) + && ArrowUtils.hasSameRootAllocator(root, arrowFormatWriter.getAllocator())) { + flush(); + nativeWriter.ensureInitialized(arrowFormatWriter.getAllocator()); + add(root); + return; } } + + BundleFormatWriter.super.writeBundle(bundleRecords); } public void add(VectorSchemaRoot vsr) throws IOException { diff --git a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java index f9329907f951..b3e29eb0a0b3 100644 --- a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java +++ b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java @@ -18,6 +18,8 @@ package org.apache.paimon.format.lance.jni; +import org.apache.paimon.arrow.ArrowUtils; + import com.lancedb.lance.file.LanceFileWriter; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; @@ -32,6 +34,7 @@ public class LanceWriter { private final String path; private final Map storageOptions; private LanceFileWriter writer; + private BufferAllocator allocator; private long bytesWritten = 0; public LanceWriter(String path, Map storageOptions) { @@ -44,12 +47,25 @@ public Long getWrittenPosition() { } public void writeVsr(VectorSchemaRoot vsr) throws IOException { - initWriteLazy(vsr.getVector(0).getAllocator()); + BufferAllocator sourceAllocator = vsr.getVector(0).getAllocator(); + initWriteLazy(sourceAllocator); + if (!ArrowUtils.hasSameRootAllocator(vsr, allocator)) { + throw new IllegalArgumentException( + "Lance writer cannot consume Arrow buffers from a different allocator root."); + } this.bytesWritten += vsr.getFieldVectors().stream().mapToLong(FieldVector::getBufferSize).sum(); this.writer.write(vsr); } + /** + * Initializes the native writer with an allocator whose lifetime is owned by the surrounding + * format writer. + */ + public void ensureInitialized(BufferAllocator bufferAllocator) throws IOException { + initWriteLazy(bufferAllocator); + } + public void close() throws IOException { if (writer != null) { try { @@ -58,6 +74,7 @@ public void close() throws IOException { throw new IOException(e); } this.writer = null; + this.allocator = null; } } @@ -68,6 +85,7 @@ public String path() { private void initWriteLazy(BufferAllocator bufferAllocator) throws IOException { if (writer == null) { writer = LanceFileWriter.open(path, bufferAllocator, null, storageOptions); + allocator = bufferAllocator; } } } diff --git a/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java new file mode 100644 index 000000000000..efd3e2a821e0 --- /dev/null +++ b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java @@ -0,0 +1,189 @@ +/* + * 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.format.lance; + +import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; +import org.apache.paimon.arrow.vector.ArrowFormatWriter; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.format.lance.jni.LanceWriter; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link LanceRecordsWriter}. */ +class LanceRecordsWriterTest { + + @Test + void testArrowBundlePreservesRowBundleRowOrder() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true); + BufferAllocator writerAllocator = arrowWriter.getAllocator(); + CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + writer.addElement(GenericRow.of(1)); + try (BufferAllocator sourceAllocator = + arrowWriter + .getAllocator() + .newChildAllocator("lance-bundle-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + setInt((IntVector) root.getVector("value"), 2); + root.setRowCount(1); + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + writer.addElement(GenericRow.of(3)); + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(3); + assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(1); + assertThat(nativeWriter.snapshots.get(1).values.get(0)).containsExactly(2); + assertThat(nativeWriter.snapshots.get(2).values.get(0)).containsExactly(3); + assertThat(nativeWriter.initializedAllocator).isSameAs(writerAllocator); + } + + @Test + void testReorderedArrowBundleFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(writerType, 1024, true); + CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + try (BufferAllocator sourceAllocator = + arrowWriter + .getAllocator() + .newChildAllocator("lance-schema-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { + setInt((IntVector) root.getVector("b"), 20); + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + Snapshot snapshot = nativeWriter.snapshots.get(0); + assertThat(snapshot.fieldNames).containsExactly("a", "b"); + assertThat(snapshot.values.get(0)).containsExactly(10); + assertThat(snapshot.values.get(1)).containsExactly(20); + } + + @Test + void testDifferentAllocatorRootFallsBackToRows() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true); + CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + try (BufferAllocator sourceAllocator = new org.apache.arrow.memory.RootAllocator(); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + nativeWriter.disallowedRoot = root; + setInt((IntVector) root.getVector("value"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + writer.close(); + + assertThat(nativeWriter.disallowedRootWrites).isZero(); + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(10); + } + + private static void setInt(IntVector vector, int value) { + vector.allocateNew(1); + vector.setSafe(0, value); + vector.setValueCount(1); + } + + private static class CapturingLanceWriter extends LanceWriter { + + private final List snapshots = new ArrayList<>(); + private BufferAllocator initializedAllocator; + private VectorSchemaRoot disallowedRoot; + private int disallowedRootWrites; + + private CapturingLanceWriter() { + super("unused", Collections.emptyMap()); + } + + @Override + public void ensureInitialized(BufferAllocator bufferAllocator) { + initializedAllocator = bufferAllocator; + } + + @Override + public void writeVsr(VectorSchemaRoot root) { + if (root == disallowedRoot) { + disallowedRootWrites++; + } + List fieldNames = + root.getSchema().getFields().stream() + .map(field -> field.getName()) + .collect(Collectors.toList()); + List> values = new ArrayList<>(); + for (int column = 0; column < root.getFieldVectors().size(); column++) { + IntVector vector = (IntVector) root.getVector(column); + List columnValues = new ArrayList<>(); + for (int row = 0; row < root.getRowCount(); row++) { + columnValues.add(vector.get(row)); + } + values.add(columnValues); + } + snapshots.add(new Snapshot(fieldNames, values)); + } + + @Override + public void close() throws IOException {} + + @Override + public String path() { + return "unused"; + } + } + + private static class Snapshot { + + private final List fieldNames; + private final List> values; + + private Snapshot(List fieldNames, List> values) { + this.fieldNames = fieldNames; + this.values = values; + } + } +} 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 944f78c4b147..3e12e626e671 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 @@ -19,6 +19,7 @@ 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.InternalRow; import org.apache.paimon.format.BundleFormatWriter; @@ -125,13 +126,19 @@ public void addElement(InternalRow internalRow) { @Override public void writeBundle(BundleRecords bundleRecords) { if (bundleRecords instanceof ArrowBundleRecords) { - flush(); - nativeWriter.write(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); - } else { - for (InternalRow row : bundleRecords) { - addElement(row); + ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; + VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot(); + if (arrowFormatWriter.isArrowBundleSchemaCompatible(arrowBundle) + && ArrowUtils.hasSameRootAllocator(root, allocator)) { + flush(); + nativeWriter.write(root); + return; } } + + for (InternalRow row : bundleRecords) { + addElement(row); + } } @Override diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java new file mode 100644 index 000000000000..eb0280329e2e --- /dev/null +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java @@ -0,0 +1,336 @@ +/* + * 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.format.mosaic; + +import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.fileindex.FileIndexOptions; +import org.apache.paimon.format.FileFormatFactory; +import org.apache.paimon.format.FormatReaderContext; +import org.apache.paimon.format.FormatReaderFactory; +import org.apache.paimon.format.FormatWriter; +import org.apache.paimon.format.FormatWriterFactory; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.PositionOutputStream; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.io.FileWriterContext; +import org.apache.paimon.io.RowDataFileWriter; +import org.apache.paimon.io.SimpleStatsProducer; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.mosaic.MosaicWriter; +import org.apache.paimon.mosaic.WriterOptions; +import org.apache.paimon.options.Options; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.LongCounter; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +/** Integration tests for bundle dispatch from core file writers to Mosaic native writes. */ +class MosaicBundleWriteIntegrationTest { + + private static final FileFormatFactory.FormatContext FORMAT_CONTEXT = + new FileFormatFactory.FormatContext(new Options(), 1024, 1024); + + @TempDir java.nio.file.Path tempDir; + + @Test + void testCompatibleArrowBundleUsesFullDirectWritePath() throws Exception { + assumeTrue(isNativeAvailable(), "Mosaic native library not available"); + + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + Path path = newPath("direct"); + LocalFileIO fileIO = new LocalFileIO(); + RootAllocator writerAllocator = new RootAllocator(); + LongCounter sequenceCounter = new LongCounter(5); + AtomicReference nativeWriterRef = new AtomicReference<>(); + + try (RowDataFileWriter writer = + createWriter( + fileIO, path, rowType, writerAllocator, sequenceCounter, nativeWriterRef)) { + TrackingMosaicWriter nativeWriter = nativeWriterRef.get(); + assertThat(nativeWriter).isNotNull(); + + try (BufferAllocator sourceAllocator = + writerAllocator.newChildAllocator( + "mosaic-direct-integration-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + setInts((IntVector) root.getVector("value"), 1, 2, 3); + root.setRowCount(3); + + TrackingDirectArrowBundleRecords bundle = + new TrackingDirectArrowBundleRecords(root, rowType); + nativeWriter.expectDirectRoot(root); + writer.writeBundle(bundle); + nativeWriter.clearExpectedRoot(); + + assertThat(bundle.iteratorCalls).isZero(); + assertThat(nativeWriter.directWrites).isEqualTo(1); + assertThat(writer.recordCount()).isEqualTo(3); + assertThat(sequenceCounter.getValue()).isEqualTo(8); + } + + // The borrowed source root has already been released. Closing the native writer must + // not access it again. + writer.close(); + DataFileMeta result = writer.result(); + assertThat(result.rowCount()).isEqualTo(3); + assertThat(result.minSequenceNumber()).isEqualTo(5); + assertThat(result.maxSequenceNumber()).isEqualTo(7); + } + + assertThat(readRows(fileIO, path, rowType)) + .containsExactly( + Collections.singletonList(1), + Collections.singletonList(2), + Collections.singletonList(3)); + } + + @Test + void testIncompatibleArrowSchemaFallsBackThroughFullWritePath() throws Exception { + assumeTrue(isNativeAvailable(), "Mosaic native library not available"); + + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + Path path = newPath("fallback"); + LocalFileIO fileIO = new LocalFileIO(); + RootAllocator writerAllocator = new RootAllocator(); + LongCounter sequenceCounter = new LongCounter(); + AtomicReference nativeWriterRef = new AtomicReference<>(); + + try (RowDataFileWriter writer = + createWriter( + fileIO, + path, + writerType, + writerAllocator, + sequenceCounter, + nativeWriterRef)) { + TrackingMosaicWriter nativeWriter = nativeWriterRef.get(); + assertThat(nativeWriter).isNotNull(); + + try (BufferAllocator sourceAllocator = + writerAllocator.newChildAllocator( + "mosaic-fallback-integration-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { + setInts((IntVector) root.getVector("b"), 20, 21); + setInts((IntVector) root.getVector("a"), 10, 11); + root.setRowCount(2); + + TrackingArrowBundleRecords bundle = + new TrackingArrowBundleRecords(root, writerType); + nativeWriter.expectDirectRoot(root); + writer.writeBundle(bundle); + nativeWriter.clearExpectedRoot(); + + assertThat(bundle.iteratorCalls).isEqualTo(1); + assertThat(nativeWriter.directWrites).isZero(); + assertThat(writer.recordCount()).isEqualTo(2); + assertThat(sequenceCounter.getValue()).isEqualTo(2); + } + + writer.close(); + assertThat(writer.result().rowCount()).isEqualTo(2); + } + + assertThat(readRows(fileIO, path, writerType)) + .containsExactly(asList(10, 20), asList(11, 21)); + } + + private Path newPath(String prefix) { + return new Path(tempDir.toUri().toString(), prefix + ".mosaic"); + } + + private static RowDataFileWriter createWriter( + LocalFileIO fileIO, + Path path, + RowType rowType, + RootAllocator allocator, + LongCounter sequenceCounter, + AtomicReference nativeWriterRef) { + FormatWriterFactory writerFactory = + new FormatWriterFactory() { + @Override + public FormatWriter create(PositionOutputStream out, String compression) { + assertThat(compression).isEqualTo("zstd"); + return new MosaicRecordsWriter( + out, + rowType, + FORMAT_CONTEXT, + Collections.emptyList(), + null, + allocator, + (outputStream, arrowSchema, options, bufferAllocator) -> { + TrackingMosaicWriter writer = + new TrackingMosaicWriter( + outputStream, + arrowSchema, + options, + bufferAllocator); + nativeWriterRef.set(writer); + return writer; + }); + } + }; + + return new RowDataFileWriter( + fileIO, + new FileWriterContext( + writerFactory, SimpleStatsProducer.disabledProducer(), "zstd"), + path, + rowType, + 1L, + () -> sequenceCounter, + new FileIndexOptions(), + FileSource.APPEND, + false, + false, + false, + null); + } + + private static List> readRows(LocalFileIO fileIO, Path path, RowType rowType) + throws IOException { + MosaicFileFormat format = new MosaicFileFormat(FORMAT_CONTEXT); + FormatReaderFactory readerFactory = + format.createReaderFactory(rowType, rowType, Collections.emptyList()); + List> rows = new ArrayList<>(); + try (RecordReader reader = + readerFactory.createReader( + new FormatReaderContext(fileIO, path, fileIO.getFileSize(path)))) { + reader.forEachRemaining( + row -> { + List values = new ArrayList<>(rowType.getFieldCount()); + for (int i = 0; i < rowType.getFieldCount(); i++) { + values.add(row.getInt(i)); + } + rows.add(values); + }); + } + return rows; + } + + private static List asList(int first, int second) { + List values = new ArrayList<>(2); + values.add(first); + values.add(second); + return values; + } + + private static void setInts(IntVector vector, int... values) { + vector.allocateNew(values.length); + for (int i = 0; i < values.length; i++) { + vector.setSafe(i, values[i]); + } + vector.setValueCount(values.length); + } + + private static boolean isNativeAvailable() { + try { + Class.forName("org.apache.paimon.mosaic.NativeLib"); + return true; + } catch (Throwable t) { + return false; + } + } + + private static class TrackingDirectArrowBundleRecords extends ArrowBundleRecords { + + private int iteratorCalls; + + private TrackingDirectArrowBundleRecords(VectorSchemaRoot root, RowType rowType) { + super(root, rowType, true); + } + + @Override + public Iterator iterator() { + iteratorCalls++; + return super.iterator(); + } + } + + private static class TrackingArrowBundleRecords extends ArrowBundleRecords { + + private int iteratorCalls; + + private TrackingArrowBundleRecords(VectorSchemaRoot root, RowType rowType) { + super(root, rowType, true); + } + + @Override + public Iterator iterator() { + iteratorCalls++; + return super.iterator(); + } + } + + private static class TrackingMosaicWriter extends MosaicWriter { + + private VectorSchemaRoot expectedDirectRoot; + private int directWrites; + + private TrackingMosaicWriter( + OutputStream outputStream, + Schema schema, + WriterOptions options, + BufferAllocator allocator) { + super(outputStream, schema, options, allocator); + } + + private void expectDirectRoot(VectorSchemaRoot root) { + expectedDirectRoot = root; + } + + private void clearExpectedRoot() { + expectedDirectRoot = null; + } + + @Override + public void write(VectorSchemaRoot root) { + if (root == expectedDirectRoot) { + directWrites++; + } + super.write(root); + } + } +} 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 58a8dc5252a2..b1435b30b5bb 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 @@ -18,12 +18,18 @@ package org.apache.paimon.format.mosaic; +import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.format.FileFormatFactory; +import org.apache.paimon.mosaic.MosaicWriter; import org.apache.paimon.options.Options; 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.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.Test; import java.io.ByteArrayOutputStream; @@ -31,15 +37,22 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.same; +import static org.mockito.Mockito.verify; /** Test for {@link MosaicRecordsWriter}. */ class MosaicRecordsWriterTest { + private static final FileFormatFactory.FormatContext FORMAT_CONTEXT = + new FileFormatFactory.FormatContext(new Options(), 1024, 1024); + @Test void testConstructorFailureClosesCreatedResources() { RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING()); - FileFormatFactory.FormatContext formatContext = - new FileFormatFactory.FormatContext(new Options(), 1024, 1024); CloseCountingRootAllocator allocator = new CloseCountingRootAllocator(); RuntimeException failure = new RuntimeException("native writer failed"); @@ -48,7 +61,7 @@ void testConstructorFailureClosesCreatedResources() { new MosaicRecordsWriter( new ByteArrayOutputStream(), rowType, - formatContext, + FORMAT_CONTEXT, Collections.emptyList(), null, allocator, @@ -60,6 +73,107 @@ void testConstructorFailureClosesCreatedResources() { assertThat(allocator.closeCount()).isEqualTo(1); } + @Test + void testCompatibleSameRootArrowBundleUsesDirectWrite() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + RootAllocator writerAllocator = new RootAllocator(); + MosaicWriter nativeWriter = mock(MosaicWriter.class); + MosaicRecordsWriter writer = createWriter(rowType, writerAllocator, nativeWriter); + + try (BufferAllocator sourceAllocator = + writerAllocator.newChildAllocator("mosaic-direct-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + setInt((IntVector) root.getVector("value"), 1); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + + verify(nativeWriter).write(same(root)); + } + writer.close(); + } + + @Test + void testDifferentAllocatorRootFallsBackToRows() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + RootAllocator writerAllocator = new RootAllocator(); + MosaicWriter nativeWriter = mock(MosaicWriter.class); + MosaicRecordsWriter writer = createWriter(rowType, writerAllocator, nativeWriter); + + try (RootAllocator sourceAllocator = new RootAllocator(); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + setInt((IntVector) root.getVector("value"), 1); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + + verify(nativeWriter, never()).write(same(root)); + } + writer.close(); + + verify(nativeWriter).write(any(VectorSchemaRoot.class)); + } + + @Test + void testReorderedArrowBundleFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + RootAllocator writerAllocator = new RootAllocator(); + MosaicWriter nativeWriter = mock(MosaicWriter.class); + MosaicRecordsWriter writer = createWriter(writerType, writerAllocator, nativeWriter); + doAnswer( + invocation -> { + VectorSchemaRoot written = invocation.getArgument(0); + assertThat(written.getSchema().getFields().get(0).getName()) + .isEqualTo("a"); + assertThat(written.getSchema().getFields().get(1).getName()) + .isEqualTo("b"); + assertThat(((IntVector) written.getVector("a")).get(0)).isEqualTo(10); + assertThat(((IntVector) written.getVector("b")).get(0)).isEqualTo(20); + return null; + }) + .when(nativeWriter) + .write(any(VectorSchemaRoot.class)); + + try (BufferAllocator sourceAllocator = + writerAllocator.newChildAllocator("mosaic-schema-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { + setInt((IntVector) root.getVector("b"), 20); + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); + + verify(nativeWriter, never()).write(same(root)); + } + writer.close(); + + verify(nativeWriter).write(any(VectorSchemaRoot.class)); + } + + private static MosaicRecordsWriter createWriter( + RowType rowType, RootAllocator allocator, MosaicWriter nativeWriter) { + return new MosaicRecordsWriter( + new ByteArrayOutputStream(), + rowType, + FORMAT_CONTEXT, + Collections.emptyList(), + null, + allocator, + (outputStream, arrowSchema, options, bufferAllocator) -> nativeWriter); + } + + private static void setInt(IntVector vector, int value) { + vector.allocateNew(1); + vector.setSafe(0, value); + vector.setValueCount(1); + } + private static class CloseCountingRootAllocator extends RootAllocator { private int closeCount; diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java index 7798fa09d419..7d83f5698906 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java @@ -102,13 +102,15 @@ public void addElement(InternalRow internalRow) throws IOException { @Override public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { - flush(); - writeBundleVsr(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); - } else { - for (InternalRow row : bundleRecords) { - addElement(row); + ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; + if (currentWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle)) { + flush(); + writeBundleVsr(arrowBundle.getVectorSchemaRoot()); + return; } } + + BundleFormatWriter.super.writeBundle(bundleRecords); } @Override diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java index e97e3cf80fcc..b90e601610f7 100644 --- a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java @@ -19,6 +19,7 @@ package org.apache.paimon.format.vortex; 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.GenericRow; @@ -45,6 +46,9 @@ import org.apache.paimon.utils.RoaringBitmap32; import dev.vortex.jni.NativeRuntime; +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 org.junit.jupiter.api.io.TempDir; @@ -320,6 +324,55 @@ public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( } } + @Test + public void testReorderedArrowBundleFallsBackToRows(@TempDir java.nio.file.Path tempDir) + throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + + Options options = new Options(); + VortexFileFormat format = + new VortexFileFormatFactory() + .create(new FileFormatFactory.FormatContext(options, 1024, 1024)); + FileIO fileIO = new LocalFileIO(); + Path testFile = + new Path(new Path(tempDir.toUri()), "test_reordered_bundle_" + UUID.randomUUID()); + + try (FormatWriter writer = + ((SupportsDirectWrite) format.createWriterFactory(writerType)) + .create(fileIO, testFile, ""); + RootAllocator sourceAllocator = new RootAllocator(); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { + setInt((IntVector) root.getVector("b"), 20); + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + + ((BundleFormatWriter) writer) + .writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + + InternalRowSerializer serializer = new InternalRowSerializer(writerType); + FormatReaderFactory readerFactory = + format.createReaderFactory(writerType, writerType, null); + try (RecordReader reader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, testFile, fileIO.getFileSize(testFile), null)); + RecordReaderIterator iterator = new RecordReaderIterator<>(reader)) { + List actualRows = new ArrayList<>(); + while (iterator.hasNext()) { + actualRows.add(serializer.copy(iterator.next())); + } + + assertEquals(1, actualRows.size()); + assertEquals(10, actualRows.get(0).getInt(0)); + assertEquals(20, actualRows.get(0).getInt(1)); + } + } + @Test public void testReadWithSelection(@TempDir java.nio.file.Path tempDir) throws Exception { RowType rowType = RowType.of(DataTypes.INT(), DataTypes.STRING()); @@ -367,6 +420,12 @@ public void testReadWithSelection(@TempDir java.nio.file.Path tempDir) throws Ex } } + private static void setInt(IntVector vector, int value) { + vector.allocateNew(1); + vector.setSafe(0, value); + vector.setValueCount(1); + } + @Test public void testReadWithVirtualRowTrackingField(@TempDir java.nio.file.Path tempDir) throws Exception { From c80ded9b6a86e1b686d6579069353b714aa2085c Mon Sep 17 00:00:00 2001 From: mingfeng Date: Sat, 8 Aug 2026 07:27:07 -0700 Subject: [PATCH 3/6] [core][mosaic] Fix RowDataFileWriter test compilation --- .../test/java/org/apache/paimon/io/RowDataFileWriterTest.java | 2 ++ .../paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java | 2 ++ 2 files changed, 4 insertions(+) diff --git a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java index 6ed4cb77ddad..57e7aa8a65e4 100644 --- a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java @@ -314,6 +314,8 @@ private static RowDataFileWriter createWriter( false, false, false, + null, + null, null); } diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java index eb0280329e2e..12a8f3cb0c76 100644 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java @@ -226,6 +226,8 @@ public FormatWriter create(PositionOutputStream out, String compression) { false, false, false, + null, + null, null); } From cc9c03d89b3378457652967710ef3d6ce7e5f71c Mon Sep 17 00:00:00 2001 From: mingfeng Date: Sun, 9 Aug 2026 09:16:32 -0700 Subject: [PATCH 4/6] [common][core][mosaic] Restrict direct bundle writes to compatible Mosaic writers Preserve existing Arrow, Lance, and Vortex writer behavior by keeping them on the row path unless a writer explicitly supports row-equivalent bundle writes. Enforce the shredding inference bound and keep Mosaic native writes behind schema and allocator-root compatibility checks with safe row fallback. --- .../org/apache/paimon/arrow/ArrowUtils.java | 5 +- .../paimon/arrow/reader/ArrowBatchReader.java | 68 ++-- .../reader/ArrowVectorizedRecordIterator.java | 35 -- .../arrow/vector/ArrowFormatWriter.java | 63 +++- .../arrow/writer/ArrowBundleWriter.java | 16 +- .../paimon/arrow/writer/NativeWriter.java | 6 - .../arrow/reader/ArrowBatchReaderTest.java | 89 ----- .../arrow/vector/ArrowFormatWriterTest.java | 29 ++ .../arrow/writer/ArrowBundleWriterTest.java | 188 ---------- .../paimon/format/BundleFormatWriter.java | 25 +- .../InferShreddingWritePlanWriter.java | 7 +- .../shredding/ShreddingFormatWriter.java | 4 +- .../InferShreddingWritePlanWriterTest.java | 16 +- .../paimon/io/RollingFileWriterImpl.java | 7 +- .../apache/paimon/io/RowDataFileWriter.java | 7 +- .../apache/paimon/io/SingleFileWriter.java | 5 + .../paimon/io/RollingFileWriterTest.java | 17 - .../paimon/io/RowDataFileWriterTest.java | 30 +- .../format/lance/LanceRecordsWriter.java | 15 +- .../paimon/format/lance/jni/LanceWriter.java | 20 +- .../format/lance/LanceRecordsWriterTest.java | 189 ---------- .../format/mosaic/MosaicRecordsWriter.java | 7 + .../MosaicBundleWriteIntegrationTest.java | 338 ------------------ .../format/mosaic/MosaicReaderWriterTest.java | 58 +++ .../mosaic/MosaicRecordsWriterTest.java | 3 +- .../format/vortex/VortexRecordsWriter.java | 12 +- .../format/vortex/VortexReaderWriterTest.java | 59 --- 27 files changed, 271 insertions(+), 1047 deletions(-) delete mode 100644 paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java delete mode 100644 paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java delete mode 100644 paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java delete mode 100644 paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java index 5176f4b4ab32..2122e00a3ba1 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java @@ -283,7 +283,10 @@ public static byte[] serializeToIpc(VectorSchemaRoot vsr) { return out.toByteArray(); } - /** Returns whether every vector in the root shares the allocator's root allocator. */ + /** + * 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) { if (vectorSchemaRoot.getFieldVectors().isEmpty()) { diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java index 5548dfb2fc65..ab939d70b4e7 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java +++ b/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 reusableBatch; + private final VectorizedColumnBatch batch; private final Arrow2PaimonVectorConverter[] convertors; private final RowType projectedRowType; private final boolean caseSensitive; @@ -57,11 +57,12 @@ public ArrowBatchReader( RowType rowType, boolean caseSensitive, Arrow2PaimonVectorConverter.Arrow2PaimonVectorConvertorVisitor visitor) { - this.reusableBatch = new VectorizedColumnBatch(new ColumnVector[rowType.getFieldCount()]); + ColumnVector[] columnVectors = new ColumnVector[rowType.getFieldCount()]; this.convertors = new Arrow2PaimonVectorConverter[rowType.getFieldCount()]; + this.batch = new VectorizedColumnBatch(columnVectors); this.projectedRowType = rowType; - for (int i = 0; i < convertors.length; i++) { + for (int i = 0; i < columnVectors.length; i++) { this.convertors[i] = Arrow2PaimonVectorConverter.construct(visitor, rowType.getTypeAt(i)); } @@ -69,41 +70,6 @@ public ArrowBatchReader( } public Iterable readBatch(VectorSchemaRoot vsr) { - populateBatch(vsr, reusableBatch); - int rowCount = reusableBatch.getNumRows(); - final ColumnarRow columnarRow = new ColumnarRow(reusableBatch); - return () -> - new Iterator() { - 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. - * - *

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 arrowFieldIndex = new HashMap<>(); @@ -117,14 +83,32 @@ private void populateBatch(VectorSchemaRoot vsr, VectorizedColumnBatch targetBat mapping[i] = arrowFieldIndex.getOrDefault(fieldName, -1); } - for (int i = 0; i < targetBatch.columns.length; i++) { + for (int i = 0; i < batch.columns.length; i++) { if (mapping[i] >= 0) { - targetBatch.columns[i] = convertors[i].convertVector(vsr.getVector(mapping[i])); + batch.columns[i] = convertors[i].convertVector(vsr.getVector(mapping[i])); } else { - targetBatch.columns[i] = AllNullColumnVector.INSTANCE; + batch.columns[i] = AllNullColumnVector.INSTANCE; } } - targetBatch.setNumRows(vsr.getRowCount()); + int rowCount = vsr.getRowCount(); + batch.setNumRows(vsr.getRowCount()); + final ColumnarRow columnarRow = new ColumnarRow(batch); + return () -> + new Iterator() { + private int position = 0; + + @Override + public boolean hasNext() { + return position < rowCount; + } + + @Override + public InternalRow next() { + columnarRow.setRowId(position); + position++; + return columnarRow; + } + }; } } diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java deleted file mode 100644 index f96f6195d98f..000000000000 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowVectorizedRecordIterator.java +++ /dev/null @@ -1,35 +0,0 @@ -/* - * 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.ArrowBundleRecords; -import org.apache.paimon.reader.VectorizedRecordIterator; - -/** A {@link VectorizedRecordIterator} which can expose its Arrow batch for direct bundle writes. */ -public interface ArrowVectorizedRecordIterator extends VectorizedRecordIterator { - - /** - * Returns a borrowed view of the Arrow vectors backing {@link #batch()}. - * - *

The caller does not own the batch and must not retain or close it. Its row order, count, - * and values correspond to {@link #batch()}, and it is valid only until {@link - * #releaseBatch()}. - */ - ArrowBundleRecords arrowBundle(); -} diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java index dfcbd72ae6eb..f120b141cd84 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java @@ -26,10 +26,14 @@ 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; @@ -310,10 +314,67 @@ public BufferAllocator getAllocator() { public boolean isArrowBundleSchemaCompatible(ArrowBundleRecords bundle) { return !bundle.getVectorSchemaRoot().getFieldVectors().isEmpty() && bundle.hasIdentityMapping() - && rowType.equals(bundle.getRowType()) + && 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 leftFields = ((RowType) left).getFields(); + List 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) { diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java index d0ebbc7ddbf8..39a006b6d7ed 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java @@ -76,21 +76,15 @@ public void addElement(InternalRow internalRow) { @Override public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { - ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; - VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot(); - if (arrowFormatWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle) - && ArrowUtils.hasSameRootAllocator(root, root.getVector(0).getAllocator())) { - flush(); - add(root); - return; - } + add(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); } else if (bundleRecords instanceof VectorizedBundleRecords) { VectorizedBundleRecords records = (VectorizedBundleRecords) bundleRecords; add(records.batch(), records.selected()); - return; + } else { + for (InternalRow row : bundleRecords) { + addElement(row); + } } - - BundleFormatWriter.super.writeBundle(bundleRecords); } public void add(VectorSchemaRoot vsr) { diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java index 3a02e6effeca..2fee3bcdc361 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/NativeWriter.java @@ -25,12 +25,6 @@ public abstract class NativeWriter { public abstract long nativeMemoryUsed(); - /** - * Writes an Arrow batch represented by C Data Interface addresses. - * - *

The implementation must consume the batch synchronously or acquire independent ownership - * before returning. Both addresses become invalid immediately after this method returns. - */ public abstract void writeIpcBytes(long arrayAddress, long schemaAddress); public abstract void close(); diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java deleted file mode 100644 index c70f21ad446e..000000000000 --- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/reader/ArrowBatchReaderTest.java +++ /dev/null @@ -1,89 +0,0 @@ -/* - * 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; - } -} diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java index 40aa87dae970..fd257450aaa9 100644 --- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java +++ b/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) { diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java index ba932bc3e1f6..a43ec4e0b65e 100644 --- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java +++ b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java @@ -18,11 +18,8 @@ package org.apache.paimon.arrow.writer; -import org.apache.paimon.arrow.ArrowBundleRecords; -import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.arrow.vector.ArrowFormatCWriter; import org.apache.paimon.arrow.vector.ArrowFormatWriter; -import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.columnar.ColumnVector; import org.apache.paimon.data.columnar.VectorizedColumnBatch; import org.apache.paimon.data.columnar.heap.HeapArrayVector; @@ -31,13 +28,10 @@ import org.apache.paimon.data.columnar.heap.HeapMapVector; import org.apache.paimon.data.columnar.heap.HeapRowVector; import org.apache.paimon.fs.PositionOutputStream; -import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; -import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; -import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.Test; @@ -45,7 +39,6 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collections; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -53,181 +46,6 @@ /** Tests for {@link ArrowBundleWriter}. */ public class ArrowBundleWriterTest { - @Test - public void testArrowBundleFlushesBufferedRowsBeforeDirectWrite() throws Exception { - RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); - ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); - List events = new ArrayList<>(); - NativeWriter nativeWriter = - new NativeWriter() { - @Override - public long nativeMemoryUsed() { - return 0; - } - - @Override - public void writeIpcBytes(long arrayAddress, long schemaAddress) { - events.add("rows"); - cWriter.release(); - } - - @Override - public void close() {} - }; - ArrowBundleWriter writer = - new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { - @Override - public void add(VectorSchemaRoot vsr) { - events.add("bundle"); - } - }; - - writer.addElement(GenericRow.of(1)); - try (RootAllocator allocator = new RootAllocator(); - VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(rowType, allocator)) { - setInt((IntVector) root.getVector("value"), 2); - root.setRowCount(1); - - writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); - } - - assertThat(events).containsExactly("rows", "bundle"); - writer.close(); - } - - @Test - public void testReorderedArrowBundleFallsBackToRows() throws Exception { - RowType writerType = - RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); - RowType sourceType = - RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); - ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); - VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); - CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); - ArrowBundleWriter writer = - new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { - @Override - public void add(VectorSchemaRoot vsr) { - throw new AssertionError("Reordered Arrow bundle must use row fallback."); - } - }; - - try (RootAllocator allocator = new RootAllocator(); - VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(sourceType, allocator)) { - setInt((IntVector) root.getVector("b"), 20); - setInt((IntVector) root.getVector("a"), 10); - root.setRowCount(1); - - writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); - } - writer.close(); - - assertThat(nativeWriter.snapshots).hasSize(1); - assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); - assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); - } - - @Test - public void testLogicalRowTypeMismatchFallsBackToRows() throws Exception { - RowType writerType = RowType.builder().field("value", DataTypes.INT()).build(); - RowType bundleType = - new RowType( - Collections.singletonList( - new DataField( - 0, "value", DataTypes.INT(), "different description"))); - ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); - VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); - CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); - ArrowBundleWriter writer = - new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { - @Override - public void add(VectorSchemaRoot vsr) { - throw new AssertionError( - "Logically incompatible Arrow bundle must use row fallback."); - } - }; - - try (RootAllocator allocator = new RootAllocator(); - VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(writerType, allocator)) { - setInt((IntVector) root.getVector("value"), 10); - root.setRowCount(1); - - writer.writeBundle(new ArrowBundleRecords(root, bundleType, true)); - } - writer.close(); - - assertThat(nativeWriter.snapshots).hasSize(1); - assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); - } - - @Test - public void testNonIdentityNameMappingFallsBackToRows() throws Exception { - RowType writerType = - RowType.builder().field("A", DataTypes.INT()).field("a", DataTypes.INT()).build(); - ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); - VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); - CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); - ArrowBundleWriter writer = - new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { - @Override - public void add(VectorSchemaRoot vsr) { - throw new AssertionError( - "Non-identity Arrow name mapping must use row fallback."); - } - }; - - try (RootAllocator allocator = new RootAllocator(); - VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(writerType, allocator)) { - setInt((IntVector) root.getVector("A"), 10); - setInt((IntVector) root.getVector("a"), 20); - root.setRowCount(1); - - writer.writeBundle(new ArrowBundleRecords(root, writerType, false)); - } - writer.close(); - - assertThat(nativeWriter.snapshots).hasSize(1); - assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(20); - assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); - } - - @Test - public void testMixedAllocatorRootsFallBackToRows() throws Exception { - RowType rowType = - RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); - ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); - VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); - CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); - ArrowBundleWriter writer = - new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { - @Override - public void add(VectorSchemaRoot vsr) { - throw new AssertionError("Mixed-root Arrow bundle must use row fallback."); - } - }; - - try (RootAllocator firstAllocator = new RootAllocator(); - RootAllocator secondAllocator = new RootAllocator()) { - FieldVector firstVector = - writerRoot.getSchema().getFields().get(0).createVector(firstAllocator); - FieldVector secondVector = - writerRoot.getSchema().getFields().get(1).createVector(secondAllocator); - try (VectorSchemaRoot root = - new VectorSchemaRoot( - writerRoot.getSchema(), Arrays.asList(firstVector, secondVector), 1)) { - setInt((IntVector) firstVector, 10); - setInt((IntVector) secondVector, 20); - - writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); - } - } - writer.close(); - - assertThat(nativeWriter.snapshots).hasSize(1); - assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); - assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); - } - @Test public void testAddBatchWithoutDeletionVector() throws IOException { RowType rowType = RowType.of(DataTypes.INT(), DataTypes.BIGINT()); @@ -800,12 +618,6 @@ static class Snapshot { } } - private static void setInt(IntVector vector, int value) { - vector.allocateNew(1); - vector.setSafe(0, value); - vector.setValueCount(1); - } - private static class NoOpPositionOutputStream extends PositionOutputStream { private long pos = 0; diff --git a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java index 1ffabd45977d..c0cd962695be 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java @@ -18,27 +18,30 @@ package org.apache.paimon.format; -import org.apache.paimon.data.InternalRow; import org.apache.paimon.io.BundleRecords; import java.io.IOException; -/** Format write with bundle interface. */ +/** Format writer with bundle interface. */ public interface BundleFormatWriter extends FormatWriter { /** - * Writes a bundle with semantics equivalent to invoking {@link #addElement} for every record. - * - *

The implementation may consume the bundle natively, convert or copy it, or fall back to - * row-by-row writes. It must not retain borrowed buffers after this method returns unless it - * has copied them or acquired independent ownership. + * Writes a bundle of records. * * @param bundle the records to be written * @throws IOException if exception happens */ - default void writeBundle(BundleRecords bundle) throws IOException { - for (InternalRow row : bundle) { - addElement(row); - } + void writeBundle(BundleRecords bundle) throws IOException; + + /** + * Returns whether {@link #writeBundle} is equivalent to invoking {@link #addElement} for every + * record. + * + *

An implementation returning {@code true} must preserve record values and order. It must + * not retain borrowed buffers after {@link #writeBundle} returns unless it has copied them or + * acquired independent ownership. + */ + default boolean supportsRowEquivalentBundleWrite() { + return false; } } diff --git a/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java b/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java index dacc6f4687f7..37b6674e0107 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriter.java @@ -77,12 +77,7 @@ public void addElement(InternalRow row) throws IOException { public void writeBundle(BundleRecords bundle) throws IOException { if (!planFinalized) { for (InternalRow row : bundle) { - bufferedRows.add( - InternalRowUtils.copyInternalRow(row, writePlanFactory.logicalRowType())); - totalBufferedRowCount++; - } - if (totalBufferedRowCount >= writePlanFactory.inferBufferRowCount()) { - finalizePlanAndFlush(); + addElement(row); } return; } diff --git a/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java b/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java index 316825b29866..97a9d4d2da9d 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/shredding/ShreddingFormatWriter.java @@ -64,7 +64,9 @@ public void writeBundle(BundleRecords bundle) throws IOException { return; } - BundleFormatWriter.super.writeBundle(bundle); + for (InternalRow row : bundle) { + addElement(row); + } } @Override diff --git a/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java b/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java index bcf10caf456a..be7a8b197ea0 100644 --- a/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/format/shredding/InferShreddingWritePlanWriterTest.java @@ -21,6 +21,7 @@ import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; import org.apache.paimon.data.shredding.ShreddingWritePlan; +import org.apache.paimon.format.BundleFormatWriter; import org.apache.paimon.format.FormatWriter; import org.apache.paimon.fs.PositionOutputStream; import org.apache.paimon.io.BundleRecords; @@ -29,6 +30,7 @@ import org.junit.jupiter.api.Test; +import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.Iterator; @@ -56,8 +58,9 @@ void testMixedRowsAndBundlesPreserveOrderAndInferenceBoundary() throws Exception writer.writeBundle(bundle(GenericRow.of(6), GenericRow.of(7))); writer.close(); - assertThat(writePlanFactory.sampleValues).containsExactly(1, 2, 3, 4); + assertThat(writePlanFactory.sampleValues).containsExactly(1, 2, 3); assertThat(writerFactory.writer.values).containsExactly(101, 102, 103, 104, 105, 106, 107); + assertThat(writerFactory.writer.bundleWriteCount).isEqualTo(1); } @Test @@ -102,15 +105,24 @@ public FormatWriter createWithShreddingWritePlan( } } - private static class TestingFormatWriter implements FormatWriter { + private static class TestingFormatWriter implements BundleFormatWriter { private final List values = new ArrayList<>(); + private int bundleWriteCount; @Override public void addElement(InternalRow element) { values.add(element.getInt(0)); } + @Override + public void writeBundle(BundleRecords bundle) throws IOException { + bundleWriteCount++; + for (InternalRow row : bundle) { + addElement(row); + } + } + @Override public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { return false; diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java index 618733b99718..fe35d51fdd5f 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java @@ -110,11 +110,10 @@ public void writeBundle(BundleRecords bundle) throws IOException { openCurrentWriter(); } - long previousRecordCount = currentWriter.recordCount(); + long rowCount = bundle.rowCount(); currentWriter.writeBundle(bundle); - long writtenRecordCount = currentWriter.recordCount() - previousRecordCount; - recordCount += writtenRecordCount; - currentFileRecordCount += writtenRecordCount; + recordCount += rowCount; + currentFileRecordCount += rowCount; if (rollingFile(true)) { closeCurrentWriter(); diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java index 42a885849e94..32db7f415a8d 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java @@ -118,10 +118,11 @@ public void write(InternalRow row) throws IOException { public void writeBundle(BundleRecords bundle) throws IOException { if (auxiliaryFileWriters.isEmpty() && sequenceNumberTracker.supportsRowCountUpdate() - && !requiresPerRecordStats()) { - long previousRecordCount = recordCount(); + && !requiresPerRecordStats() + && supportsRowEquivalentBundleWrite()) { + long rowCount = bundle.rowCount(); super.writeBundle(bundle); - sequenceNumberTracker.updateByRowCount(recordCount() - previousRecordCount); + sequenceNumberTracker.updateByRowCount(rowCount); return; } diff --git a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java index 96f273091dd7..04720e503291 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java @@ -164,6 +164,11 @@ public void writeBundle(BundleRecords bundle) throws IOException { } } + protected final boolean supportsRowEquivalentBundleWrite() { + return writer instanceof BundleFormatWriter + && ((BundleFormatWriter) writer).supportsRowEquivalentBundleWrite(); + } + protected InternalRow writeImpl(T record) throws IOException { if (closed) { throw new RuntimeException("Writer has already closed!"); diff --git a/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java index 33b04908560e..5eb6fd8a1d31 100644 --- a/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/io/RollingFileWriterTest.java @@ -171,21 +171,6 @@ public void testRollingByRowsWithBundle() throws IOException { assertThat(files.get(2).rowCount()).isEqualTo(30); } - @Test - public void testRollingWriterDoesNotReadBundleRowCountAgain() throws IOException { - initialize("parquet", false, 1024L * 1024 * 1024, 100L); - SingleUseBundleRecords bundle = bundle(150); - - rollingFileWriter.writeBundle(bundle); - rollingFileWriter.close(); - - assertThat(bundle.rowCountCalls).isEqualTo(1); - assertThat(rollingFileWriter.result()) - .singleElement() - .extracting(DataFileMeta::rowCount) - .isEqualTo(150L); - } - private static SingleUseBundleRecords bundle(int rowCount) { List rows = new ArrayList<>(); for (int i = 0; i < rowCount; i++) { @@ -331,7 +316,6 @@ private static class SingleUseBundleRecords implements BundleRecords { private final List rows; private boolean iterated; - private int rowCountCalls; private SingleUseBundleRecords(List rows) { this.rows = rows; @@ -348,7 +332,6 @@ public Iterator iterator() { @Override public long rowCount() { - rowCountCalls++; return rows.size(); } } diff --git a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java index 57e7aa8a65e4..706bc18f9acb 100644 --- a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java @@ -74,7 +74,6 @@ void testEligibleBundleIsForwardedWithoutIteration() throws Exception { writer.writeBundle(bundle); - assertThat(bundle.rowCountCalls).isEqualTo(1); assertThat(bundle.iteratorCalls).isZero(); assertThat(formatWriter.writtenBundle).isSameAs(bundle); assertThat(formatWriter.bundleWrites).isEqualTo(1); @@ -107,7 +106,6 @@ void testExtractorStatsAllowBundleForwarding() throws Exception { writer.writeBundle(bundle); - assertThat(bundle.rowCountCalls).isEqualTo(1); assertThat(bundle.iteratorCalls).isZero(); assertThat(formatWriter.writtenBundle).isSameAs(bundle); assertThat(formatWriter.bundleWrites).isEqualTo(1); @@ -144,7 +142,7 @@ void testPlainFormatWriterFallsBackToRows() throws Exception { } @Test - void testBundleFormatWriterCanChooseRowFallback() throws Exception { + void testBundleFormatWriterWithoutOptInFallsBackToRows() throws Exception { TestingFallbackBundleFormatWriter formatWriter = new TestingFallbackBundleFormatWriter(); LongCounter sequenceCounter = new LongCounter(); RowDataFileWriter writer = @@ -158,6 +156,7 @@ void testBundleFormatWriterCanChooseRowFallback() throws Exception { writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); + assertThat(formatWriter.bundleWrites).isZero(); assertThat(formatWriter.rowWrites).isEqualTo(2); assertThat(writer.recordCount()).isEqualTo(2); assertThat(sequenceCounter.getValue()).isEqualTo(2); @@ -379,10 +378,26 @@ public void writeBundle(BundleRecords bundle) { bundleWrites++; writtenBundle = bundle; } + + @Override + public boolean supportsRowEquivalentBundleWrite() { + return true; + } } private static class TestingFallbackBundleFormatWriter extends TestingFormatWriter - implements BundleFormatWriter {} + implements BundleFormatWriter { + + private int bundleWrites; + + @Override + public void writeBundle(BundleRecords bundle) { + bundleWrites++; + for (InternalRow row : bundle) { + addElement(row); + } + } + } private static class TestingThrowingBundleFormatWriter extends TestingFormatWriter implements BundleFormatWriter { @@ -397,6 +412,11 @@ private TestingThrowingBundleFormatWriter(IOException failure) { public void writeBundle(BundleRecords bundle) throws IOException { throw failure; } + + @Override + public boolean supportsRowEquivalentBundleWrite() { + return true; + } } private static class TestingStatsProducer implements SimpleStatsProducer { @@ -473,7 +493,6 @@ private static class TrackingBundleRecords implements BundleRecords { private final List rows; private int iteratorCalls; - private int rowCountCalls; private TrackingBundleRecords(List rows) { this.rows = rows; @@ -487,7 +506,6 @@ public Iterator iterator() { @Override public long rowCount() { - rowCountCalls++; return rows.size(); } } diff --git a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java index d315f44633ec..ee1f572462d7 100644 --- a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java +++ b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java @@ -19,7 +19,6 @@ package org.apache.paimon.format.lance; import org.apache.paimon.arrow.ArrowBundleRecords; -import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.arrow.vector.ArrowFormatWriter; import org.apache.paimon.data.InternalRow; import org.apache.paimon.format.BundleFormatWriter; @@ -66,18 +65,12 @@ public void addElement(InternalRow internalRow) throws IOException { @Override public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { - ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; - VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot(); - if (arrowFormatWriter.isArrowBundleSchemaCompatible(arrowBundle) - && ArrowUtils.hasSameRootAllocator(root, arrowFormatWriter.getAllocator())) { - flush(); - nativeWriter.ensureInitialized(arrowFormatWriter.getAllocator()); - add(root); - return; + add(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); + } else { + for (InternalRow row : bundleRecords) { + addElement(row); } } - - BundleFormatWriter.super.writeBundle(bundleRecords); } public void add(VectorSchemaRoot vsr) throws IOException { diff --git a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java index b3e29eb0a0b3..f9329907f951 100644 --- a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java +++ b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java @@ -18,8 +18,6 @@ package org.apache.paimon.format.lance.jni; -import org.apache.paimon.arrow.ArrowUtils; - import com.lancedb.lance.file.LanceFileWriter; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; @@ -34,7 +32,6 @@ public class LanceWriter { private final String path; private final Map storageOptions; private LanceFileWriter writer; - private BufferAllocator allocator; private long bytesWritten = 0; public LanceWriter(String path, Map storageOptions) { @@ -47,25 +44,12 @@ public Long getWrittenPosition() { } public void writeVsr(VectorSchemaRoot vsr) throws IOException { - BufferAllocator sourceAllocator = vsr.getVector(0).getAllocator(); - initWriteLazy(sourceAllocator); - if (!ArrowUtils.hasSameRootAllocator(vsr, allocator)) { - throw new IllegalArgumentException( - "Lance writer cannot consume Arrow buffers from a different allocator root."); - } + initWriteLazy(vsr.getVector(0).getAllocator()); this.bytesWritten += vsr.getFieldVectors().stream().mapToLong(FieldVector::getBufferSize).sum(); this.writer.write(vsr); } - /** - * Initializes the native writer with an allocator whose lifetime is owned by the surrounding - * format writer. - */ - public void ensureInitialized(BufferAllocator bufferAllocator) throws IOException { - initWriteLazy(bufferAllocator); - } - public void close() throws IOException { if (writer != null) { try { @@ -74,7 +58,6 @@ public void close() throws IOException { throw new IOException(e); } this.writer = null; - this.allocator = null; } } @@ -85,7 +68,6 @@ public String path() { private void initWriteLazy(BufferAllocator bufferAllocator) throws IOException { if (writer == null) { writer = LanceFileWriter.open(path, bufferAllocator, null, storageOptions); - allocator = bufferAllocator; } } } diff --git a/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java deleted file mode 100644 index efd3e2a821e0..000000000000 --- a/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java +++ /dev/null @@ -1,189 +0,0 @@ -/* - * 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.format.lance; - -import org.apache.paimon.arrow.ArrowBundleRecords; -import org.apache.paimon.arrow.ArrowUtils; -import org.apache.paimon.arrow.vector.ArrowFormatWriter; -import org.apache.paimon.data.GenericRow; -import org.apache.paimon.format.lance.jni.LanceWriter; -import org.apache.paimon.types.DataTypes; -import org.apache.paimon.types.RowType; - -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.vector.IntVector; -import org.apache.arrow.vector.VectorSchemaRoot; -import org.junit.jupiter.api.Test; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import java.util.stream.Collectors; - -import static org.assertj.core.api.Assertions.assertThat; - -/** Tests for {@link LanceRecordsWriter}. */ -class LanceRecordsWriterTest { - - @Test - void testArrowBundlePreservesRowBundleRowOrder() throws Exception { - RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); - ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true); - BufferAllocator writerAllocator = arrowWriter.getAllocator(); - CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); - LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); - - writer.addElement(GenericRow.of(1)); - try (BufferAllocator sourceAllocator = - arrowWriter - .getAllocator() - .newChildAllocator("lance-bundle-test", 0, Long.MAX_VALUE); - VectorSchemaRoot root = - ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { - setInt((IntVector) root.getVector("value"), 2); - root.setRowCount(1); - writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); - } - writer.addElement(GenericRow.of(3)); - writer.close(); - - assertThat(nativeWriter.snapshots).hasSize(3); - assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(1); - assertThat(nativeWriter.snapshots.get(1).values.get(0)).containsExactly(2); - assertThat(nativeWriter.snapshots.get(2).values.get(0)).containsExactly(3); - assertThat(nativeWriter.initializedAllocator).isSameAs(writerAllocator); - } - - @Test - void testReorderedArrowBundleFallsBackToRows() throws Exception { - RowType writerType = - RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); - RowType sourceType = - RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); - ArrowFormatWriter arrowWriter = new ArrowFormatWriter(writerType, 1024, true); - CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); - LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); - - try (BufferAllocator sourceAllocator = - arrowWriter - .getAllocator() - .newChildAllocator("lance-schema-test", 0, Long.MAX_VALUE); - VectorSchemaRoot root = - ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { - setInt((IntVector) root.getVector("b"), 20); - setInt((IntVector) root.getVector("a"), 10); - root.setRowCount(1); - writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); - } - writer.close(); - - assertThat(nativeWriter.snapshots).hasSize(1); - Snapshot snapshot = nativeWriter.snapshots.get(0); - assertThat(snapshot.fieldNames).containsExactly("a", "b"); - assertThat(snapshot.values.get(0)).containsExactly(10); - assertThat(snapshot.values.get(1)).containsExactly(20); - } - - @Test - void testDifferentAllocatorRootFallsBackToRows() throws Exception { - RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); - ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true); - CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); - LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); - - try (BufferAllocator sourceAllocator = new org.apache.arrow.memory.RootAllocator(); - VectorSchemaRoot root = - ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { - nativeWriter.disallowedRoot = root; - setInt((IntVector) root.getVector("value"), 10); - root.setRowCount(1); - - writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); - } - writer.close(); - - assertThat(nativeWriter.disallowedRootWrites).isZero(); - assertThat(nativeWriter.snapshots).hasSize(1); - assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(10); - } - - private static void setInt(IntVector vector, int value) { - vector.allocateNew(1); - vector.setSafe(0, value); - vector.setValueCount(1); - } - - private static class CapturingLanceWriter extends LanceWriter { - - private final List snapshots = new ArrayList<>(); - private BufferAllocator initializedAllocator; - private VectorSchemaRoot disallowedRoot; - private int disallowedRootWrites; - - private CapturingLanceWriter() { - super("unused", Collections.emptyMap()); - } - - @Override - public void ensureInitialized(BufferAllocator bufferAllocator) { - initializedAllocator = bufferAllocator; - } - - @Override - public void writeVsr(VectorSchemaRoot root) { - if (root == disallowedRoot) { - disallowedRootWrites++; - } - List fieldNames = - root.getSchema().getFields().stream() - .map(field -> field.getName()) - .collect(Collectors.toList()); - List> values = new ArrayList<>(); - for (int column = 0; column < root.getFieldVectors().size(); column++) { - IntVector vector = (IntVector) root.getVector(column); - List columnValues = new ArrayList<>(); - for (int row = 0; row < root.getRowCount(); row++) { - columnValues.add(vector.get(row)); - } - values.add(columnValues); - } - snapshots.add(new Snapshot(fieldNames, values)); - } - - @Override - public void close() throws IOException {} - - @Override - public String path() { - return "unused"; - } - } - - private static class Snapshot { - - private final List fieldNames; - private final List> values; - - private Snapshot(List fieldNames, List> values) { - this.fieldNames = fieldNames; - this.values = values; - } - } -} 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 3e12e626e671..86300a820d5a 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 @@ -128,6 +128,8 @@ public void writeBundle(BundleRecords bundleRecords) { if (bundleRecords instanceof ArrowBundleRecords) { ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot(); + // Mosaic exports the borrowed vectors through the writer allocator, so direct writes + // require every source vector to share its root; otherwise preserve semantics via rows. if (arrowFormatWriter.isArrowBundleSchemaCompatible(arrowBundle) && ArrowUtils.hasSameRootAllocator(root, allocator)) { flush(); @@ -141,6 +143,11 @@ public void writeBundle(BundleRecords bundleRecords) { } } + @Override + public boolean supportsRowEquivalentBundleWrite() { + return true; + } + @Override public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { if (!suggestedCheck) { diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java deleted file mode 100644 index 12a8f3cb0c76..000000000000 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicBundleWriteIntegrationTest.java +++ /dev/null @@ -1,338 +0,0 @@ -/* - * 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.format.mosaic; - -import org.apache.paimon.arrow.ArrowBundleRecords; -import org.apache.paimon.arrow.ArrowUtils; -import org.apache.paimon.data.InternalRow; -import org.apache.paimon.fileindex.FileIndexOptions; -import org.apache.paimon.format.FileFormatFactory; -import org.apache.paimon.format.FormatReaderContext; -import org.apache.paimon.format.FormatReaderFactory; -import org.apache.paimon.format.FormatWriter; -import org.apache.paimon.format.FormatWriterFactory; -import org.apache.paimon.fs.Path; -import org.apache.paimon.fs.PositionOutputStream; -import org.apache.paimon.fs.local.LocalFileIO; -import org.apache.paimon.io.DataFileMeta; -import org.apache.paimon.io.FileWriterContext; -import org.apache.paimon.io.RowDataFileWriter; -import org.apache.paimon.io.SimpleStatsProducer; -import org.apache.paimon.manifest.FileSource; -import org.apache.paimon.mosaic.MosaicWriter; -import org.apache.paimon.mosaic.WriterOptions; -import org.apache.paimon.options.Options; -import org.apache.paimon.reader.RecordReader; -import org.apache.paimon.types.DataTypes; -import org.apache.paimon.types.RowType; -import org.apache.paimon.utils.LongCounter; - -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.vector.IntVector; -import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.types.pojo.Schema; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.io.TempDir; - -import java.io.IOException; -import java.io.OutputStream; -import java.util.ArrayList; -import java.util.Collections; -import java.util.Iterator; -import java.util.List; -import java.util.concurrent.atomic.AtomicReference; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.junit.jupiter.api.Assumptions.assumeTrue; - -/** Integration tests for bundle dispatch from core file writers to Mosaic native writes. */ -class MosaicBundleWriteIntegrationTest { - - private static final FileFormatFactory.FormatContext FORMAT_CONTEXT = - new FileFormatFactory.FormatContext(new Options(), 1024, 1024); - - @TempDir java.nio.file.Path tempDir; - - @Test - void testCompatibleArrowBundleUsesFullDirectWritePath() throws Exception { - assumeTrue(isNativeAvailable(), "Mosaic native library not available"); - - RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); - Path path = newPath("direct"); - LocalFileIO fileIO = new LocalFileIO(); - RootAllocator writerAllocator = new RootAllocator(); - LongCounter sequenceCounter = new LongCounter(5); - AtomicReference nativeWriterRef = new AtomicReference<>(); - - try (RowDataFileWriter writer = - createWriter( - fileIO, path, rowType, writerAllocator, sequenceCounter, nativeWriterRef)) { - TrackingMosaicWriter nativeWriter = nativeWriterRef.get(); - assertThat(nativeWriter).isNotNull(); - - try (BufferAllocator sourceAllocator = - writerAllocator.newChildAllocator( - "mosaic-direct-integration-test", 0, Long.MAX_VALUE); - VectorSchemaRoot root = - ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { - setInts((IntVector) root.getVector("value"), 1, 2, 3); - root.setRowCount(3); - - TrackingDirectArrowBundleRecords bundle = - new TrackingDirectArrowBundleRecords(root, rowType); - nativeWriter.expectDirectRoot(root); - writer.writeBundle(bundle); - nativeWriter.clearExpectedRoot(); - - assertThat(bundle.iteratorCalls).isZero(); - assertThat(nativeWriter.directWrites).isEqualTo(1); - assertThat(writer.recordCount()).isEqualTo(3); - assertThat(sequenceCounter.getValue()).isEqualTo(8); - } - - // The borrowed source root has already been released. Closing the native writer must - // not access it again. - writer.close(); - DataFileMeta result = writer.result(); - assertThat(result.rowCount()).isEqualTo(3); - assertThat(result.minSequenceNumber()).isEqualTo(5); - assertThat(result.maxSequenceNumber()).isEqualTo(7); - } - - assertThat(readRows(fileIO, path, rowType)) - .containsExactly( - Collections.singletonList(1), - Collections.singletonList(2), - Collections.singletonList(3)); - } - - @Test - void testIncompatibleArrowSchemaFallsBackThroughFullWritePath() throws Exception { - assumeTrue(isNativeAvailable(), "Mosaic native library not available"); - - RowType writerType = - RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); - RowType sourceType = - RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); - Path path = newPath("fallback"); - LocalFileIO fileIO = new LocalFileIO(); - RootAllocator writerAllocator = new RootAllocator(); - LongCounter sequenceCounter = new LongCounter(); - AtomicReference nativeWriterRef = new AtomicReference<>(); - - try (RowDataFileWriter writer = - createWriter( - fileIO, - path, - writerType, - writerAllocator, - sequenceCounter, - nativeWriterRef)) { - TrackingMosaicWriter nativeWriter = nativeWriterRef.get(); - assertThat(nativeWriter).isNotNull(); - - try (BufferAllocator sourceAllocator = - writerAllocator.newChildAllocator( - "mosaic-fallback-integration-test", 0, Long.MAX_VALUE); - VectorSchemaRoot root = - ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { - setInts((IntVector) root.getVector("b"), 20, 21); - setInts((IntVector) root.getVector("a"), 10, 11); - root.setRowCount(2); - - TrackingArrowBundleRecords bundle = - new TrackingArrowBundleRecords(root, writerType); - nativeWriter.expectDirectRoot(root); - writer.writeBundle(bundle); - nativeWriter.clearExpectedRoot(); - - assertThat(bundle.iteratorCalls).isEqualTo(1); - assertThat(nativeWriter.directWrites).isZero(); - assertThat(writer.recordCount()).isEqualTo(2); - assertThat(sequenceCounter.getValue()).isEqualTo(2); - } - - writer.close(); - assertThat(writer.result().rowCount()).isEqualTo(2); - } - - assertThat(readRows(fileIO, path, writerType)) - .containsExactly(asList(10, 20), asList(11, 21)); - } - - private Path newPath(String prefix) { - return new Path(tempDir.toUri().toString(), prefix + ".mosaic"); - } - - private static RowDataFileWriter createWriter( - LocalFileIO fileIO, - Path path, - RowType rowType, - RootAllocator allocator, - LongCounter sequenceCounter, - AtomicReference nativeWriterRef) { - FormatWriterFactory writerFactory = - new FormatWriterFactory() { - @Override - public FormatWriter create(PositionOutputStream out, String compression) { - assertThat(compression).isEqualTo("zstd"); - return new MosaicRecordsWriter( - out, - rowType, - FORMAT_CONTEXT, - Collections.emptyList(), - null, - allocator, - (outputStream, arrowSchema, options, bufferAllocator) -> { - TrackingMosaicWriter writer = - new TrackingMosaicWriter( - outputStream, - arrowSchema, - options, - bufferAllocator); - nativeWriterRef.set(writer); - return writer; - }); - } - }; - - return new RowDataFileWriter( - fileIO, - new FileWriterContext( - writerFactory, SimpleStatsProducer.disabledProducer(), "zstd"), - path, - rowType, - 1L, - () -> sequenceCounter, - new FileIndexOptions(), - FileSource.APPEND, - false, - false, - false, - null, - null, - null); - } - - private static List> readRows(LocalFileIO fileIO, Path path, RowType rowType) - throws IOException { - MosaicFileFormat format = new MosaicFileFormat(FORMAT_CONTEXT); - FormatReaderFactory readerFactory = - format.createReaderFactory(rowType, rowType, Collections.emptyList()); - List> rows = new ArrayList<>(); - try (RecordReader reader = - readerFactory.createReader( - new FormatReaderContext(fileIO, path, fileIO.getFileSize(path)))) { - reader.forEachRemaining( - row -> { - List values = new ArrayList<>(rowType.getFieldCount()); - for (int i = 0; i < rowType.getFieldCount(); i++) { - values.add(row.getInt(i)); - } - rows.add(values); - }); - } - return rows; - } - - private static List asList(int first, int second) { - List values = new ArrayList<>(2); - values.add(first); - values.add(second); - return values; - } - - private static void setInts(IntVector vector, int... values) { - vector.allocateNew(values.length); - for (int i = 0; i < values.length; i++) { - vector.setSafe(i, values[i]); - } - vector.setValueCount(values.length); - } - - private static boolean isNativeAvailable() { - try { - Class.forName("org.apache.paimon.mosaic.NativeLib"); - return true; - } catch (Throwable t) { - return false; - } - } - - private static class TrackingDirectArrowBundleRecords extends ArrowBundleRecords { - - private int iteratorCalls; - - private TrackingDirectArrowBundleRecords(VectorSchemaRoot root, RowType rowType) { - super(root, rowType, true); - } - - @Override - public Iterator iterator() { - iteratorCalls++; - return super.iterator(); - } - } - - private static class TrackingArrowBundleRecords extends ArrowBundleRecords { - - private int iteratorCalls; - - private TrackingArrowBundleRecords(VectorSchemaRoot root, RowType rowType) { - super(root, rowType, true); - } - - @Override - public Iterator iterator() { - iteratorCalls++; - return super.iterator(); - } - } - - private static class TrackingMosaicWriter extends MosaicWriter { - - private VectorSchemaRoot expectedDirectRoot; - private int directWrites; - - private TrackingMosaicWriter( - OutputStream outputStream, - Schema schema, - WriterOptions options, - BufferAllocator allocator) { - super(outputStream, schema, options, allocator); - } - - private void expectDirectRoot(VectorSchemaRoot root) { - expectedDirectRoot = root; - } - - private void clearExpectedRoot() { - expectedDirectRoot = null; - } - - @Override - public void write(VectorSchemaRoot root) { - if (root == expectedDirectRoot) { - directWrites++; - } - super.write(root); - } - } -} diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java index 60efceed08e3..931f4d703318 100644 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java @@ -18,10 +18,13 @@ package org.apache.paimon.format.mosaic; +import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; import org.apache.paimon.data.serializer.InternalRowSerializer; +import org.apache.paimon.format.BundleFormatWriter; import org.apache.paimon.format.FileFormatFactory; import org.apache.paimon.format.FormatReaderContext; import org.apache.paimon.format.FormatReaderFactory; @@ -37,6 +40,9 @@ 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.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -44,6 +50,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; +import java.util.Iterator; import java.util.List; import java.util.UUID; @@ -80,6 +87,38 @@ void testWriteAndRead() throws IOException { assertThat(result.get(1).getString(1).toString()).isEqualTo("world"); } + @Test + void testPublicWriterFactoryFallsBackForCrossRootArrowBundle() throws IOException { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + Path path = newPath(); + MosaicFileFormat format = createFormat(); + FormatWriterFactory writerFactory = format.createWriterFactory(rowType); + LocalFileIO fileIO = new LocalFileIO(); + + try (RootAllocator sourceAllocator = new RootAllocator(); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator); + BundleFormatWriter writer = + (BundleFormatWriter) + writerFactory.create(fileIO.newOutputStream(path, false), "zstd")) { + IntVector vector = (IntVector) root.getVector("value"); + vector.allocateNew(2); + vector.setSafe(0, 7); + vector.setSafe(1, 9); + vector.setValueCount(2); + root.setRowCount(2); + + CountingArrowBundleRecords bundle = new CountingArrowBundleRecords(root, rowType); + writer.writeBundle(bundle); + assertThat(bundle.iteratorCalls()).isEqualTo(1); + } + + List result = readAll(rowType, rowType, path, null); + assertThat(result).hasSize(2); + assertThat(result.get(0).getInt(0)).isEqualTo(7); + assertThat(result.get(1).getInt(0)).isEqualTo(9); + } + @Test void testNullValues() throws IOException { RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING()); @@ -358,4 +397,23 @@ private static boolean isNativeAvailable() { return false; } } + + private static class CountingArrowBundleRecords extends ArrowBundleRecords { + + private int iteratorCalls; + + private CountingArrowBundleRecords(VectorSchemaRoot root, RowType rowType) { + super(root, rowType, true); + } + + @Override + public Iterator iterator() { + iteratorCalls++; + return super.iterator(); + } + + private int iteratorCalls() { + return iteratorCalls; + } + } } 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 b1435b30b5bb..344d5a0022b6 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 @@ -140,7 +140,8 @@ void testReorderedArrowBundleFallsBackToRows() throws Exception { .write(any(VectorSchemaRoot.class)); try (BufferAllocator sourceAllocator = - writerAllocator.newChildAllocator("mosaic-schema-test", 0, Long.MAX_VALUE); + writerAllocator.newChildAllocator( + "mosaic-reordered-test", 0, Long.MAX_VALUE); VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { setInt((IntVector) root.getVector("b"), 20); diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java index 7d83f5698906..7798fa09d419 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java @@ -102,15 +102,13 @@ public void addElement(InternalRow internalRow) throws IOException { @Override public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { - ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; - if (currentWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle)) { - flush(); - writeBundleVsr(arrowBundle.getVectorSchemaRoot()); - return; + flush(); + writeBundleVsr(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); + } else { + for (InternalRow row : bundleRecords) { + addElement(row); } } - - BundleFormatWriter.super.writeBundle(bundleRecords); } @Override diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java index b90e601610f7..e97e3cf80fcc 100644 --- a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java @@ -19,7 +19,6 @@ package org.apache.paimon.format.vortex; 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.GenericRow; @@ -46,9 +45,6 @@ import org.apache.paimon.utils.RoaringBitmap32; import dev.vortex.jni.NativeRuntime; -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 org.junit.jupiter.api.io.TempDir; @@ -324,55 +320,6 @@ public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( } } - @Test - public void testReorderedArrowBundleFallsBackToRows(@TempDir java.nio.file.Path tempDir) - throws Exception { - RowType writerType = - RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); - RowType sourceType = - RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); - - Options options = new Options(); - VortexFileFormat format = - new VortexFileFormatFactory() - .create(new FileFormatFactory.FormatContext(options, 1024, 1024)); - FileIO fileIO = new LocalFileIO(); - Path testFile = - new Path(new Path(tempDir.toUri()), "test_reordered_bundle_" + UUID.randomUUID()); - - try (FormatWriter writer = - ((SupportsDirectWrite) format.createWriterFactory(writerType)) - .create(fileIO, testFile, ""); - RootAllocator sourceAllocator = new RootAllocator(); - VectorSchemaRoot root = - ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { - setInt((IntVector) root.getVector("b"), 20); - setInt((IntVector) root.getVector("a"), 10); - root.setRowCount(1); - - ((BundleFormatWriter) writer) - .writeBundle(new ArrowBundleRecords(root, writerType, true)); - } - - InternalRowSerializer serializer = new InternalRowSerializer(writerType); - FormatReaderFactory readerFactory = - format.createReaderFactory(writerType, writerType, null); - try (RecordReader reader = - readerFactory.createReader( - new FormatReaderContext( - fileIO, testFile, fileIO.getFileSize(testFile), null)); - RecordReaderIterator iterator = new RecordReaderIterator<>(reader)) { - List actualRows = new ArrayList<>(); - while (iterator.hasNext()) { - actualRows.add(serializer.copy(iterator.next())); - } - - assertEquals(1, actualRows.size()); - assertEquals(10, actualRows.get(0).getInt(0)); - assertEquals(20, actualRows.get(0).getInt(1)); - } - } - @Test public void testReadWithSelection(@TempDir java.nio.file.Path tempDir) throws Exception { RowType rowType = RowType.of(DataTypes.INT(), DataTypes.STRING()); @@ -420,12 +367,6 @@ public void testReadWithSelection(@TempDir java.nio.file.Path tempDir) throws Ex } } - private static void setInt(IntVector vector, int value) { - vector.allocateNew(1); - vector.setSafe(0, value); - vector.setValueCount(1); - } - @Test public void testReadWithVirtualRowTrackingField(@TempDir java.nio.file.Path tempDir) throws Exception { From abc1e0b5bca362866b3fb7e1f77ab75e7ebdd05b Mon Sep 17 00:00:00 2001 From: mingfeng Date: Mon, 10 Aug 2026 00:54:53 -0700 Subject: [PATCH 5/6] [core][format] Restore safe bundle writer forwarding --- .../org/apache/paimon/arrow/ArrowUtils.java | 13 + .../arrow/writer/ArrowBundleWriter.java | 26 +- .../arrow/writer/ArrowBundleWriterTest.java | 293 ++++++++++++++++- .../paimon/format/BundleFormatWriter.java | 23 +- .../apache/paimon/io/RowDataFileWriter.java | 3 +- .../apache/paimon/io/SingleFileWriter.java | 5 - .../paimon/io/RowDataFileWriterTest.java | 69 +++- .../format/lance/LanceRecordsWriter.java | 17 +- .../paimon/format/lance/jni/LanceWriter.java | 20 +- .../format/lance/LanceRecordsWriterTest.java | 297 ++++++++++++++++++ .../format/mosaic/MosaicRecordsWriter.java | 5 - .../format/vortex/VortexRecordsWriter.java | 89 ++++-- .../format/vortex/VortexReaderWriterTest.java | 127 +++++++- 13 files changed, 902 insertions(+), 85 deletions(-) create mode 100644 paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java index 2122e00a3ba1..8013d9ec8f82 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java @@ -277,6 +277,19 @@ 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); diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java index 39a006b6d7ed..f8c7a22bb74c 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/writer/ArrowBundleWriter.java +++ b/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); diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java index a43ec4e0b65e..e7a82d14484c 100644 --- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java +++ b/paimon-arrow/src/test/java/org/apache/paimon/arrow/writer/ArrowBundleWriterTest.java @@ -18,8 +18,11 @@ package org.apache.paimon.arrow.writer; +import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.arrow.vector.ArrowFormatCWriter; import org.apache.paimon.arrow.vector.ArrowFormatWriter; +import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.columnar.ColumnVector; import org.apache.paimon.data.columnar.VectorizedColumnBatch; import org.apache.paimon.data.columnar.heap.HeapArrayVector; @@ -28,10 +31,14 @@ import org.apache.paimon.data.columnar.heap.HeapMapVector; import org.apache.paimon.data.columnar.heap.HeapRowVector; import org.apache.paimon.fs.PositionOutputStream; +import org.apache.paimon.io.VectorizedBundleRecords; +import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.Test; @@ -39,13 +46,289 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link ArrowBundleWriter}. */ public class ArrowBundleWriterTest { + @Test + public void testArrowBundleFlushesBufferedRowsBeforeDirectWrite() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); + List events = new ArrayList<>(); + NativeWriter nativeWriter = + new NativeWriter() { + @Override + public long nativeMemoryUsed() { + return 0; + } + + @Override + public void writeIpcBytes(long arrayAddress, long schemaAddress) { + events.add("rows"); + cWriter.release(); + } + + @Override + public void close() {} + }; + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + events.add("bundle"); + } + }; + + writer.addElement(GenericRow.of(1)); + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(rowType, allocator)) { + setInt((IntVector) root.getVector("value"), 2); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + + assertThat(events).containsExactly("rows", "bundle"); + writer.close(); + } + + @Test + public void testReorderedArrowBundleFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError("Reordered Arrow bundle must use row fallback."); + } + }; + + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(sourceType, allocator)) { + setInt((IntVector) root.getVector("b"), 20); + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); + } + + @Test + public void testMissingColumnArrowBundleFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = RowType.builder().field("a", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError( + "Arrow bundle with a missing column must use row fallback."); + } + }; + + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(sourceType, allocator)) { + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)) + .containsExactly((Object) null); + } + + @Test + public void testDescriptionOnlyRowTypeDifferenceUsesDirectWrite() throws Exception { + RowType writerType = RowType.builder().field("value", DataTypes.INT()).build(); + RowType bundleType = + new RowType( + Collections.singletonList( + new DataField( + 0, "value", DataTypes.INT(), "different description"))); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + int[] directWrites = {0}; + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + directWrites[0]++; + } + }; + + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(writerType, allocator)) { + setInt((IntVector) root.getVector("value"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, bundleType, true)); + } + writer.close(); + + assertThat(directWrites[0]).isEqualTo(1); + assertThat(nativeWriter.snapshots).isEmpty(); + } + + @Test + public void testNonIdentityNameMappingFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("A", DataTypes.INT()).field("a", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(writerType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError( + "Non-identity Arrow name mapping must use row fallback."); + } + }; + + try (RootAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = ArrowUtils.createVectorSchemaRoot(writerType, allocator)) { + setInt((IntVector) root.getVector("A"), 10); + setInt((IntVector) root.getVector("a"), 20); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, writerType, false)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(20); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); + } + + @Test + public void testMixedAllocatorRootsFallBackToRows() throws Exception { + RowType rowType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); + VectorSchemaRoot writerRoot = cWriter.getVectorSchemaRoot(); + CapturingNativeWriter nativeWriter = new CapturingNativeWriter(writerRoot, cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter) { + @Override + public void add(VectorSchemaRoot vsr) { + throw new AssertionError("Mixed-root Arrow bundle must use row fallback."); + } + }; + + try (RootAllocator firstAllocator = new RootAllocator(); + RootAllocator secondAllocator = new RootAllocator()) { + FieldVector firstVector = + writerRoot.getSchema().getFields().get(0).createVector(firstAllocator); + FieldVector secondVector = + writerRoot.getSchema().getFields().get(1).createVector(secondAllocator); + try (VectorSchemaRoot root = + new VectorSchemaRoot( + writerRoot.getSchema(), Arrays.asList(firstVector, secondVector), 1)) { + setInt((IntVector) firstVector, 10); + setInt((IntVector) secondVector, 20); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(10); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(1)).containsExactly(20); + } + + @Test + public void testDirectWriteFailureReleasesExportedReferences() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); + RuntimeException failure = new RuntimeException("Expected native write failure."); + NativeWriter nativeWriter = + new NativeWriter() { + @Override + public long nativeMemoryUsed() { + return 0; + } + + @Override + public void writeIpcBytes(long arrayAddress, long schemaAddress) { + throw failure; + } + + @Override + public void close() {} + }; + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter); + + try { + try (RootAllocator allocator = new RootAllocator()) { + try (VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, allocator)) { + setInt((IntVector) root.getVector("value"), 10); + root.setRowCount(1); + + assertThatThrownBy( + () -> + writer.writeBundle( + new ArrowBundleRecords(root, rowType, true))) + .isSameAs(failure); + } + + assertThat(allocator.getAllocatedMemory()).isZero(); + } + } finally { + writer.close(); + } + } + + @Test + public void testVectorizedBundleFlushesBufferedRowsBeforeWrite() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatCWriter cWriter = new ArrowFormatCWriter(rowType, 1024, true); + CapturingNativeWriter nativeWriter = + new CapturingNativeWriter(cWriter.getVectorSchemaRoot(), cWriter); + ArrowBundleWriter writer = + new ArrowBundleWriter(new NoOpPositionOutputStream(), cWriter, nativeWriter); + + writer.addElement(GenericRow.of(1)); + HeapIntVector vector = new HeapIntVector(2); + vector.setInt(0, 2); + vector.setInt(1, 3); + VectorizedColumnBatch batch = new VectorizedColumnBatch(new ColumnVector[] {vector}); + batch.setNumRows(2); + writer.writeBundle(new VectorizedBundleRecords(batch, null)); + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(2); + assertThat(nativeWriter.snapshots.get(0).objectColumns.get(0)).containsExactly(1); + assertThat(nativeWriter.snapshots.get(1).objectColumns.get(0)).containsExactly(2, 3); + } + @Test public void testAddBatchWithoutDeletionVector() throws IOException { RowType rowType = RowType.of(DataTypes.INT(), DataTypes.BIGINT()); @@ -578,12 +861,12 @@ public void writeIpcBytes(long arrayAddress, long schemaAddress) { if (fv instanceof IntVector) { intValues = new int[rowCount]; for (int i = 0; i < rowCount; i++) { - intValues[i] = ((IntVector) fv).get(i); + intValues[i] = fv.isNull(i) ? 0 : ((IntVector) fv).get(i); } } else if (fv instanceof BigIntVector) { longValues = new long[rowCount]; for (int i = 0; i < rowCount; i++) { - longValues[i] = ((BigIntVector) fv).get(i); + longValues[i] = fv.isNull(i) ? 0L : ((BigIntVector) fv).get(i); } } List colObjects = new ArrayList<>(); @@ -618,6 +901,12 @@ static class Snapshot { } } + private static void setInt(IntVector vector, int value) { + vector.allocateNew(1); + vector.setSafe(0, value); + vector.setValueCount(1); + } + private static class NoOpPositionOutputStream extends PositionOutputStream { private long pos = 0; diff --git a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java index c0cd962695be..95be4e9339b9 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/BundleFormatWriter.java @@ -22,26 +22,21 @@ import java.io.IOException; -/** Format writer with bundle interface. */ +/** + * Format writer with a row-equivalent bundle interface. + * + *

Implementations may consume a compatible bundle natively, convert or copy it, or fall back to + * row-by-row writes. {@link #writeBundle} must preserve the same record values and order as + * invoking {@link #addElement} for every record. It must not retain borrowed buffers after the + * method returns unless it has copied them or acquired independent ownership. + */ public interface BundleFormatWriter extends FormatWriter { /** - * Writes a bundle of records. + * Writes a bundle with semantics equivalent to invoking {@link #addElement} for every record. * * @param bundle the records to be written * @throws IOException if exception happens */ void writeBundle(BundleRecords bundle) throws IOException; - - /** - * Returns whether {@link #writeBundle} is equivalent to invoking {@link #addElement} for every - * record. - * - *

An implementation returning {@code true} must preserve record values and order. It must - * not retain borrowed buffers after {@link #writeBundle} returns unless it has copied them or - * acquired independent ownership. - */ - default boolean supportsRowEquivalentBundleWrite() { - return false; - } } diff --git a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java index 32db7f415a8d..8d5fb5ede353 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/RowDataFileWriter.java @@ -118,8 +118,7 @@ public void write(InternalRow row) throws IOException { public void writeBundle(BundleRecords bundle) throws IOException { if (auxiliaryFileWriters.isEmpty() && sequenceNumberTracker.supportsRowCountUpdate() - && !requiresPerRecordStats() - && supportsRowEquivalentBundleWrite()) { + && !requiresPerRecordStats()) { long rowCount = bundle.rowCount(); super.writeBundle(bundle); sequenceNumberTracker.updateByRowCount(rowCount); diff --git a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java index 04720e503291..96f273091dd7 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java @@ -164,11 +164,6 @@ public void writeBundle(BundleRecords bundle) throws IOException { } } - protected final boolean supportsRowEquivalentBundleWrite() { - return writer instanceof BundleFormatWriter - && ((BundleFormatWriter) writer).supportsRowEquivalentBundleWrite(); - } - protected InternalRow writeImpl(T record) throws IOException { if (closed) { throw new RuntimeException("Writer has already closed!"); diff --git a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java index 706bc18f9acb..192ba65e6977 100644 --- a/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/io/RowDataFileWriterTest.java @@ -23,6 +23,7 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.fileindex.FileIndexOptions; import org.apache.paimon.format.BundleFormatWriter; +import org.apache.paimon.format.FileFormat; import org.apache.paimon.format.FormatWriter; import org.apache.paimon.format.FormatWriterFactory; import org.apache.paimon.format.SimpleColStats; @@ -142,7 +143,7 @@ void testPlainFormatWriterFallsBackToRows() throws Exception { } @Test - void testBundleFormatWriterWithoutOptInFallsBackToRows() throws Exception { + void testBundleFormatWriterCanChooseRowFallback() throws Exception { TestingFallbackBundleFormatWriter formatWriter = new TestingFallbackBundleFormatWriter(); LongCounter sequenceCounter = new LongCounter(); RowDataFileWriter writer = @@ -156,7 +157,7 @@ void testBundleFormatWriterWithoutOptInFallsBackToRows() throws Exception { writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); - assertThat(formatWriter.bundleWrites).isZero(); + assertThat(formatWriter.bundleWrites).isEqualTo(1); assertThat(formatWriter.rowWrites).isEqualTo(2); assertThat(writer.recordCount()).isEqualTo(2); assertThat(sequenceCounter.getValue()).isEqualTo(2); @@ -287,6 +288,36 @@ void testFileIndexFallsBackToRows() throws Exception { assertThat(writer.result().embeddedIndex()).isNotNull(); } + @Test + void testRowSidecarFallsBackToRows() throws Exception { + FileIO fileIO = fileIO(); + TestingBundleFormatWriter formatWriter = new TestingBundleFormatWriter(); + TestingFormatWriter sidecarWriter = new TestingFormatWriter(); + FileFormat rowSidecarFormat = mock(FileFormat.class); + when(rowSidecarFormat.createWriterFactory(ROW_TYPE)) + .thenReturn(new TestingFormatWriterFactory(sidecarWriter)); + Path rowSidecarPath = new Path("file:/tmp/data-file.row"); + RowDataFileWriter writer = + createWriter( + fileIO, + ROW_TYPE, + formatWriter, + SimpleStatsProducer.disabledProducer(), + new LongCounter(), + new FileIndexOptions(), + rowSidecarFormat, + rowSidecarPath); + + writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2))); + + assertThat(formatWriter.bundleWrites).isZero(); + assertThat(formatWriter.rowWrites).isEqualTo(2); + assertThat(sidecarWriter.rowWrites).isEqualTo(2); + + writer.close(); + assertThat(writer.result().extraFiles()).containsExactly(rowSidecarPath.getName()); + } + private static FileIO fileIO() throws IOException { FileIO fileIO = mock(FileIO.class); when(fileIO.getFileSize(PATH)).thenReturn(123L); @@ -300,6 +331,26 @@ private static RowDataFileWriter createWriter( SimpleStatsProducer statsProducer, LongCounter sequenceCounter, FileIndexOptions fileIndexOptions) { + return createWriter( + fileIO, + rowType, + formatWriter, + statsProducer, + sequenceCounter, + fileIndexOptions, + null, + null); + } + + private static RowDataFileWriter createWriter( + FileIO fileIO, + RowType rowType, + FormatWriter formatWriter, + SimpleStatsProducer statsProducer, + LongCounter sequenceCounter, + FileIndexOptions fileIndexOptions, + FileFormat rowSidecarFormat, + Path rowSidecarPath) { return new RowDataFileWriter( fileIO, new FileWriterContext( @@ -314,8 +365,8 @@ private static RowDataFileWriter createWriter( false, false, null, - null, - null); + rowSidecarFormat, + rowSidecarPath); } private static BundleRecords rows(InternalRow... rows) { @@ -378,11 +429,6 @@ public void writeBundle(BundleRecords bundle) { bundleWrites++; writtenBundle = bundle; } - - @Override - public boolean supportsRowEquivalentBundleWrite() { - return true; - } } private static class TestingFallbackBundleFormatWriter extends TestingFormatWriter @@ -412,11 +458,6 @@ private TestingThrowingBundleFormatWriter(IOException failure) { public void writeBundle(BundleRecords bundle) throws IOException { throw failure; } - - @Override - public boolean supportsRowEquivalentBundleWrite() { - return true; - } } private static class TestingStatsProducer implements SimpleStatsProducer { diff --git a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java index ee1f572462d7..a23d95489f9e 100644 --- a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java +++ b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java @@ -19,6 +19,7 @@ package org.apache.paimon.format.lance; import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.arrow.vector.ArrowFormatWriter; import org.apache.paimon.data.InternalRow; import org.apache.paimon.format.BundleFormatWriter; @@ -65,12 +66,20 @@ public void addElement(InternalRow internalRow) throws IOException { @Override public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { - add(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); - } else { - for (InternalRow row : bundleRecords) { - addElement(row); + ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; + VectorSchemaRoot root = arrowBundle.getVectorSchemaRoot(); + if (arrowFormatWriter.isArrowBundleSchemaCompatible(arrowBundle) + && ArrowUtils.hasSameRootAllocator(root, arrowFormatWriter.getAllocator())) { + flush(); + nativeWriter.ensureInitialized(arrowFormatWriter.getAllocator()); + add(root); + return; } } + + for (InternalRow row : bundleRecords) { + addElement(row); + } } public void add(VectorSchemaRoot vsr) throws IOException { diff --git a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java index f9329907f951..b3e29eb0a0b3 100644 --- a/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java +++ b/paimon-lance/src/main/java/org/apache/paimon/format/lance/jni/LanceWriter.java @@ -18,6 +18,8 @@ package org.apache.paimon.format.lance.jni; +import org.apache.paimon.arrow.ArrowUtils; + import com.lancedb.lance.file.LanceFileWriter; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; @@ -32,6 +34,7 @@ public class LanceWriter { private final String path; private final Map storageOptions; private LanceFileWriter writer; + private BufferAllocator allocator; private long bytesWritten = 0; public LanceWriter(String path, Map storageOptions) { @@ -44,12 +47,25 @@ public Long getWrittenPosition() { } public void writeVsr(VectorSchemaRoot vsr) throws IOException { - initWriteLazy(vsr.getVector(0).getAllocator()); + BufferAllocator sourceAllocator = vsr.getVector(0).getAllocator(); + initWriteLazy(sourceAllocator); + if (!ArrowUtils.hasSameRootAllocator(vsr, allocator)) { + throw new IllegalArgumentException( + "Lance writer cannot consume Arrow buffers from a different allocator root."); + } this.bytesWritten += vsr.getFieldVectors().stream().mapToLong(FieldVector::getBufferSize).sum(); this.writer.write(vsr); } + /** + * Initializes the native writer with an allocator whose lifetime is owned by the surrounding + * format writer. + */ + public void ensureInitialized(BufferAllocator bufferAllocator) throws IOException { + initWriteLazy(bufferAllocator); + } + public void close() throws IOException { if (writer != null) { try { @@ -58,6 +74,7 @@ public void close() throws IOException { throw new IOException(e); } this.writer = null; + this.allocator = null; } } @@ -68,6 +85,7 @@ public String path() { private void initWriteLazy(BufferAllocator bufferAllocator) throws IOException { if (writer == null) { writer = LanceFileWriter.open(path, bufferAllocator, null, storageOptions); + allocator = bufferAllocator; } } } diff --git a/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java new file mode 100644 index 000000000000..7d5161d5dca2 --- /dev/null +++ b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java @@ -0,0 +1,297 @@ +/* + * 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.format.lance; + +import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; +import org.apache.paimon.arrow.vector.ArrowFormatWriter; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.data.serializer.InternalRowSerializer; +import org.apache.paimon.format.lance.jni.LanceWriter; +import org.apache.paimon.fs.Path; +import org.apache.paimon.reader.RecordReaderIterator; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link LanceRecordsWriter}. */ +class LanceRecordsWriterTest { + + @Test + void testArrowBundlePreservesRowBundleRowOrder() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true); + BufferAllocator writerAllocator = arrowWriter.getAllocator(); + CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + writer.addElement(GenericRow.of(1)); + try (BufferAllocator sourceAllocator = + arrowWriter + .getAllocator() + .newChildAllocator("lance-bundle-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + setInt((IntVector) root.getVector("value"), 2); + root.setRowCount(1); + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + writer.addElement(GenericRow.of(3)); + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(3); + assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(1); + assertThat(nativeWriter.snapshots.get(1).values.get(0)).containsExactly(2); + assertThat(nativeWriter.snapshots.get(2).values.get(0)).containsExactly(3); + assertThat(nativeWriter.initializedAllocator).isSameAs(writerAllocator); + } + + @Test + void testNativeWriterUsesDirectSameRootAndFallbackCrossRoot(@TempDir java.nio.file.Path tempDir) + throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + String file = tempDir.resolve("bundle_" + UUID.randomUUID()).toString(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true); + TrackingLanceWriter nativeWriter = new TrackingLanceWriter(file); + try (LanceRecordsWriter writer = + new LanceRecordsWriter( + nativeWriter::getWrittenPosition, arrowWriter, nativeWriter)) { + writer.addElement(GenericRow.of(1)); + try (BufferAllocator sourceAllocator = + arrowWriter + .getAllocator() + .newChildAllocator( + "lance-direct-native-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + nativeWriter.sameRootBundle = root; + setInt((IntVector) root.getVector("value"), 2); + root.setRowCount(1); + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + + try (BufferAllocator sourceAllocator = new org.apache.arrow.memory.RootAllocator(); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + nativeWriter.crossRootBundle = root; + setInt((IntVector) root.getVector("value"), 3); + root.setRowCount(1); + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + writer.addElement(GenericRow.of(4)); + } + + assertThat(nativeWriter.sameRootBundleWrites).isEqualTo(1); + assertThat(nativeWriter.crossRootBundleWrites).isZero(); + + InternalRowSerializer serializer = new InternalRowSerializer(rowType); + List actual = new ArrayList<>(); + try (RecordReaderIterator iterator = + new RecordReaderIterator<>( + new LanceRecordsReader( + new Path(file), null, rowType, 1024, Collections.emptyMap()))) { + while (iterator.hasNext()) { + actual.add(serializer.copy(iterator.next()).getInt(0)); + } + } + assertThat(actual).containsExactly(1, 2, 3, 4); + } + + @Test + void testReorderedArrowBundleFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(writerType, 1024, true); + CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + try (BufferAllocator sourceAllocator = + arrowWriter + .getAllocator() + .newChildAllocator("lance-schema-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { + setInt((IntVector) root.getVector("b"), 20); + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + Snapshot snapshot = nativeWriter.snapshots.get(0); + assertThat(snapshot.fieldNames).containsExactly("a", "b"); + assertThat(snapshot.values.get(0)).containsExactly(10); + assertThat(snapshot.values.get(1)).containsExactly(20); + } + + @Test + void testMissingColumnArrowBundleFallsBackToRows() throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = RowType.builder().field("a", DataTypes.INT()).build(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(writerType, 1024, true); + CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + try (BufferAllocator sourceAllocator = + arrowWriter + .getAllocator() + .newChildAllocator("lance-missing-column-test", 0, Long.MAX_VALUE); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + writer.writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + writer.close(); + + assertThat(nativeWriter.snapshots).hasSize(1); + Snapshot snapshot = nativeWriter.snapshots.get(0); + assertThat(snapshot.fieldNames).containsExactly("a", "b"); + assertThat(snapshot.values.get(0)).containsExactly(10); + assertThat(snapshot.values.get(1)).containsExactly((Integer) null); + } + + @Test + void testDifferentAllocatorRootFallsBackToRows() throws Exception { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true); + CapturingLanceWriter nativeWriter = new CapturingLanceWriter(); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + try (BufferAllocator sourceAllocator = new org.apache.arrow.memory.RootAllocator(); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(rowType, sourceAllocator)) { + nativeWriter.disallowedRoot = root; + setInt((IntVector) root.getVector("value"), 10); + root.setRowCount(1); + + writer.writeBundle(new ArrowBundleRecords(root, rowType, true)); + } + writer.close(); + + assertThat(nativeWriter.disallowedRootWrites).isZero(); + assertThat(nativeWriter.snapshots).hasSize(1); + assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(10); + } + + private static void setInt(IntVector vector, int value) { + vector.allocateNew(1); + vector.setSafe(0, value); + vector.setValueCount(1); + } + + private static class CapturingLanceWriter extends LanceWriter { + + private final List snapshots = new ArrayList<>(); + private BufferAllocator initializedAllocator; + private VectorSchemaRoot disallowedRoot; + private int disallowedRootWrites; + + private CapturingLanceWriter() { + super("unused", Collections.emptyMap()); + } + + @Override + public void ensureInitialized(BufferAllocator bufferAllocator) { + initializedAllocator = bufferAllocator; + } + + @Override + public void writeVsr(VectorSchemaRoot root) { + if (root == disallowedRoot) { + disallowedRootWrites++; + } + List fieldNames = + root.getSchema().getFields().stream() + .map(field -> field.getName()) + .collect(Collectors.toList()); + List> values = new ArrayList<>(); + for (int column = 0; column < root.getFieldVectors().size(); column++) { + IntVector vector = (IntVector) root.getVector(column); + List columnValues = new ArrayList<>(); + for (int row = 0; row < root.getRowCount(); row++) { + columnValues.add(vector.isNull(row) ? null : vector.get(row)); + } + values.add(columnValues); + } + snapshots.add(new Snapshot(fieldNames, values)); + } + + @Override + public void close() throws IOException {} + + @Override + public String path() { + return "unused"; + } + } + + private static class TrackingLanceWriter extends LanceWriter { + + private VectorSchemaRoot sameRootBundle; + private VectorSchemaRoot crossRootBundle; + private int sameRootBundleWrites; + private int crossRootBundleWrites; + + private TrackingLanceWriter(String path) { + super(path, Collections.emptyMap()); + } + + @Override + public void writeVsr(VectorSchemaRoot root) throws IOException { + if (root == sameRootBundle) { + sameRootBundleWrites++; + } + if (root == crossRootBundle) { + crossRootBundleWrites++; + } + super.writeVsr(root); + } + } + + private static class Snapshot { + + private final List fieldNames; + private final List> values; + + private Snapshot(List fieldNames, List> values) { + this.fieldNames = fieldNames; + this.values = values; + } + } +} 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 86300a820d5a..1a3e4d4149b6 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 @@ -143,11 +143,6 @@ public void writeBundle(BundleRecords bundleRecords) { } } - @Override - public boolean supportsRowEquivalentBundleWrite() { - return true; - } - @Override public boolean reachTargetSize(boolean suggestedCheck, long targetSize) { if (!suggestedCheck) { diff --git a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java index 7798fa09d419..502b4185e870 100644 --- a/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java +++ b/paimon-vortex/paimon-vortex-format/src/main/java/org/apache/paimon/format/vortex/VortexRecordsWriter.java @@ -19,6 +19,7 @@ package org.apache.paimon.format.vortex; import org.apache.paimon.arrow.ArrowBundleRecords; +import org.apache.paimon.arrow.ArrowUtils; import org.apache.paimon.arrow.vector.ArrowCStruct; import org.apache.paimon.arrow.vector.ArrowFormatCWriter; import org.apache.paimon.data.InternalRow; @@ -38,6 +39,9 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import javax.annotation.Nullable; + +import java.io.ByteArrayInputStream; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -65,6 +69,11 @@ public class VortexRecordsWriter implements BundleFormatWriter { private final List retainedResources; private ArrowFormatCWriter currentWriter; + // Direct bundles are copied into this writer-owned allocator. Vortex synchronously imports + // independent Arrow buffer references before writeBatch returns, while nativeWriter.close() + // completes outstanding asynchronous writes before this allocator is closed. + @Nullable private RootAllocator bundleAllocator; + private long jniCost = 0; private long ffiBytes = 0; @@ -102,13 +111,17 @@ public void addElement(InternalRow internalRow) throws IOException { @Override public void writeBundle(BundleRecords bundleRecords) throws IOException { if (bundleRecords instanceof ArrowBundleRecords) { - flush(); - writeBundleVsr(((ArrowBundleRecords) bundleRecords).getVectorSchemaRoot()); - } else { - for (InternalRow row : bundleRecords) { - addElement(row); + ArrowBundleRecords arrowBundle = (ArrowBundleRecords) bundleRecords; + if (currentWriter.formatWriter().isArrowBundleSchemaCompatible(arrowBundle)) { + flush(); + writeBundleVsr(arrowBundle.getVectorSchemaRoot()); + return; } } + + for (InternalRow row : bundleRecords) { + addElement(row); + } } @Override @@ -136,10 +149,26 @@ public void close() throws IOException { // Release all retained resources now that async writes are done. for (AutoCloseable res : retainedResources) { - closeQuietly(res); + try { + closeQuietly(res); + } catch (Throwable t) { + throwable = addSuppressed(throwable, t); + } } retainedResources.clear(); - closeQuietly(currentWriter); + try { + closeQuietly(currentWriter); + } catch (Throwable t) { + throwable = addSuppressed(throwable, t); + } + + if (bundleAllocator != null) { + try { + bundleAllocator.close(); + } catch (Throwable t) { + throwable = addSuppressed(throwable, t); + } + } try { session.close(); @@ -170,34 +199,38 @@ private void flush() throws IOException { } } - /** Write an external VSR (from writeBundle) via IPC copy into an independent allocator. */ + /** Write an external VSR (from writeBundle) via IPC copy into the writer-owned allocator. */ private void writeBundleVsr(VectorSchemaRoot vsr) throws IOException { ffiBytes += bufferBytes(vsr); - byte[] ipc = org.apache.paimon.arrow.ArrowUtils.serializeToIpc(vsr); - RootAllocator bundleAllocator = new RootAllocator(Long.MAX_VALUE); - try { - ArrowStreamReader ipcReader = - new ArrowStreamReader(new java.io.ByteArrayInputStream(ipc), bundleAllocator); - ipcReader.loadNextBatch(); - VectorSchemaRoot copy = ipcReader.getVectorSchemaRoot(); - - ArrowArray arrowArray = ArrowArray.allocateNew(bundleAllocator); - ArrowSchema arrowSchema = ArrowSchema.allocateNew(bundleAllocator); - Data.exportVectorSchemaRoot(bundleAllocator, copy, null, arrowArray, arrowSchema); + byte[] ipc = ArrowUtils.serializeToIpc(vsr); + RootAllocator allocator = bundleAllocator(); + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(ipc), allocator); + ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + if (!reader.loadNextBatch()) { + throw new IOException("Arrow IPC copy did not contain a record batch."); + } + Data.exportVectorSchemaRoot( + allocator, reader.getVectorSchemaRoot(), null, array, schema); + long t1 = System.currentTimeMillis(); - nativeWriter.writeBatch(arrowArray.memoryAddress(), arrowSchema.memoryAddress()); + try { + nativeWriter.writeBatch(array.memoryAddress(), schema.memoryAddress()); + } finally { + ArrowUtils.releaseCDataIfNeeded(array, schema); + } jniCost += (System.currentTimeMillis() - t1); - // Retain all resources that own the exported C Data buffers and release - // callbacks. Rust holds async zero-copy references via Arc. - // Order matters: close ipcReader (owns VectorSchemaRoot) before allocator. - retainedResources.add(ipcReader); - retainedResources.add(bundleAllocator); - } catch (Exception e) { - closeQuietly(bundleAllocator); - throw e instanceof IOException ? (IOException) e : new IOException(e); } } + private RootAllocator bundleAllocator() { + if (bundleAllocator == null) { + bundleAllocator = new RootAllocator(Long.MAX_VALUE); + } + return bundleAllocator; + } + private static long bufferBytes(VectorSchemaRoot vsr) { long bytes = 0; for (int i = 0; i < vsr.getFieldVectors().size(); i++) { diff --git a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java index e97e3cf80fcc..81dab2774d19 100644 --- a/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java +++ b/paimon-vortex/paimon-vortex-format/src/test/java/org/apache/paimon/format/vortex/VortexReaderWriterTest.java @@ -19,6 +19,7 @@ package org.apache.paimon.format.vortex; 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.GenericRow; @@ -45,6 +46,9 @@ import org.apache.paimon.utils.RoaringBitmap32; import dev.vortex.jni.NativeRuntime; +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 org.junit.jupiter.api.io.TempDir; @@ -264,7 +268,7 @@ public void testWriteDoesNotLeakArrowMemoryOnClose(@TempDir java.nio.file.Path t } @Test - public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( + public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffersAcrossBatches( @TempDir java.nio.file.Path tempDir) throws Exception { RowType rowType = RowType.builder() @@ -289,6 +293,16 @@ public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( arrowWriter.write(GenericRow.of(2, BinaryString.fromString("world"))); arrowWriter.flush(); + ((BundleFormatWriter) writer) + .writeBundle( + new ArrowBundleRecords( + arrowWriter.getVectorSchemaRoot(), rowType, true)); + + arrowWriter.reset(); + arrowWriter.write(GenericRow.of(3, BinaryString.fromString("second"))); + arrowWriter.write(GenericRow.of(4, BinaryString.fromString("batch"))); + arrowWriter.flush(); + ((BundleFormatWriter) writer) .writeBundle( new ArrowBundleRecords( @@ -312,11 +326,114 @@ public void testArrowBundleRecordsWriteDoesNotBorrowCallerBuffers( actualRows.add(serializer.copy(iterator.next())); } - assertEquals(2, actualRows.size()); + assertEquals(4, actualRows.size()); assertEquals(1, actualRows.get(0).getInt(0)); assertEquals(BinaryString.fromString("hello"), actualRows.get(0).getString(1)); assertEquals(2, actualRows.get(1).getInt(0)); assertEquals(BinaryString.fromString("world"), actualRows.get(1).getString(1)); + assertEquals(3, actualRows.get(2).getInt(0)); + assertEquals(BinaryString.fromString("second"), actualRows.get(2).getString(1)); + assertEquals(4, actualRows.get(3).getInt(0)); + assertEquals(BinaryString.fromString("batch"), actualRows.get(3).getString(1)); + } + } + + @Test + public void testArrowBundlePreservesRowBundleRowOrder(@TempDir java.nio.file.Path tempDir) + throws Exception { + RowType rowType = + RowType.builder() + .field("f_int", DataTypes.INT()) + .field("f_string", DataTypes.STRING()) + .build(); + VortexFileFormat format = + new VortexFileFormatFactory() + .create(new FileFormatFactory.FormatContext(new Options(), 1024, 1024)); + FileIO fileIO = new LocalFileIO(); + Path testFile = + new Path(new Path(tempDir.toUri()), "test_bundle_order_" + UUID.randomUUID()); + + try (FormatWriter writer = + ((SupportsDirectWrite) format.createWriterFactory(rowType)) + .create(fileIO, testFile, ""); + ArrowFormatWriter arrowWriter = new ArrowFormatWriter(rowType, 1024, true)) { + writer.addElement(GenericRow.of(1, BinaryString.fromString("row"))); + + arrowWriter.write(GenericRow.of(2, BinaryString.fromString("bundle"))); + arrowWriter.flush(); + ((BundleFormatWriter) writer) + .writeBundle( + new ArrowBundleRecords( + arrowWriter.getVectorSchemaRoot(), rowType, true)); + + writer.addElement(GenericRow.of(3, BinaryString.fromString("row"))); + } + + InternalRowSerializer serializer = new InternalRowSerializer(rowType); + FormatReaderFactory readerFactory = format.createReaderFactory(rowType, rowType, null); + try (RecordReader reader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, testFile, fileIO.getFileSize(testFile), null)); + RecordReaderIterator iterator = new RecordReaderIterator<>(reader)) { + List actualRows = new ArrayList<>(); + while (iterator.hasNext()) { + actualRows.add(serializer.copy(iterator.next())); + } + + assertEquals(3, actualRows.size()); + assertEquals(1, actualRows.get(0).getInt(0)); + assertEquals(2, actualRows.get(1).getInt(0)); + assertEquals(3, actualRows.get(2).getInt(0)); + } + } + + @Test + public void testReorderedArrowBundleFallsBackToRows(@TempDir java.nio.file.Path tempDir) + throws Exception { + RowType writerType = + RowType.builder().field("a", DataTypes.INT()).field("b", DataTypes.INT()).build(); + RowType sourceType = + RowType.builder().field("b", DataTypes.INT()).field("a", DataTypes.INT()).build(); + + Options options = new Options(); + VortexFileFormat format = + new VortexFileFormatFactory() + .create(new FileFormatFactory.FormatContext(options, 1024, 1024)); + FileIO fileIO = new LocalFileIO(); + Path testFile = + new Path(new Path(tempDir.toUri()), "test_reordered_bundle_" + UUID.randomUUID()); + + try (FormatWriter writer = + ((SupportsDirectWrite) format.createWriterFactory(writerType)) + .create(fileIO, testFile, ""); + RootAllocator sourceAllocator = new RootAllocator(); + VectorSchemaRoot root = + ArrowUtils.createVectorSchemaRoot(sourceType, sourceAllocator)) { + setInt((IntVector) root.getVector("b"), 20); + setInt((IntVector) root.getVector("a"), 10); + root.setRowCount(1); + + ((BundleFormatWriter) writer) + .writeBundle(new ArrowBundleRecords(root, writerType, true)); + } + + InternalRowSerializer serializer = new InternalRowSerializer(writerType); + FormatReaderFactory readerFactory = + format.createReaderFactory(writerType, writerType, null); + try (RecordReader reader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, testFile, fileIO.getFileSize(testFile), null)); + RecordReaderIterator iterator = new RecordReaderIterator<>(reader)) { + List actualRows = new ArrayList<>(); + while (iterator.hasNext()) { + actualRows.add(serializer.copy(iterator.next())); + } + + assertEquals(1, actualRows.size()); + assertEquals(10, actualRows.get(0).getInt(0)); + assertEquals(20, actualRows.get(0).getInt(1)); } } @@ -367,6 +484,12 @@ public void testReadWithSelection(@TempDir java.nio.file.Path tempDir) throws Ex } } + private static void setInt(IntVector vector, int value) { + vector.allocateNew(1); + vector.setSafe(0, value); + vector.setValueCount(1); + } + @Test public void testReadWithVirtualRowTrackingField(@TempDir java.nio.file.Path tempDir) throws Exception { From 935fe229f03f81bc6b8c4b7e8f38d32763d3e228 Mon Sep 17 00:00:00 2001 From: mingfeng Date: Mon, 10 Aug 2026 04:44:41 -0700 Subject: [PATCH 6/6] [arrow] Simplify allocator root validation --- .../org/apache/paimon/arrow/ArrowUtils.java | 33 ++++--------------- 1 file changed, 7 insertions(+), 26 deletions(-) diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java index 8013d9ec8f82..5ea4ac8aaac9 100644 --- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java +++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java @@ -302,17 +302,8 @@ public static byte[] serializeToIpc(VectorSchemaRoot vsr) { */ public static boolean hasSameRootAllocator( VectorSchemaRoot vectorSchemaRoot, BufferAllocator allocator) { - if (vectorSchemaRoot.getFieldVectors().isEmpty()) { - return false; - } - - BufferAllocator expectedRoot = rootAllocator(allocator); - for (FieldVector vector : vectorSchemaRoot.getFieldVectors()) { - if (!hasSameRootAllocator(vector, expectedRoot)) { - return false; - } - } - return true; + List vectors = vectorSchemaRoot.getFieldVectors(); + return !vectors.isEmpty() && allVectorsShareRootWith(vectors, allocator.getRoot()); } public static void serializeToIpc(VectorSchemaRoot vsr, OutputStream out) { @@ -349,21 +340,11 @@ private static long zoneCastedTimestampZoneCastToEpoch( } } - private static BufferAllocator rootAllocator(BufferAllocator allocator) { - BufferAllocator current = allocator; - while (current.getParentAllocator() != null) { - current = current.getParentAllocator(); - } - return current; - } - - private static boolean hasSameRootAllocator(FieldVector vector, BufferAllocator expectedRoot) { - if (rootAllocator(vector.getAllocator()) != expectedRoot) { - return false; - } - - for (FieldVector child : vector.getChildrenFromFields()) { - if (!hasSameRootAllocator(child, expectedRoot)) { + private static boolean allVectorsShareRootWith( + List vectors, BufferAllocator expectedRoot) { + for (FieldVector vector : vectors) { + if (vector.getAllocator().getRoot() != expectedRoot + || !allVectorsShareRootWith(vector.getChildrenFromFields(), expectedRoot)) { return false; } }