diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/ColumnarRowIterator.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/ColumnarRowIterator.java index 49a0fe5e714f..1aa9a88aad17 100644 --- a/paimon-common/src/main/java/org/apache/paimon/data/columnar/ColumnarRowIterator.java +++ b/paimon-common/src/main/java/org/apache/paimon/data/columnar/ColumnarRowIterator.java @@ -115,6 +115,11 @@ public ColumnarRowIterator copy(ColumnVector[] vectors) { public ColumnarRowIterator mapping( @Nullable PartitionInfo partitionInfo, @Nullable int[] indexMapping) { + // Preserve concrete iterator capabilities when a full identity mapping changes no columns. + if (partitionInfo == null + && isFullIdentityMapping(indexMapping, row.batch().columns.length)) { + return this; + } if (partitionInfo != null || indexMapping != null) { VectorizedColumnBatch vectorizedColumnBatch = row.batch(); ColumnVector[] vectors = vectorizedColumnBatch.columns; @@ -129,6 +134,18 @@ public ColumnarRowIterator mapping( return this; } + private static boolean isFullIdentityMapping(@Nullable int[] indexMapping, int columnCount) { + if (indexMapping == null || indexMapping.length != columnCount) { + return false; + } + for (int i = 0; i < indexMapping.length; i++) { + if (indexMapping[i] != i) { + return false; + } + } + return true; + } + public ColumnarRowIterator assignRowTracking( Long firstRowId, Long snapshotId, Map meta) { VectorizedColumnBatch vectorizedColumnBatch = row.batch(); diff --git a/paimon-common/src/main/java/org/apache/paimon/reader/BundleRecordIterator.java b/paimon-common/src/main/java/org/apache/paimon/reader/BundleRecordIterator.java new file mode 100644 index 000000000000..077073cd3b6f --- /dev/null +++ b/paimon-common/src/main/java/org/apache/paimon/reader/BundleRecordIterator.java @@ -0,0 +1,35 @@ +/* + * 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.reader; + +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.io.BundleRecords; + +/** A record iterator whose complete remaining logical rows can be exposed as a bundle. */ +public interface BundleRecordIterator extends RecordReader.RecordIterator { + + /** + * Returns all unconsumed logical rows as bundle records. + * + *

This method must be called before the first {@link #next()} invocation. Once called, + * {@code next()} must not be used for the same batch. The returned bundle must be consumed + * synchronously before {@link #releaseBatch()}. + */ + BundleRecords bundleRecords(); +} diff --git a/paimon-common/src/test/java/org/apache/paimon/data/columnar/ColumnarRowIteratorTest.java b/paimon-common/src/test/java/org/apache/paimon/data/columnar/ColumnarRowIteratorTest.java index 8ab926a193cd..3436efa88594 100644 --- a/paimon-common/src/test/java/org/apache/paimon/data/columnar/ColumnarRowIteratorTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/data/columnar/ColumnarRowIteratorTest.java @@ -18,8 +18,16 @@ package org.apache.paimon.data.columnar; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.BinaryRowWriter; +import org.apache.paimon.data.PartitionInfo; import org.apache.paimon.data.columnar.heap.HeapIntVector; import org.apache.paimon.fs.Path; +import org.apache.paimon.io.BundleRecords; +import org.apache.paimon.io.VectorizedBundleRecords; +import org.apache.paimon.reader.BundleRecordIterator; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; import org.apache.paimon.utils.LongIterator; import org.junit.jupiter.api.Test; @@ -62,4 +70,95 @@ public void testRowIterator() { assertThat(rowIterator.returnedPosition()).isEqualTo(positions[rowIterator.index - 1]); } } + + @Test + public void testIdentityMappingKeepsBundleIterator() { + HeapIntVector heapIntVector = new HeapIntVector(1); + VectorizedColumnBatch vectorizedColumnBatch = + new VectorizedColumnBatch(new ColumnVector[] {heapIntVector}); + vectorizedColumnBatch.setNumRows(1); + ColumnarRowIterator rowIterator = + new TestingBundleColumnarRowIterator(vectorizedColumnBatch); + rowIterator.reset(0); + + assertThat(rowIterator.mapping(null, new int[] {0})).isSameAs(rowIterator); + } + + @Test + public void testNonIdentityMappingCopiesAndReordersIterator() { + HeapIntVector first = new HeapIntVector(1); + first.setInt(0, 1); + HeapIntVector second = new HeapIntVector(1); + second.setInt(0, 2); + VectorizedColumnBatch vectorizedColumnBatch = + new VectorizedColumnBatch(new ColumnVector[] {first, second}); + vectorizedColumnBatch.setNumRows(1); + ColumnarRowIterator rowIterator = + new ColumnarRowIterator( + new Path("test"), new ColumnarRow(vectorizedColumnBatch), null); + rowIterator.reset(0); + + ColumnarRowIterator mapped = rowIterator.mapping(null, new int[] {1, 0}); + + assertThat(mapped).isNotSameAs(rowIterator); + assertThat(mapped.next().getInt(0)).isEqualTo(2); + assertThat(mapped.next()).isNull(); + } + + @Test + public void testIdentityPrefixMappingStillCopiesIterator() { + HeapIntVector first = new HeapIntVector(1); + first.setInt(0, 1); + HeapIntVector second = new HeapIntVector(1); + second.setInt(0, 2); + VectorizedColumnBatch vectorizedColumnBatch = + new VectorizedColumnBatch(new ColumnVector[] {first, second}); + vectorizedColumnBatch.setNumRows(1); + ColumnarRowIterator rowIterator = + new ColumnarRowIterator( + new Path("test"), new ColumnarRow(vectorizedColumnBatch), null); + rowIterator.reset(0); + + ColumnarRowIterator mapped = rowIterator.mapping(null, new int[] {0}); + + assertThat(mapped).isNotSameAs(rowIterator); + assertThat(mapped.next().getFieldCount()).isEqualTo(1); + } + + @Test + public void testPartitionMappingStillCopiesIterator() { + HeapIntVector dataVector = new HeapIntVector(1); + dataVector.setInt(0, 1); + VectorizedColumnBatch vectorizedColumnBatch = + new VectorizedColumnBatch(new ColumnVector[] {dataVector}); + vectorizedColumnBatch.setNumRows(1); + ColumnarRowIterator rowIterator = + new TestingBundleColumnarRowIterator(vectorizedColumnBatch); + rowIterator.reset(0); + + BinaryRow partition = new BinaryRow(1); + BinaryRowWriter partitionWriter = new BinaryRowWriter(partition); + partitionWriter.writeInt(0, 9); + partitionWriter.complete(); + PartitionInfo partitionInfo = + new PartitionInfo(new int[] {-1, 1, 0}, RowType.of(DataTypes.INT()), partition); + + ColumnarRowIterator mapped = rowIterator.mapping(partitionInfo, new int[] {0}); + + assertThat(mapped).isNotSameAs(rowIterator); + assertThat(mapped.next().getInt(0)).isEqualTo(9); + } + + private static class TestingBundleColumnarRowIterator extends ColumnarRowIterator + implements BundleRecordIterator { + + private TestingBundleColumnarRowIterator(VectorizedColumnBatch batch) { + super(new Path("test"), new ColumnarRow(batch), null); + } + + @Override + public BundleRecords bundleRecords() { + return new VectorizedBundleRecords(batch(), null); + } + } } diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/BaseAppendFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/BaseAppendFileStoreWrite.java index 2e734552bd71..eee777b0ddb7 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/BaseAppendFileStoreWrite.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/BaseAppendFileStoreWrite.java @@ -20,6 +20,7 @@ import org.apache.paimon.AppendOnlyFileStore; import org.apache.paimon.CoreOptions; +import org.apache.paimon.annotation.VisibleForTesting; import org.apache.paimon.append.AppendOnlyWriter; import org.apache.paimon.append.cluster.Sorter; import org.apache.paimon.compact.CompactManager; @@ -34,10 +35,13 @@ import org.apache.paimon.io.BundleRecords; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.io.DataFilePathFactory; +import org.apache.paimon.io.RollingFileWriter; import org.apache.paimon.io.RowDataRollingFileWriter; import org.apache.paimon.manifest.FileSource; import org.apache.paimon.metrics.MetricRegistry; import org.apache.paimon.operation.metrics.BlobFetchMetrics; +import org.apache.paimon.reader.BundleRecordIterator; +import org.apache.paimon.reader.RecordReader; import org.apache.paimon.reader.RecordReaderIterator; import org.apache.paimon.statistics.SimpleColStatsCollector; import org.apache.paimon.types.DataField; @@ -292,7 +296,9 @@ public List compactRewrite( } } try { - rewriter.write(createFilesIterator(partition, bucket, toCompact, dvFactories)); + writeCompactReader( + readForCompact.createReader(partition, bucket, toCompact, dvFactories), + rewriter); } catch (Exception e) { collectedExceptions = e; } finally { @@ -308,6 +314,33 @@ public List compactRewrite( return rewriter.result(); } + private static void writeCompactReader( + RecordReader reader, RollingFileWriter rewriter) + throws Exception { + try { + RecordReader.RecordIterator batch; + while ((batch = reader.readBatch()) != null) { + try { + if (batch instanceof BundleRecordIterator) { + BundleRecords bundle = ((BundleRecordIterator) batch).bundleRecords(); + if (bundle.rowCount() > 0) { + rewriter.writeBundle(bundle); + } + } else { + InternalRow row; + while ((row = batch.next()) != null) { + rewriter.write(row); + } + } + } finally { + batch.releaseBatch(); + } + } + } finally { + reader.close(); + } + } + public List clusterRewrite( BinaryRow partition, int bucket, List toCluster) throws Exception { RecordReaderIterator reader = @@ -346,7 +379,8 @@ public List clusterRewrite( return rewriter.result(); } - private RowDataRollingFileWriter createRollingFileWriter( + @VisibleForTesting + RowDataRollingFileWriter createRollingFileWriter( BinaryRow partition, int bucket, Supplier seqNumCounterSupplier) { return new RowDataRollingFileWriter( fileIO, diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/BaseAppendFileStoreWriteTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/BaseAppendFileStoreWriteTest.java new file mode 100644 index 000000000000..46e9fe5ff11a --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/operation/BaseAppendFileStoreWriteTest.java @@ -0,0 +1,293 @@ +/* + * 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.operation; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.compact.CompactManager; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.data.columnar.ColumnVector; +import org.apache.paimon.data.columnar.ColumnarRow; +import org.apache.paimon.data.columnar.VectorizedColumnBatch; +import org.apache.paimon.data.columnar.VectorizedRowIterator; +import org.apache.paimon.data.columnar.heap.HeapIntVector; +import org.apache.paimon.data.columnar.heap.HeapLongVector; +import org.apache.paimon.deletionvectors.BucketedDvMaintainer; +import org.apache.paimon.deletionvectors.DeletionVector; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.io.BundleRecords; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.io.RowDataRollingFileWriter; +import org.apache.paimon.io.VectorizedBundleRecords; +import org.apache.paimon.options.Options; +import org.apache.paimon.reader.BundleRecordIterator; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.table.SpecialFields; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.FileStorePathFactory; +import org.apache.paimon.utils.IOExceptionSupplier; +import org.apache.paimon.utils.LongCounter; +import org.apache.paimon.utils.SnapshotManager; + +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.function.Function; +import java.util.function.Supplier; + +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.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests for {@link BaseAppendFileStoreWrite}. */ +class BaseAppendFileStoreWriteTest { + + @Test + void testCompactRewriteWritesBatchAsBundle() throws Exception { + TestingBundleReader reader = new TestingBundleReader(2); + RawFileSplitRead rawRead = mock(RawFileSplitRead.class); + RowDataRollingFileWriter writer = mock(RowDataRollingFileWriter.class); + List inputFiles = Collections.singletonList(mock(DataFileMeta.class)); + List outputFiles = Collections.singletonList(mock(DataFileMeta.class)); + when(rawRead.createReader( + BinaryRow.EMPTY_ROW, + 0, + inputFiles, + (Map>) null)) + .thenReturn(reader); + doAnswer( + invocation -> { + assertThat(reader.batchReleased).isFalse(); + assertThat(reader.closed).isFalse(); + return null; + }) + .when(writer) + .writeBundle(any()); + when(writer.result()).thenReturn(outputFiles); + TestingAppendWrite write = new TestingAppendWrite(rawRead, writer); + + assertThat(write.compactRewrite(BinaryRow.EMPTY_ROW, 0, null, inputFiles)) + .isSameAs(outputFiles); + + ArgumentCaptor bundle = ArgumentCaptor.forClass(BundleRecords.class); + verify(writer).writeBundle(bundle.capture()); + assertThat(bundle.getValue().rowCount()).isEqualTo(2); + verify(writer, never()).write(any(InternalRow.class)); + verify(writer).close(); + assertThat(reader.nextCalls).isZero(); + assertThat(reader.batchReleased).isTrue(); + assertThat(reader.closed).isTrue(); + } + + @Test + void testCompactRewriteReleasesBorrowedBatchAfterBundleFailure() throws Exception { + IOException failure = new IOException("bundle failure"); + TestingBundleReader reader = new TestingBundleReader(2); + RawFileSplitRead rawRead = mock(RawFileSplitRead.class); + RowDataRollingFileWriter writer = mock(RowDataRollingFileWriter.class); + List inputFiles = Collections.singletonList(mock(DataFileMeta.class)); + when(rawRead.createReader( + BinaryRow.EMPTY_ROW, + 0, + inputFiles, + (Map>) null)) + .thenReturn(reader); + doThrow(failure).when(writer).writeBundle(any()); + TestingAppendWrite write = new TestingAppendWrite(rawRead, writer); + + assertThatThrownBy(() -> write.compactRewrite(BinaryRow.EMPTY_ROW, 0, null, inputFiles)) + .isSameAs(failure); + verify(writer).close(); + assertThat(reader.batchReleased).isTrue(); + assertThat(reader.closed).isTrue(); + } + + @Test + void testCompactRewriteConsumesStatefulVectorizedIteratorAsRows() throws Exception { + HeapIntVector idVector = new HeapIntVector(2); + idVector.setInt(0, 10); + idVector.setInt(1, 20); + HeapLongVector rowIdVector = new HeapLongVector(2); + rowIdVector.fillWithNulls(); + VectorizedColumnBatch batch = + new VectorizedColumnBatch(new ColumnVector[] {idVector, rowIdVector}); + batch.setNumRows(2); + VectorizedRowIterator tracked = + new VectorizedRowIterator( + new Path("file:///tracked.data"), new ColumnarRow(batch), null); + tracked.reset(5L); + tracked.assignRowTracking( + 100L, 200L, Collections.singletonMap(SpecialFields.ROW_ID.name(), 1)); + + RecordReader reader = singleBatchReader(tracked); + RawFileSplitRead rawRead = mock(RawFileSplitRead.class); + RowDataRollingFileWriter writer = mock(RowDataRollingFileWriter.class); + List inputFiles = Collections.singletonList(mock(DataFileMeta.class)); + List outputFiles = Collections.singletonList(mock(DataFileMeta.class)); + when(rawRead.createReader( + BinaryRow.EMPTY_ROW, + 0, + inputFiles, + (Map>) null)) + .thenReturn(reader); + when(writer.result()).thenReturn(outputFiles); + List> written = new ArrayList<>(); + doAnswer( + invocation -> { + InternalRow row = invocation.getArgument(0); + written.add(Arrays.asList((long) row.getInt(0), row.getLong(1))); + return null; + }) + .when(writer) + .write(any(InternalRow.class)); + + TestingAppendWrite write = new TestingAppendWrite(rawRead, writer); + assertThat(write.compactRewrite(BinaryRow.EMPTY_ROW, 0, null, inputFiles)) + .isSameAs(outputFiles); + + assertThat(written).containsExactly(Arrays.asList(10L, 105L), Arrays.asList(20L, 106L)); + verify(writer, never()).writeBundle(any()); + verify(writer).close(); + } + + private static RecordReader singleBatchReader( + RecordReader.RecordIterator iterator) { + return new RecordReader() { + private boolean returned; + + @Override + public RecordIterator readBatch() { + if (returned) { + return null; + } + returned = true; + return iterator; + } + + @Override + public void close() {} + }; + } + + private static class TestingBundleReader implements RecordReader { + + private final VectorizedColumnBatch batch; + private boolean returned; + private boolean batchReleased; + private boolean closed; + private int nextCalls; + + private TestingBundleReader(int rowCount) { + HeapIntVector vector = new HeapIntVector(rowCount); + for (int i = 0; i < rowCount; i++) { + vector.setInt(i, i); + } + batch = new VectorizedColumnBatch(new ColumnVector[] {vector}); + batch.setNumRows(rowCount); + } + + @Override + public RecordIterator readBatch() { + if (returned) { + return null; + } + returned = true; + return new BundleRecordIterator() { + @Override + public BundleRecords bundleRecords() { + return new VectorizedBundleRecords(batch, null); + } + + @Override + public InternalRow next() { + nextCalls++; + throw new AssertionError("Bundle compaction must not consume rows"); + } + + @Override + public void releaseBatch() { + batchReleased = true; + } + }; + } + + @Override + public void close() { + closed = true; + } + } + + private static class TestingAppendWrite extends BaseAppendFileStoreWrite { + + private final RowDataRollingFileWriter writer; + + private TestingAppendWrite( + RawFileSplitRead readForCompact, RowDataRollingFileWriter writer) { + super( + mock(FileIO.class), + readForCompact, + 0, + RowType.of(DataTypes.INT()), + RowType.of(), + mock(FileStorePathFactory.class), + mock(SnapshotManager.class), + mock(FileStoreScan.class), + new CoreOptions(new Options()), + null, + "test"); + this.writer = writer; + } + + @Override + RowDataRollingFileWriter createRollingFileWriter( + BinaryRow partition, int bucket, Supplier seqNumCounterSupplier) { + return writer; + } + + @Override + protected CompactManager getCompactManager( + BinaryRow partition, + int bucket, + List restoredFiles, + ExecutorService compactExecutor, + BucketedDvMaintainer dvMaintainer) { + return null; + } + + @Override + protected Function, Boolean> createWriterCleanChecker() { + return null; + } + } +}