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 d41ec7357f [core] Read primary-key vector results by position (#8576)
d41ec7357f is described below
commit d41ec7357f0c4de9f6ee5c83ee35b4be3c72b4b3
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Jul 12 20:43:32 2026 +0800
[core] Read primary-key vector results by position (#8576)
Adds the physical-row materialization path for primary-key vector search
results. Bucket-local ANN and exact hits can now reuse `IndexedSplit`
for distributed transport and be read from their source data files while
preserving physical positions and scores.
---
.../apache/paimon/globalindex/IndexedSplit.java | 8 +-
.../operation/PrimaryKeyIndexedSplitRead.java | 121 +++++++++++++
.../apache/paimon/operation/RawFileSplitRead.java | 88 ++++++---
.../paimon/table/source/KeyValueTableRead.java | 2 +
.../source/PrimaryKeyVectorPositionReader.java | 106 +++++++++++
.../PrimaryKeyIndexedSplitReadProvider.java | 55 ++++++
.../operation/PrimaryKeyIndexedSplitReadTest.java | 98 ++++++++++
.../source/PrimaryKeyVectorPositionReaderTest.java | 197 +++++++++++++++++++++
.../PrimaryKeyVectorRawFileSplitReadTest.java | 108 +++++++++++
9 files changed, 762 insertions(+), 21 deletions(-)
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 d9b6fef28c..49802d5d4c 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.List;
import java.util.Objects;
import java.util.OptionalLong;
-/** Indexed split for global index. */
+/**
+ * Indexed split for global index.
+ *
+ * <p>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 0000000000..bb8d82e32c
--- /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<InternalRow> {
+
+ private final RawFileSplitRead rawRead;
+
+ public PrimaryKeyIndexedSplitRead(RawFileSplitRead rawRead) {
+ this.rawRead = rawRead;
+ }
+
+ @Override
+ public SplitRead<InternalRow> forceKeepDelete() {
+ rawRead.forceKeepDelete();
+ return this;
+ }
+
+ @Override
+ public SplitRead<InternalRow> withIOManager(@Nullable IOManager ioManager)
{
+ rawRead.withIOManager(ioManager);
+ return this;
+ }
+
+ @Override
+ public SplitRead<InternalRow> withReadType(RowType readType) {
+ rawRead.withReadType(readType);
+ return this;
+ }
+
+ @Override
+ public SplitRead<InternalRow> withFilter(@Nullable Predicate predicate) {
+ rawRead.withFilter(predicate);
+ return this;
+ }
+
+ @Override
+ public RecordReader<InternalRow> 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<Integer, Float> 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<Float> scoreGetter =
+ scoreByPosition == null ? position -> Float.NaN :
scoreByPosition::get;
+ FileRecordReader<InternalRow> 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 464d9ea0e5..56680aeb83 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 class RawFileSplitRead implements
SplitRead<InternalRow> {
pathFactory.createDataFilePathFactory(partition, bucket);
List<ReaderSupplier<InternalRow>> 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 class RawFileSplitRead implements
SplitRead<InternalRow> {
dataFilePathFactory,
file,
formatReaderMappingBuilder,
- dvFactories));
+ dvFactories,
+ null));
}
return ConcatRecordReader.create(suppliers);
}
+ FileRecordReader<InternalRow> 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<String, IOExceptionSupplier<DeletionVector>> dvFactories = new
HashMap<>();
+ dvFactories.put(
+ dataFile.fileName(), () ->
dvFactory.create(dataFile.fileName()).orElse(null));
+ DataFilePathFactory dataFilePathFactory =
+ pathFactory.createDataFilePathFactory(dataSplit.partition(),
dataSplit.bucket());
+ return (FileRecordReader<InternalRow>)
+ 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<InternalRow> createFileReader(
BinaryRow partition,
DataFilePathFactory dataFilePathFactory,
DataFileMeta file,
Builder formatBuilder,
- @Nullable Map<String, IOExceptionSupplier<DeletionVector>>
dvFactories) {
+ @Nullable Map<String, IOExceptionSupplier<DeletionVector>>
dvFactories,
+ @Nullable RoaringBitmap32 selectedPositions) {
String formatIdentifier =
DataFilePathFactory.formatIdentifier(file.fileName());
long schemaId = file.schemaId();
@@ -248,7 +281,12 @@ public class RawFileSplitRead implements
SplitRead<InternalRow> {
dvFactories == null ? null : dvFactories.get(file.fileName());
return () ->
createFileReader(
- partition, file, dataFilePathFactory,
formatReaderMapping, dvFactory);
+ partition,
+ file,
+ dataFilePathFactory,
+ formatReaderMapping,
+ dvFactory,
+ selectedPositions);
}
private FileRecordReader<InternalRow> createFileReader(
@@ -256,7 +294,8 @@ public class RawFileSplitRead implements
SplitRead<InternalRow> {
DataFileMeta file,
DataFilePathFactory dataFilePathFactory,
FormatReaderMapping formatReaderMapping,
- IOExceptionSupplier<DeletionVector> dvFactory)
+ IOExceptionSupplier<DeletionVector> dvFactory,
+ @Nullable RoaringBitmap32 selectedPositions)
throws IOException {
FileIndexResult fileIndexResult = null;
DeletionVector deletionVector = dvFactory == null ? null :
dvFactory.get();
@@ -280,6 +319,15 @@ public class RawFileSplitRead implements
SplitRead<InternalRow> {
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 fda7d70ffd..48b08df0c1 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.schema.TableSchema;
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 final class KeyValueTableRead extends
AbstractDataTableRead {
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 0000000000..6470fc26eb
--- /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<InternalRow> {
+
+ private final FileRecordReader<InternalRow> reader;
+ private final RoaringBitmap32 rowPositions;
+ private final IntFunction<Float> scoreGetter;
+ private final int lastPosition;
+ private boolean exhausted;
+
+ public PrimaryKeyVectorPositionReader(
+ FileRecordReader<InternalRow> reader,
+ RoaringBitmap32 rowPositions,
+ IntFunction<Float> 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<InternalRow> readBatch() throws IOException {
+ if (exhausted) {
+ return null;
+ }
+ FileRecordIterator<InternalRow> batch = reader.readBatch();
+ if (batch == null) {
+ return null;
+ }
+ FileRecordIterator<InternalRow> selected =
batch.selection(rowPositions);
+ return new ScoreRecordIterator<InternalRow>() {
+
+ 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 0000000000..c96e995a34
--- /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<PrimaryKeyIndexedSplitRead> splitRead;
+
+ public PrimaryKeyIndexedSplitReadProvider(
+ Supplier<RawFileSplitRead> 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<PrimaryKeyIndexedSplitRead> 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 0000000000..3d4ef2413f
--- /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<InternalRow> fileReader =
mock(FileRecordReader.class);
+ when(rawRead.createFileReader(eq(dataSplit),
any(RoaringBitmap32.class)))
+ .thenReturn(fileReader);
+
+ RecordReader<InternalRow> reader =
+ new PrimaryKeyIndexedSplitRead(rawRead).createReader(split);
+
+ assertThat(reader).isInstanceOf(PrimaryKeyVectorPositionReader.class);
+ ArgumentCaptor<RoaringBitmap32> 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 0000000000..c99f1d0c41
--- /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<InternalRow> 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<InternalRow> firstBatch = reader.readBatch();
+ assertThat(firstBatch.next().getInt(0)).isEqualTo(1);
+ assertThat(firstBatch.returnedRowId()).isEqualTo(1);
+ assertThat(firstBatch.next()).isNull();
+
+ ScoreRecordIterator<InternalRow> 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<InternalRow> {
+
+ 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<InternalRow> 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<InternalRow> {
+
+ 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 0000000000..73b75cfec5
--- /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<InternalRow> reader =
+ table.newRead().withLimit(2).createReader(vectorSplit)) {
+ ScoreRecordIterator<InternalRow> batch =
+ (ScoreRecordIterator<InternalRow>) 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();
+ }
+ }
+}
