diff --git a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java index d9b6fef28c46..49802d5d4c0c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java +++ b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java @@ -37,7 +37,13 @@ import java.util.Objects; import java.util.OptionalLong; -/** Indexed split for global index. */ +/** + * Indexed split for global index. + * + *

Row ranges use the coordinate system of the table read path. Data Evolution reads them as + * stable row IDs, while primary-key raw reads use them as physical positions in the split's single + * data file. Scores, when present, follow the expanded row-range order. + */ public class IndexedSplit implements Split { private static final long serialVersionUID = 1L; diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitRead.java b/paimon-core/src/main/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitRead.java new file mode 100644 index 000000000000..bb8d82e32c42 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitRead.java @@ -0,0 +1,121 @@ +/* + * 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.data.InternalRow; +import org.apache.paimon.disk.IOManager; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.PrimaryKeyVectorPositionReader; +import org.apache.paimon.table.source.Split; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringBitmap32; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.function.IntFunction; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Reads an {@link IndexedSplit} as physical positions in one primary-key data file. */ +public class PrimaryKeyIndexedSplitRead implements SplitRead { + + private final RawFileSplitRead rawRead; + + public PrimaryKeyIndexedSplitRead(RawFileSplitRead rawRead) { + this.rawRead = rawRead; + } + + @Override + public SplitRead forceKeepDelete() { + rawRead.forceKeepDelete(); + return this; + } + + @Override + public SplitRead withIOManager(@Nullable IOManager ioManager) { + rawRead.withIOManager(ioManager); + return this; + } + + @Override + public SplitRead withReadType(RowType readType) { + rawRead.withReadType(readType); + return this; + } + + @Override + public SplitRead withFilter(@Nullable Predicate predicate) { + rawRead.withFilter(predicate); + return this; + } + + @Override + public RecordReader createReader(Split split) throws IOException { + IndexedSplit indexedSplit = (IndexedSplit) split; + DataSplit dataSplit = indexedSplit.dataSplit(); + checkArgument( + dataSplit.dataFiles().size() == 1, + "Indexed split for a primary-key table must contain exactly one data file."); + DataFileMeta dataFile = dataSplit.dataFiles().get(0); + RoaringBitmap32 rowPositions = new RoaringBitmap32(); + float[] scores = indexedSplit.scores(); + Map scoreByPosition = scores == null ? null : new HashMap<>(); + int scoreIndex = 0; + for (Range range : indexedSplit.rowRanges()) { + checkArgument( + range.from >= 0 && range.to < dataFile.rowCount(), + "Indexed row range %s is outside data file %s row count %s.", + range, + dataFile.fileName(), + dataFile.rowCount()); + checkArgument( + range.to <= Integer.MAX_VALUE, + "Indexed row range %s exceeds supported physical positions.", + range); + for (long position = range.from; position <= range.to; position++) { + rowPositions.add((int) position); + if (scoreByPosition != null) { + checkArgument( + scoreIndex < scores.length, + "Scores length does not match row ranges in indexed split."); + scoreByPosition.put((int) position, scores[scoreIndex++]); + } + } + } + checkArgument(!rowPositions.isEmpty(), "Indexed split must select at least one row."); + if (scores != null) { + checkArgument( + scoreIndex == scores.length, + "Scores length does not match row ranges in indexed split."); + } + IntFunction scoreGetter = + scoreByPosition == null ? position -> Float.NaN : scoreByPosition::get; + FileRecordReader reader = rawRead.createFileReader(dataSplit, rowPositions); + return new PrimaryKeyVectorPositionReader(reader, rowPositions, scoreGetter); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java index 464d9ea0e5d7..56680aeb835e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java @@ -194,22 +194,7 @@ public RecordReader createReader( pathFactory.createDataFilePathFactory(partition, bucket); List> suppliers = new ArrayList<>(); - Builder formatReaderMappingBuilder = - new Builder( - formatDiscover, - readRowType.getFields(), - schema -> { - if (rowTrackingEnabled) { - // maybe file has no row id and sequence number, but in manifest - // entry - return rowTypeWithRowTracking(schema.logicalRowType(), true, true) - .getFields(); - } - return schema.fields(); - }, - filters, - topN, - limit); + Builder formatReaderMappingBuilder = createFormatReaderMappingBuilder(); for (DataFileMeta file : files) { suppliers.add( @@ -218,18 +203,66 @@ public RecordReader createReader( dataFilePathFactory, file, formatReaderMappingBuilder, - dvFactories)); + dvFactories, + null)); } return ConcatRecordReader.create(suppliers); } + FileRecordReader createFileReader( + DataSplit dataSplit, RoaringBitmap32 selectedPositions) throws IOException { + DataFileMeta dataFile = dataSplit.dataFiles().get(0); + DeletionVector.Factory dvFactory = + DeletionVector.factory( + fileIO, dataSplit.dataFiles(), dataSplit.deletionFiles().orElse(null)); + Map> dvFactories = new HashMap<>(); + dvFactories.put( + dataFile.fileName(), () -> dvFactory.create(dataFile.fileName()).orElse(null)); + DataFilePathFactory dataFilePathFactory = + pathFactory.createDataFilePathFactory(dataSplit.partition(), dataSplit.bucket()); + return (FileRecordReader) + createFileReader( + dataSplit.partition(), + dataFilePathFactory, + dataFile, + // The caller has already selected the rows. Applying a regular + // TopN or limit before position filtering can drop hits. + createFormatReaderMappingBuilder(null, null), + dvFactories, + selectedPositions) + .get(); + } + + private Builder createFormatReaderMappingBuilder() { + return createFormatReaderMappingBuilder(topN, limit); + } + + private Builder createFormatReaderMappingBuilder( + @Nullable TopN pushDownTopN, @Nullable Integer pushDownLimit) { + return new Builder( + formatDiscover, + readRowType.getFields(), + schema -> { + if (rowTrackingEnabled) { + // maybe file has no row id and sequence number, but in manifest entry + return rowTypeWithRowTracking(schema.logicalRowType(), true, true) + .getFields(); + } + return schema.fields(); + }, + filters, + pushDownTopN, + pushDownLimit); + } + private ReaderSupplier createFileReader( BinaryRow partition, DataFilePathFactory dataFilePathFactory, DataFileMeta file, Builder formatBuilder, - @Nullable Map> dvFactories) { + @Nullable Map> dvFactories, + @Nullable RoaringBitmap32 selectedPositions) { String formatIdentifier = DataFilePathFactory.formatIdentifier(file.fileName()); long schemaId = file.schemaId(); @@ -248,7 +281,12 @@ private ReaderSupplier createFileReader( dvFactories == null ? null : dvFactories.get(file.fileName()); return () -> createFileReader( - partition, file, dataFilePathFactory, formatReaderMapping, dvFactory); + partition, + file, + dataFilePathFactory, + formatReaderMapping, + dvFactory, + selectedPositions); } private FileRecordReader createFileReader( @@ -256,7 +294,8 @@ private FileRecordReader createFileReader( DataFileMeta file, DataFilePathFactory dataFilePathFactory, FormatReaderMapping formatReaderMapping, - IOExceptionSupplier dvFactory) + IOExceptionSupplier dvFactory, + @Nullable RoaringBitmap32 selectedPositions) throws IOException { FileIndexResult fileIndexResult = null; DeletionVector deletionVector = dvFactory == null ? null : dvFactory.get(); @@ -280,6 +319,15 @@ private FileRecordReader createFileReader( if (fileIndexResult instanceof BitmapIndexResult) { selection = ((BitmapIndexResult) fileIndexResult).get(); } + if (selectedPositions != null) { + selection = + selection == null + ? selectedPositions.clone() + : RoaringBitmap32.and(selection, selectedPositions); + if (selection.isEmpty()) { + return new EmptyFileRecordReader<>(); + } + } FormatReaderContext formatReaderContext = new FormatReaderContext( diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java index fda7d70ffdf6..48b08df0c1dc 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java @@ -33,6 +33,7 @@ import org.apache.paimon.table.source.splitread.IncrementalChangelogReadProvider; import org.apache.paimon.table.source.splitread.IncrementalDiffReadProvider; import org.apache.paimon.table.source.splitread.MergeFileSplitReadProvider; +import org.apache.paimon.table.source.splitread.PrimaryKeyIndexedSplitReadProvider; import org.apache.paimon.table.source.splitread.PrimaryKeyTableRawFileSplitReadProvider; import org.apache.paimon.table.source.splitread.SplitReadProvider; import org.apache.paimon.types.RowType; @@ -67,6 +68,7 @@ public KeyValueTableRead( super(schema); this.readProviders = Arrays.asList( + new PrimaryKeyIndexedSplitReadProvider(batchRawReadSupplier, this::config), new PrimaryKeyTableRawFileSplitReadProvider( batchRawReadSupplier, this::config), new MergeFileSplitReadProvider(mergeReadSupplier, this::config), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java new file mode 100644 index 000000000000..6470fc26ebd7 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java @@ -0,0 +1,106 @@ +/* + * 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.table.source; + +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.reader.ScoreRecordIterator; +import org.apache.paimon.utils.RoaringBitmap32; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.function.IntFunction; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Reads selected physical file positions and exposes their vector-search scores. */ +public class PrimaryKeyVectorPositionReader implements RecordReader { + + private final FileRecordReader reader; + private final RoaringBitmap32 rowPositions; + private final IntFunction scoreGetter; + private final int lastPosition; + private boolean exhausted; + + public PrimaryKeyVectorPositionReader( + FileRecordReader reader, + RoaringBitmap32 rowPositions, + IntFunction scoreGetter) { + checkArgument(!rowPositions.isEmpty(), "Selected row positions must not be empty."); + this.reader = reader; + this.rowPositions = rowPositions.clone(); + this.scoreGetter = scoreGetter; + this.lastPosition = rowPositions.last(); + } + + @Nullable + @Override + public ScoreRecordIterator readBatch() throws IOException { + if (exhausted) { + return null; + } + FileRecordIterator batch = reader.readBatch(); + if (batch == null) { + return null; + } + FileRecordIterator selected = batch.selection(rowPositions); + return new ScoreRecordIterator() { + + private long returnedPosition = -1; + private float returnedScore = Float.NaN; + + @Override + public float returnedScore() { + return returnedScore; + } + + @Override + public long returnedRowId() { + return returnedPosition; + } + + @Nullable + @Override + public InternalRow next() throws IOException { + InternalRow row = selected.next(); + if (row != null) { + returnedPosition = selected.returnedPosition(); + returnedScore = scoreGetter.apply((int) returnedPosition); + if (returnedPosition >= lastPosition) { + exhausted = true; + } + } + return row; + } + + @Override + public void releaseBatch() { + selected.releaseBatch(); + } + }; + } + + @Override + public void close() throws IOException { + reader.close(); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyIndexedSplitReadProvider.java b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyIndexedSplitReadProvider.java new file mode 100644 index 000000000000..c96e995a349d --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyIndexedSplitReadProvider.java @@ -0,0 +1,55 @@ +/* + * 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.table.source.splitread; + +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.operation.PrimaryKeyIndexedSplitRead; +import org.apache.paimon.operation.RawFileSplitRead; +import org.apache.paimon.table.source.Split; +import org.apache.paimon.utils.LazyField; + +import java.util.function.Supplier; + +/** Provides physical-position reads for primary-key {@link IndexedSplit}s. */ +public class PrimaryKeyIndexedSplitReadProvider implements SplitReadProvider { + + private final LazyField splitRead; + + public PrimaryKeyIndexedSplitReadProvider( + Supplier supplier, SplitReadConfig splitReadConfig) { + this.splitRead = + new LazyField<>( + () -> { + PrimaryKeyIndexedSplitRead read = + new PrimaryKeyIndexedSplitRead(supplier.get()); + splitReadConfig.config(read); + return read; + }); + } + + @Override + public boolean match(Split split, Context context) { + return !context.forceKeepDelete() && split instanceof IndexedSplit; + } + + @Override + public LazyField get() { + return splitRead; + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitReadTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitReadTest.java new file mode 100644 index 000000000000..3d4ef2413f5e --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitReadTest.java @@ -0,0 +1,98 @@ +/* + * 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.data.BinaryRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.PrimaryKeyVectorPositionReader; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringBitmap32; + +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import java.util.Arrays; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests for {@link PrimaryKeyIndexedSplitRead}. */ +class PrimaryKeyIndexedSplitReadTest { + + @Test + void testConvertsRangesToPhysicalPositionsAndDelegates() throws Exception { + DataSplit dataSplit = dataSplit(); + IndexedSplit split = + new IndexedSplit( + dataSplit, + Arrays.asList(new Range(1, 1), new Range(3, 3)), + new float[] {0.5F, 0.25F}); + RawFileSplitRead rawRead = mock(RawFileSplitRead.class); + FileRecordReader fileReader = mock(FileRecordReader.class); + when(rawRead.createFileReader(eq(dataSplit), any(RoaringBitmap32.class))) + .thenReturn(fileReader); + + RecordReader reader = + new PrimaryKeyIndexedSplitRead(rawRead).createReader(split); + + assertThat(reader).isInstanceOf(PrimaryKeyVectorPositionReader.class); + ArgumentCaptor positions = ArgumentCaptor.forClass(RoaringBitmap32.class); + verify(rawRead).createFileReader(eq(dataSplit), positions.capture()); + assertThat(positions.getValue()).isEqualTo(RoaringBitmap32.bitmapOf(1, 3)); + } + + private static DataSplit dataSplit() { + DataFileMeta dataFile = + DataFileMeta.forAppend( + "data-file", + 100, + 5, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.APPEND, + null, + null, + null, + null); + return DataSplit.builder() + .withSnapshot(3) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(1) + .withDataFiles(Collections.singletonList(dataFile)) + .build(); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java new file mode 100644 index 000000000000..c99f1d0c41f2 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java @@ -0,0 +1,197 @@ +/* + * 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.table.source; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.fs.Path; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.operation.MergeFileSplitRead; +import org.apache.paimon.operation.RawFileSplitRead; +import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.ScoreRecordIterator; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringBitmap32; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Answers.RETURNS_SELF; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +/** Tests physical row-position filtering with vector scores. */ +class PrimaryKeyVectorPositionReaderTest { + + @Test + void testSelectsPhysicalPositionsAndReturnsScores() throws Exception { + PrimaryKeyVectorPositionReader reader = + new PrimaryKeyVectorPositionReader( + new TestingFileReader(5), + RoaringBitmap32.bitmapOf(1, 3), + position -> position / 10F); + + ScoreRecordIterator batch = reader.readBatch(); + assertThat(batch.next().getInt(0)).isEqualTo(1); + assertThat(batch.returnedRowId()).isEqualTo(1); + assertThat(batch.returnedScore()).isEqualTo(0.1F); + assertThat(batch.next().getInt(0)).isEqualTo(3); + assertThat(batch.returnedRowId()).isEqualTo(3); + assertThat(batch.returnedScore()).isEqualTo(0.3F); + assertThat(batch.next()).isNull(); + assertThat(reader.readBatch()).isNull(); + } + + @Test + void testSelectedPositionsAcrossBatches() throws Exception { + PrimaryKeyVectorPositionReader reader = + new PrimaryKeyVectorPositionReader( + new TestingFileReader(5, 2), + RoaringBitmap32.bitmapOf(1, 3), + position -> position / 10F); + + ScoreRecordIterator firstBatch = reader.readBatch(); + assertThat(firstBatch.next().getInt(0)).isEqualTo(1); + assertThat(firstBatch.returnedRowId()).isEqualTo(1); + assertThat(firstBatch.next()).isNull(); + + ScoreRecordIterator secondBatch = reader.readBatch(); + assertThat(secondBatch.next().getInt(0)).isEqualTo(3); + assertThat(secondBatch.returnedScore()).isEqualTo(0.3F); + assertThat(secondBatch.next()).isNull(); + assertThat(reader.readBatch()).isNull(); + } + + @Test + void testKeyValueTableReadRoutesPhysicalSplitToRawReader() throws Exception { + RawFileSplitRead rawRead = mock(RawFileSplitRead.class, RETURNS_SELF); + IndexedSplit split = selectedSplit(); + KeyValueTableRead tableRead = + new KeyValueTableRead( + () -> mock(MergeFileSplitRead.class), + () -> rawRead, + mock(TableSchema.class)); + + assertThat(tableRead.createReader(split)) + .isInstanceOf(PrimaryKeyVectorPositionReader.class); + verify(rawRead, never()).createReader(split); + } + + private static IndexedSplit selectedSplit() { + DataFileMeta dataFile = + DataFileMeta.forAppend( + "data-file", + 100, + 5, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.APPEND, + null, + null, + null, + null); + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(3) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(1) + .withDataFiles(Collections.singletonList(dataFile)) + .build(); + return new IndexedSplit( + dataSplit, Collections.singletonList(new Range(1, 1)), new float[] {0.5F}); + } + + private static class TestingFileReader implements FileRecordReader { + + private final int rowCount; + private final int batchSize; + private int nextPosition; + + private TestingFileReader(int rowCount) { + this(rowCount, rowCount); + } + + private TestingFileReader(int rowCount, int batchSize) { + this.rowCount = rowCount; + this.batchSize = batchSize; + } + + @Override + public FileRecordIterator readBatch() { + if (nextPosition == rowCount) { + return null; + } + int startPosition = nextPosition; + nextPosition = Math.min(rowCount, nextPosition + batchSize); + return new TestingFileIterator(startPosition, nextPosition); + } + + @Override + public void close() {} + } + + private static class TestingFileIterator implements FileRecordIterator { + + private final int endPosition; + private int nextPosition; + private int returnedPosition = -1; + + private TestingFileIterator(int startPosition, int endPosition) { + this.nextPosition = startPosition; + this.endPosition = endPosition; + } + + @Override + public long returnedPosition() { + return returnedPosition; + } + + @Override + public Path filePath() { + return new Path("/test"); + } + + @Override + public InternalRow next() { + if (nextPosition == endPosition) { + return null; + } + returnedPosition = nextPosition; + return GenericRow.of(nextPosition++); + } + + @Override + public void releaseBatch() {} + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java new file mode 100644 index 000000000000..73b75cfec576 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java @@ -0,0 +1,108 @@ +/* + * 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.table.source; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.options.Options; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.reader.ScoreRecordIterator; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.SchemaUtils; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.FileStoreTableFactory; +import org.apache.paimon.table.sink.BatchTableCommit; +import org.apache.paimon.table.sink.BatchTableWrite; +import org.apache.paimon.table.sink.BatchWriteBuilder; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.utils.Range; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.util.Arrays; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests raw-file materialization of primary-key vector results. */ +class PrimaryKeyVectorRawFileSplitReadTest { + + @TempDir java.nio.file.Path tempDir; + + @Test + void testRegularLimitDoesNotDropSelectedPhysicalPositions() throws Exception { + Path tablePath = new Path(tempDir.toUri()); + Options options = new Options(); + options.set(CoreOptions.PATH, tablePath.toString()); + options.set(CoreOptions.BUCKET, 1); + Schema schema = + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("value", DataTypes.INT()) + .primaryKey("id") + .options(options.toMap()) + .build(); + TableSchema tableSchema = + SchemaUtils.forceCommit(new SchemaManager(LocalFileIO.create(), tablePath), schema); + FileStoreTable table = + FileStoreTableFactory.create(LocalFileIO.create(), tablePath, tableSchema); + + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + for (int i = 0; i < 5; i++) { + write.write(GenericRow.of(i, i * 10)); + } + commit.commit(write.prepareCommit()); + } + + DataSplit dataSplit = table.newSnapshotReader().read().dataSplits().get(0); + assertThat(dataSplit.dataFiles()).hasSize(1); + IndexedSplit vectorSplit = + new IndexedSplit( + dataSplit, + Arrays.asList(new Range(1, 1), new Range(3, 3)), + new float[] {0.5F, 0.25F}); + + try (RecordReader reader = + table.newRead().withLimit(2).createReader(vectorSplit)) { + ScoreRecordIterator batch = + (ScoreRecordIterator) reader.readBatch(); + InternalRow row = batch.next(); + assertThat(row.getInt(0)).isEqualTo(1); + assertThat(row.getInt(1)).isEqualTo(10); + assertThat(batch.returnedRowId()).isEqualTo(1); + assertThat(batch.returnedScore()).isEqualTo(0.5F); + row = batch.next(); + assertThat(row.getInt(0)).isEqualTo(3); + assertThat(row.getInt(1)).isEqualTo(30); + assertThat(batch.returnedRowId()).isEqualTo(3); + assertThat(batch.returnedScore()).isEqualTo(0.25F); + assertThat(batch.next()).isNull(); + batch.releaseBatch(); + assertThat(reader.readBatch()).isNull(); + } + } +}