This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 9431b894f5 [core] Support batch writes with file indexes (#10603)
9431b894f5 is described below
commit 9431b894f53338ed331c2d79c3a4824cf58398b0
Author: jianguotian <[email protected]>
AuthorDate: Fri Oct 16 13:57:31 2026 +0800
[core] Support batch writes with file indexes (#10603)
---
.../org/apache/paimon/io/RowDataFileWriter.java | 47 +++-
.../apache/paimon/io/RowDataFileWriterTest.java | 257 ++++++++++++++++++++-
.../format/mosaic/MosaicReaderWriterTest.java | 105 +++++++++
3 files changed, 397 insertions(+), 12 deletions(-)
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 8d5fb5ede3..9dfbaf7ea1 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
@@ -116,17 +116,48 @@ public class RowDataFileWriter extends
StatsCollectingSingleFileWriter<InternalR
@Override
public void writeBundle(BundleRecords bundle) throws IOException {
- if (auxiliaryFileWriters.isEmpty()
- && sequenceNumberTracker.supportsRowCountUpdate()
- && !requiresPerRecordStats()) {
- long rowCount = bundle.rowCount();
- super.writeBundle(bundle);
- sequenceNumberTracker.updateByRowCount(rowCount);
+ if (!canForwardBundle(bundle)) {
+ for (InternalRow row : bundle) {
+ write(row);
+ }
return;
}
- for (InternalRow row : bundle) {
- write(row);
+ long rowCount = bundle.rowCount();
+ writeFileIndexes(bundle);
+ super.writeBundle(bundle);
+ sequenceNumberTracker.updateByRowCount(rowCount);
+ }
+
+ private boolean canForwardBundle(BundleRecords bundle) {
+ if (!sequenceNumberTracker.supportsRowCountUpdate() ||
requiresPerRecordStats()) {
+ return false;
+ }
+ return auxiliaryFileWriters.isEmpty()
+ || (bundle instanceof VectorizedBundleRecords
+ && auxiliaryFileWriters.stream()
+ .allMatch(
+ writer -> writer instanceof
DataFileIndexAuxiliaryWriter));
+ }
+
+ private void writeFileIndexes(BundleRecords bundle) throws IOException {
+ if (auxiliaryFileWriters.isEmpty()) {
+ return;
+ }
+ if (closed) {
+ throw new RuntimeException("Writer has already closed!");
+ }
+ // Vectorized bundles expose repeatable row views without consuming
the batch.
+ // Read index fields before the format writer may transfer the
borrowed buffers.
+ try {
+ for (InternalRow row : bundle) {
+ for (DataFileAuxiliaryWriter auxiliaryFileWriter :
auxiliaryFileWriters) {
+ auxiliaryFileWriter.write(row);
+ }
+ }
+ } catch (Throwable e) {
+ abort();
+ throw e;
}
}
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 192ba65e69..887b969ceb 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
@@ -21,33 +21,50 @@ 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.data.columnar.ColumnVector;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
+import org.apache.paimon.data.columnar.heap.HeapIntVector;
+import org.apache.paimon.fileindex.FileIndexFormat;
import org.apache.paimon.fileindex.FileIndexOptions;
+import org.apache.paimon.fileindex.FileIndexReader;
+import org.apache.paimon.fileindex.bitmap.BitmapIndexResult;
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;
import org.apache.paimon.format.SupportsDirectWrite;
+import org.apache.paimon.fs.ByteArraySeekableStream;
import org.apache.paimon.fs.FileIO;
import org.apache.paimon.fs.Path;
import org.apache.paimon.fs.PositionOutputStream;
+import org.apache.paimon.fs.local.LocalFileIO;
import org.apache.paimon.manifest.FileSource;
import org.apache.paimon.options.Options;
+import org.apache.paimon.predicate.FieldRef;
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.apache.paimon.utils.RoaringBitmap32;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
+import static org.apache.paimon.io.DataFilePathFactory.dataFileToFileIndexPath;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -57,6 +74,8 @@ 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();
+ @TempDir java.nio.file.Path tempDir;
+
@Test
void testEligibleBundleIsForwardedWithoutIteration() throws Exception {
FileIO fileIO = fileIO();
@@ -224,9 +243,9 @@ class RowDataFileWriterTest {
formatWriter,
statsProducer,
sequenceCounter,
- new FileIndexOptions());
+ bitmapOptions());
- writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2)));
+ writer.writeBundle(vectorizedRows(1, 2));
assertThat(formatWriter.bundleWrites).isZero();
assertThat(formatWriter.rowWrites).isEqualTo(2);
@@ -288,6 +307,165 @@ class RowDataFileWriterTest {
assertThat(writer.result().embeddedIndex()).isNotNull();
}
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testVectorizedBundlesPreserveBitmapPositions(boolean externalIndex)
throws Exception {
+ Options options = new Options();
+ options.set("file-index.bitmap.columns", "id");
+ options.set("file-index.in-manifest-threshold", externalIndex ? "0B" :
"1MB");
+ FileIndexOptions indexOptions = new FileIndexOptions(new
CoreOptions(options));
+ LongCounter sequenceCounter = new LongCounter(5);
+ List<TestingBundleFormatWriter> formatWriters = new ArrayList<>();
+ RollingFileWriterImpl<InternalRow, DataFileMeta> writer =
+ new RollingFileWriterImpl<>(
+ () -> {
+ TestingBundleFormatWriter formatWriter =
+ new TestingBundleFormatWriter();
+ formatWriters.add(formatWriter);
+ Path path =
+ new Path(
+ tempDir.resolve("data-" +
formatWriters.size())
+ .toUri());
+ return createIndexedWriter(
+ path, formatWriter, sequenceCounter,
indexOptions);
+ },
+ Long.MAX_VALUE,
+ 7L);
+ // The first file mixes row writes, selected batches and an empty
batch.
+ writer.write(GenericRow.of(1));
+ VectorizedBundleRecords source = vectorizedRows(1, null, 2, 1);
+ writer.writeBundle(new VectorizedBundleRecords(source.batch(), new
int[] {2, 0, 1, 2}));
+ writer.writeBundle(vectorizedRows());
+ writer.writeBundle(vectorizedRows(2, 1));
+ // Index positions must restart at zero after rolling to the second
file.
+ writer.writeBundle(vectorizedRows(2, 1));
+ writer.close();
+
+ assertThat(formatWriters).extracting(w ->
w.bundleWrites).containsExactly(3, 1);
+ assertThat(formatWriters).extracting(w ->
w.rowWrites).containsExactly(1, 0);
+ assertThat(writer.recordCount()).isEqualTo(9);
+ assertThat(sequenceCounter.getValue()).isEqualTo(14);
+ List<DataFileMeta> files = writer.result();
+
assertThat(files).extracting(DataFileMeta::rowCount).containsExactly(7L, 2L);
+
assertThat(files).extracting(DataFileMeta::minSequenceNumber).containsExactly(5L,
12L);
+
assertThat(files).extracting(DataFileMeta::maxSequenceNumber).containsExactly(11L,
13L);
+ try (FileIndexFormat.Reader reader = readFileIndex(files.get(0),
externalIndex)) {
+ FileIndexReader index =
reader.readColumnIndex("id").iterator().next();
+ assertBitmapPositions(index, 1, 0, 2, 6);
+ assertBitmapPositions(index, 2, 1, 4, 5);
+ assertBitmapPositions(index, null, 3);
+ }
+ try (FileIndexFormat.Reader reader = readFileIndex(files.get(1),
externalIndex)) {
+ FileIndexReader index =
reader.readColumnIndex("id").iterator().next();
+ assertBitmapPositions(index, 1, 1);
+ assertBitmapPositions(index, 2, 0);
+ assertBitmapPositions(index, null);
+ }
+ }
+
+ @Test
+ void testIndexedVectorizedBundleAllowsFormatRowFallback() throws Exception
{
+ TestingFallbackBundleFormatWriter formatWriter = new
TestingFallbackBundleFormatWriter();
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO(),
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ new LongCounter(),
+ bitmapOptions());
+ writer.writeBundle(vectorizedRows(1, 2));
+ assertThat(formatWriter.bundleWrites).isEqualTo(1);
+ assertThat(formatWriter.rowWrites).isEqualTo(2);
+ writer.close();
+ assertThat(writer.result().embeddedIndex()).isNotNull();
+ }
+
+ @Test
+ void testIndexFailureAbortsBeforeWritingBatch() throws Exception {
+ FileIO fileIO = fileIO();
+ TestingBundleFormatWriter formatWriter = new
TestingBundleFormatWriter();
+ LongCounter sequenceCounter = new LongCounter(5);
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO,
+ ROW_TYPE,
+ formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ sequenceCounter,
+ bitmapOptions());
+ RuntimeException failure = new RuntimeException("index field access
failed");
+ HeapIntVector vector =
+ new HeapIntVector(2) {
+ @Override
+ public int getInt(int rowId) {
+ if (rowId == 1) {
+ throw failure;
+ }
+ return 1;
+ }
+ };
+ VectorizedColumnBatch batch = new VectorizedColumnBatch(new
ColumnVector[] {vector});
+ batch.setNumRows(2);
+ assertThatThrownBy(() -> writer.writeBundle(new
VectorizedBundleRecords(batch, null)))
+ .isSameAs(failure);
+ assertThat(formatWriter.bundleWrites).isZero();
+ assertThat(formatWriter.closeCalls).isEqualTo(1);
+ assertThat(writer.recordCount()).isZero();
+ assertThat(sequenceCounter.getValue()).isEqualTo(5);
+ verify(fileIO).deleteQuietly(PATH);
+ verify(fileIO,
atLeastOnce()).deleteQuietly(dataFileToFileIndexPath(PATH));
+ }
+
+ @Test
+ void testClosedWriterRejectsBundleBeforeReadingIndexFields() throws
Exception {
+ FileIO fileIO = fileIO();
+ RowDataFileWriter writer =
+ createWriter(
+ fileIO,
+ ROW_TYPE,
+ new TestingBundleFormatWriter(),
+ SimpleStatsProducer.disabledProducer(),
+ new LongCounter(),
+ bitmapOptions());
+ writer.close();
+ VectorizedBundleRecords bundle =
+ new VectorizedBundleRecords(vectorizedRows(1).batch(), null) {
+ @Override
+ public Iterator<InternalRow> iterator() {
+ throw new AssertionError("Closed writer must not
consume the index batch");
+ }
+ };
+ assertThatThrownBy(() -> writer.writeBundle(bundle))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessage("Writer has already closed!");
+ verify(fileIO, never()).deleteQuietly(PATH);
+ verify(fileIO, never()).deleteQuietly(dataFileToFileIndexPath(PATH));
+ }
+
+ @Test
+ void testBundleFailureAlsoAbortsPreparedIndex() 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,
+ bitmapOptions());
+ assertThatThrownBy(() -> writer.writeBundle(vectorizedRows(1,
2))).isSameAs(failure);
+ assertThat(formatWriter.closeCalls).isEqualTo(1);
+ assertThat(writer.recordCount()).isZero();
+ assertThat(sequenceCounter.getValue()).isEqualTo(5);
+ verify(fileIO).deleteQuietly(PATH);
+ verify(fileIO,
atLeastOnce()).deleteQuietly(dataFileToFileIndexPath(PATH));
+ }
+
@Test
void testRowSidecarFallsBackToRows() throws Exception {
FileIO fileIO = fileIO();
@@ -304,11 +482,11 @@ class RowDataFileWriterTest {
formatWriter,
SimpleStatsProducer.disabledProducer(),
new LongCounter(),
- new FileIndexOptions(),
+ bitmapOptions(),
rowSidecarFormat,
rowSidecarPath);
- writer.writeBundle(rows(GenericRow.of(1), GenericRow.of(2)));
+ writer.writeBundle(vectorizedRows(1, 2));
assertThat(formatWriter.bundleWrites).isZero();
assertThat(formatWriter.rowWrites).isEqualTo(2);
@@ -324,6 +502,77 @@ class RowDataFileWriterTest {
return fileIO;
}
+ private static RowDataFileWriter createIndexedWriter(
+ Path path,
+ FormatWriter formatWriter,
+ LongCounter sequenceCounter,
+ FileIndexOptions indexOptions) {
+ return new RowDataFileWriter(
+ LocalFileIO.create(),
+ new FileWriterContext(
+ (out, compression) -> formatWriter,
+ SimpleStatsProducer.disabledProducer(),
+ "none"),
+ path,
+ ROW_TYPE,
+ 1L,
+ () -> sequenceCounter,
+ indexOptions,
+ FileSource.APPEND,
+ false,
+ false,
+ false,
+ null,
+ null,
+ null);
+ }
+
+ private FileIndexFormat.Reader readFileIndex(DataFileMeta file, boolean
externalIndex)
+ throws IOException {
+ if (externalIndex) {
+ assertThat(file.embeddedIndex()).isNull();
+ assertThat(file.extraFiles()).hasSize(1);
+ return FileIndexFormat.createReader(
+ LocalFileIO.create()
+ .newInputStream(
+ new
Path(tempDir.resolve(file.extraFiles().get(0)).toUri())),
+ ROW_TYPE);
+ }
+ assertThat(file.extraFiles()).isEmpty();
+ assertThat(file.embeddedIndex()).isNotNull();
+ return FileIndexFormat.createReader(
+ new ByteArraySeekableStream(file.embeddedIndex()), ROW_TYPE);
+ }
+
+ private static void assertBitmapPositions(
+ FileIndexReader index, Integer value, int... positions) {
+ FieldRef field = new FieldRef(0, "id", DataTypes.INT());
+ BitmapIndexResult result =
+ (BitmapIndexResult)
+ (value == null ? index.visitIsNull(field) :
index.visitEqual(field, value));
+
assertThat(result.get()).isEqualTo(RoaringBitmap32.bitmapOf(positions));
+ }
+
+ private static FileIndexOptions bitmapOptions() {
+ Options options = new Options();
+ options.set("file-index.bitmap.columns", "id");
+ return new FileIndexOptions(new CoreOptions(options));
+ }
+
+ private static VectorizedBundleRecords vectorizedRows(Integer... values) {
+ HeapIntVector vector = new HeapIntVector(values.length);
+ for (int i = 0; i < values.length; i++) {
+ if (values[i] == null) {
+ vector.setNullAt(i);
+ } else {
+ vector.setInt(i, values[i]);
+ }
+ }
+ VectorizedColumnBatch batch = new VectorizedColumnBatch(new
ColumnVector[] {vector});
+ batch.setNumRows(values.length);
+ return new VectorizedBundleRecords(batch, null);
+ }
+
private static RowDataFileWriter createWriter(
FileIO fileIO,
RowType rowType,
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 3037987232..d9b3d70011 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,33 +18,47 @@
package org.apache.paimon.format.mosaic;
+import org.apache.paimon.CoreOptions;
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.fileindex.FileIndexFormat;
+import org.apache.paimon.fileindex.FileIndexOptions;
+import org.apache.paimon.fileindex.FileIndexReader;
+import org.apache.paimon.fileindex.bitmap.BitmapIndexResult;
import org.apache.paimon.format.BundleFormatWriter;
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.ByteArraySeekableStream;
import org.apache.paimon.fs.Path;
import org.apache.paimon.fs.SeekableInputStream;
import org.apache.paimon.fs.local.LocalFileIO;
+import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.io.DataFileRecordReader;
+import org.apache.paimon.io.FileWriterContext;
+import org.apache.paimon.io.RollingFileWriter;
+import org.apache.paimon.io.RowDataFileWriter;
import org.apache.paimon.io.VectorizedBundleRecords;
+import org.apache.paimon.manifest.FileSource;
import org.apache.paimon.options.MemorySize;
import org.apache.paimon.options.Options;
+import org.apache.paimon.predicate.FieldRef;
import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.predicate.PredicateBuilder;
import org.apache.paimon.reader.FileRecordIterator;
import org.apache.paimon.reader.FileRecordReader;
import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.reader.VectorizedRecordIterator;
+import org.apache.paimon.statistics.SimpleColStatsCollector;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
+import org.apache.paimon.utils.LongCounter;
import org.apache.paimon.utils.RoaringBitmap32;
import org.apache.arrow.memory.RootAllocator;
@@ -145,6 +159,97 @@ class MosaicReaderWriterTest {
assertThat(result.get(1).getString(1).toString()).isEqualTo("world");
}
+ @Test
+ void testColumnBatchRewritePreservesVinBitmap() throws IOException {
+ RowType rowType = RowType.builder().field("vin",
DataTypes.STRING()).build();
+ Path sourcePath = newPath();
+ Path targetPath = newPath();
+ BinaryString a = BinaryString.fromString("A");
+ BinaryString b = BinaryString.fromString("B");
+ writeRows(rowType, sourcePath, GenericRow.of(a), GenericRow.of(b),
GenericRow.of(a));
+
+ LocalFileIO fileIO = new LocalFileIO();
+ MosaicFileFormat format = createFormat("vin");
+ Options options = new Options();
+ options.set("file-index.bitmap.columns", "vin");
+ FormatWriterFactory factory =
+ (out, compression) ->
+ new MosaicRecordsWriter(
+ out,
+ rowType,
+ new FileFormatFactory.FormatContext(new
Options(), 2, 2),
+ Collections.singletonList("vin"),
+ null) {
+ @Override
+ public void addElement(InternalRow row) {
+ throw new AssertionError(
+ "Indexed column batches must not fall
back to row writes");
+ }
+ };
+ RowDataFileWriter writer =
+ new RowDataFileWriter(
+ fileIO,
+ new FileWriterContext(
+ factory,
+ RollingFileWriter.createStatsProducer(
+ format,
+ rowType,
+
SimpleColStatsCollector.createFullStatsFactories(1)),
+ "zstd"),
+ targetPath,
+ rowType,
+ 0L,
+ LongCounter::new,
+ new FileIndexOptions(new CoreOptions(options)),
+ FileSource.COMPACT,
+ false,
+ false,
+ false,
+ null,
+ null,
+ null);
+ try (RecordReader<InternalRow> reader =
+ format.createReaderFactory(rowType, rowType, null)
+ .createReader(
+ new FormatReaderContext(
+ fileIO,
+ sourcePath,
+ fileIO.getFileSize(sourcePath),
+ null,
+ null))) {
+ RecordReader.RecordIterator<InternalRow> records;
+ while ((records = reader.readBatch()) != null) {
+ try {
+
assertThat(records).isInstanceOf(VectorizedRecordIterator.class);
+ writer.writeBundle(
+ new VectorizedBundleRecords(
+ ((VectorizedRecordIterator)
records).batch(), null));
+ } finally {
+ records.releaseBatch();
+ }
+ }
+ }
+ writer.close();
+
+ assertThat(readAll(rowType, rowType, targetPath, null))
+ .extracting(row -> row.getString(0))
+ .containsExactly(a, b, a);
+ DataFileMeta file = writer.result();
+ assertThat(file.rowCount()).isEqualTo(3);
+ assertThat(file.valueStats().minValues().getString(0)).isEqualTo(a);
+ assertThat(file.valueStats().maxValues().getString(0)).isEqualTo(b);
+ try (FileIndexFormat.Reader indexReader =
+ FileIndexFormat.createReader(
+ new ByteArraySeekableStream(file.embeddedIndex()),
rowType)) {
+ FileIndexReader index =
indexReader.readColumnIndex("vin").iterator().next();
+ FieldRef field = new FieldRef(0, "vin", DataTypes.STRING());
+ assertThat(((BitmapIndexResult) index.visitEqual(field, a)).get())
+ .isEqualTo(RoaringBitmap32.bitmapOf(0, 2));
+ assertThat(((BitmapIndexResult) index.visitEqual(field, b)).get())
+ .isEqualTo(RoaringBitmap32.bitmapOf(1));
+ }
+ }
+
@Test
void testPublicWriterFactoryFallsBackForCrossRootArrowBundle() throws
IOException {
RowType rowType = RowType.builder().field("value",
DataTypes.INT()).build();
