[core] Read primary-key vector results by position

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.
+ *
+ * <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;

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);
+    }
+}

paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java

@@ -194,22 +194,7 @@ public RecordReader<InternalRow> createReader(
                 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 RecordReader<InternalRow> createReader(
                             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,15 +281,21 @@ private ReaderSupplier<InternalRow> createFileReader(
                 dvFactories == null ? null : dvFactories.get(file.fileName());
         return () ->
                 createFileReader(
-                        partition, file, dataFilePathFactory, formatReaderMapping, dvFactory);
+                        partition,
+                        file,
+                        dataFilePathFactory,
+                        formatReaderMapping,
+                        dvFactory,
+                        selectedPositions);
     }
 
     private FileRecordReader<InternalRow> createFileReader(
             BinaryRow partition,
             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 @@ private FileRecordReader<InternalRow> 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(

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),

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();
+    }
+}