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 005dbff9bf [core] Optimize vector refine raw reads (#8528)
005dbff9bf is described below
commit 005dbff9bf394f3a117370d8f82d0abc38b2a429
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Jul 9 21:51:06 2026 +0800
[core] Optimize vector refine raw reads (#8528)
Optimize vector raw refine by keeping candidate bitmaps as the internal
source of truth, reading only the columns needed for refine, and
maintaining bounded topK state while scoring. Batch vector refine now
unions candidates across queries and reads raw vectors once before
scoring each query's own candidate set.
---
.../paimon/table/source/AbstractVectorRead.java | 212 +++++++++++++++++---
.../paimon/table/source/BatchVectorReadImpl.java | 30 ++-
.../table/source/VectorSearchBuilderTest.java | 133 +++++++++++++
.../pypaimon/table/source/vector_search_read.py | 219 +++++++++++++++++----
.../pypaimon/tests/vector_search_filter_test.py | 48 ++++-
.../paimon/spark/read/SparkVectorReadImplTest.java | 10 +-
6 files changed, 573 insertions(+), 79 deletions(-)
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
index 0e3adf3b6a..93341f3f1b 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
@@ -38,6 +38,7 @@ import org.apache.paimon.index.IndexPathFactory;
import org.apache.paimon.partition.PartitionPredicate;
import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.predicate.PredicateVisitor;
import org.apache.paimon.predicate.VectorSearch;
import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.table.FileStoreTable;
@@ -60,6 +61,7 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
+import java.util.PriorityQueue;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.CompletableFuture;
@@ -79,6 +81,10 @@ public abstract class AbstractVectorRead implements
Serializable {
protected final DataField vectorColumn;
protected final Map<String, String> options;
+ private static final Comparator<long[]> WEAKEST_SCORE_FIRST =
+ Comparator.<long[]>comparingDouble(a -> Float.intBitsToFloat((int)
a[1]))
+ .thenComparing((a, b) -> Long.compare(b[0], a[0]));
+
protected AbstractVectorRead(
FileStoreTable table,
@Nullable PartitionPredicate partitionFilter,
@@ -316,11 +322,7 @@ public abstract class AbstractVectorRead implements
Serializable {
return result;
}
ScoredGlobalIndexResult candidates =
result.topK(indexedSearchLimit(indexType));
- return readRawSearch(
- candidates.results().toRangeList(),
- candidates.results(),
- globalIndexer,
- queryVector);
+ return readRawRefineSearch(candidates.results(), globalIndexer,
queryVector);
}
protected String vectorIndexType(List<IndexVectorSearchSplit> splits) {
@@ -376,48 +378,208 @@ public abstract class AbstractVectorRead implements
Serializable {
@Nullable RoaringNavigableMap64 preFilter,
String metric,
float[] queryVector) {
- RowType readType = SpecialFields.rowTypeWithRowId(table.rowType());
- if (preFilter != null) {
- rawRowRanges =
- Range.and(
- Range.sortAndMergeOverlap(rawRowRanges, true),
- Range.sortAndMergeOverlap(preFilter.toRangeList(),
true));
- }
+ return readRawSearch(rawRowRanges, preFilter, null, metric,
queryVector, true);
+ }
+
+ protected ScoredGlobalIndexResult readRawRefineSearch(
+ RoaringNavigableMap64 candidates,
+ @Nullable GlobalIndexer globalIndexer,
+ float[] queryVector) {
+ return readRawCandidateSearch(
+ candidates.toRangeList(),
+ candidates,
+ rawSearchMetric(globalIndexer),
+ queryVector,
+ false);
+ }
+
+ @Deprecated
+ protected ScoredGlobalIndexResult readRawRefineSearch(
+ List<Range> rawRowRanges,
+ @Nullable RoaringNavigableMap64 candidates,
+ @Nullable GlobalIndexer globalIndexer,
+ float[] queryVector) {
+ return readRawSearch(
+ rawRowRanges, null, candidates,
rawSearchMetric(globalIndexer), queryVector, false);
+ }
+
+ protected ScoredGlobalIndexResult readRawCandidateSearch(
+ List<Range> rawRowRanges,
+ RoaringNavigableMap64 candidates,
+ String metric,
+ float[] queryVector,
+ boolean includeFilter) {
+ return readRawSearch(rawRowRanges, null, candidates, metric,
queryVector, includeFilter);
+ }
+
+ private ScoredGlobalIndexResult readRawSearch(
+ List<Range> rawRowRanges,
+ @Nullable RoaringNavigableMap64 preFilter,
+ @Nullable RoaringNavigableMap64 scoreCandidates,
+ String metric,
+ float[] queryVector,
+ boolean includeFilter) {
+ RowType readType = rawSearchReadType(includeFilter);
+ rawRowRanges = filteredRawRowRanges(rawRowRanges, preFilter);
if (rawRowRanges.isEmpty()) {
return ScoredGlobalIndexResult.createEmpty();
}
TableScan.Plan plan =
newRawReadBuilder(readType,
false).withRowRanges(rawRowRanges).newScan().plan();
- ReadBuilder readBuilder = newRawReadBuilder(readType, true);
- RoaringNavigableMap64 resultBitmap = new RoaringNavigableMap64();
- Map<Long, Float> scoreMap = new HashMap<>();
+ ReadBuilder readBuilder = newRawReadBuilder(readType, includeFilter);
int vectorIndex = readType.getFieldIndex(vectorColumn.name());
int rowIdIndex = readType.getFieldIndex(SpecialFields.ROW_ID.name());
+ PriorityQueue<long[]> topKHeap =
+ new PriorityQueue<>(Math.max(1, limit + 1),
WEAKEST_SCORE_FIRST);
try (RecordReader<InternalRow> reader =
readBuilder.newRead().executeFilter().createReader(plan)) {
reader.forEachRemaining(
row -> {
+ long rowId = row.getLong(rowIdIndex);
+ if (scoreCandidates != null &&
!scoreCandidates.contains(rowId)) {
+ return;
+ }
if (row.isNullAt(vectorIndex)) {
return;
}
float[] stored = getVector(row, vectorIndex);
- if (stored.length != queryVector.length) {
- throw new IllegalArgumentException(
- String.format(
- "Query vector dimension mismatch:
expected %d, got %d",
- stored.length,
queryVector.length));
- }
+ checkVectorDimension(queryVector, stored);
+ offerScore(
+ topKHeap, limit, rowId,
computeScore(queryVector, stored, metric));
+ });
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to read raw vectors for vector
search.", e);
+ }
+
+ return scoredResult(topKHeap);
+ }
+
+ protected Map<Long, float[]> readRawVectors(
+ RoaringNavigableMap64 candidates, boolean includeFilter) {
+ return readRawVectors(candidates.toRangeList(), candidates,
includeFilter);
+ }
+
+ protected Map<Long, float[]> readRawVectors(
+ List<Range> rawRowRanges,
+ @Nullable RoaringNavigableMap64 candidates,
+ boolean includeFilter) {
+ RowType readType = rawSearchReadType(includeFilter);
+ rawRowRanges = filteredRawRowRanges(rawRowRanges, null);
+ if (rawRowRanges.isEmpty()) {
+ return Collections.emptyMap();
+ }
+
+ TableScan.Plan plan =
+ newRawReadBuilder(readType,
false).withRowRanges(rawRowRanges).newScan().plan();
+ ReadBuilder readBuilder = newRawReadBuilder(readType, includeFilter);
+ Map<Long, float[]> rawVectors = new HashMap<>();
+ int vectorIndex = readType.getFieldIndex(vectorColumn.name());
+ int rowIdIndex = readType.getFieldIndex(SpecialFields.ROW_ID.name());
+
+ try (RecordReader<InternalRow> reader =
+ readBuilder.newRead().executeFilter().createReader(plan)) {
+ reader.forEachRemaining(
+ row -> {
long rowId = row.getLong(rowIdIndex);
- resultBitmap.add(rowId);
- scoreMap.put(rowId, computeScore(queryVector, stored,
metric));
+ if (candidates != null && !candidates.contains(rowId))
{
+ return;
+ }
+ if (row.isNullAt(vectorIndex)) {
+ return;
+ }
+ float[] stored = getVector(row, vectorIndex);
+ rawVectors.put(rowId, stored);
});
} catch (IOException e) {
throw new RuntimeException("Failed to read raw vectors for vector
search.", e);
}
- return ScoredGlobalIndexResult.create(resultBitmap,
scoreMap::get).topK(limit);
+ return rawVectors;
+ }
+
+ private List<Range> filteredRawRowRanges(
+ List<Range> rawRowRanges, @Nullable RoaringNavigableMap64
preFilter) {
+ if (preFilter == null) {
+ return rawRowRanges;
+ }
+ return Range.and(
+ Range.sortAndMergeOverlap(rawRowRanges, true),
+ Range.sortAndMergeOverlap(preFilter.toRangeList(), true));
+ }
+
+ protected ScoredGlobalIndexResult scoreRawVectors(
+ RoaringNavigableMap64 candidates,
+ Map<Long, float[]> rawVectors,
+ float[] queryVector,
+ String metric,
+ int topK) {
+ PriorityQueue<long[]> topKHeap =
+ new PriorityQueue<>(Math.max(1, topK + 1),
WEAKEST_SCORE_FIRST);
+ for (long rowId : candidates) {
+ float[] stored = rawVectors.get(rowId);
+ if (stored == null) {
+ continue;
+ }
+ checkVectorDimension(queryVector, stored);
+ offerScore(topKHeap, topK, rowId, computeScore(queryVector,
stored, metric));
+ }
+ return scoredResult(topKHeap);
+ }
+
+ private static void offerScore(
+ PriorityQueue<long[]> topKHeap, int topK, long rowId, float score)
{
+ if (topK <= 0) {
+ return;
+ }
+ long[] entry = new long[] {rowId, Float.floatToRawIntBits(score)};
+ if (topKHeap.size() < topK) {
+ topKHeap.offer(entry);
+ } else if (WEAKEST_SCORE_FIRST.compare(entry, topKHeap.peek()) > 0) {
+ topKHeap.poll();
+ topKHeap.offer(entry);
+ }
+ }
+
+ private static ScoredGlobalIndexResult scoredResult(PriorityQueue<long[]>
topKHeap) {
+ if (topKHeap.isEmpty()) {
+ return ScoredGlobalIndexResult.createEmpty();
+ }
+ RoaringNavigableMap64 resultBitmap = new RoaringNavigableMap64();
+ Map<Long, Float> scoreMap = new HashMap<>();
+ for (long[] entry : topKHeap) {
+ long rowId = entry[0];
+ resultBitmap.add(rowId);
+ scoreMap.put(rowId, Float.intBitsToFloat((int) entry[1]));
+ }
+ return ScoredGlobalIndexResult.create(resultBitmap, scoreMap::get);
+ }
+
+ private static void checkVectorDimension(float[] queryVector, float[]
stored) {
+ if (stored.length != queryVector.length) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Query vector dimension mismatch: expected %d, got
%d",
+ stored.length, queryVector.length));
+ }
+ }
+
+ protected RowType rawSearchReadType(boolean includeFilter) {
+ RowType tableRowType = table.rowType();
+ List<String> readFields = new ArrayList<>();
+ readFields.add(vectorColumn.name());
+
+ if (includeFilter && filter != null) {
+ Set<String> filterFields =
PredicateVisitor.collectFieldNames(filter);
+ for (String field : tableRowType.getFieldNames()) {
+ if (filterFields.contains(field) &&
!readFields.contains(field)) {
+ readFields.add(field);
+ }
+ }
+ }
+
+ return
SpecialFields.rowTypeWithRowId(tableRowType.project(readFields));
}
private ReadBuilder newRawReadBuilder(RowType readType, boolean
includeFilter) {
@@ -554,7 +716,7 @@ public abstract class AbstractVectorRead implements
Serializable {
return metric.toLowerCase().replace('-', '_');
}
- private int configuredRefineFactor(String indexType) {
+ protected int configuredRefineFactor(String indexType) {
String value = configuredRefineFactor(options, indexType);
if (value == null) {
value = configuredRefineFactor(table.options(), indexType);
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorReadImpl.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorReadImpl.java
index 739b16260d..72003cff82 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorReadImpl.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorReadImpl.java
@@ -138,9 +138,35 @@ public class BatchVectorReadImpl extends
AbstractVectorRead implements BatchVect
}
}
}
+ return maybeRerankIndexedBatchResults(merged, indexType,
globalIndexer);
+ }
+
+ protected ScoredGlobalIndexResult[] maybeRerankIndexedBatchResults(
+ ScoredGlobalIndexResult[] results, String indexType, GlobalIndexer
globalIndexer) {
+ if (configuredRefineFactor(indexType) == 0) {
+ return results;
+ }
+
+ int n = results.length;
+ int searchLimit = indexedSearchLimit(indexType);
+ ScoredGlobalIndexResult[] candidates = new ScoredGlobalIndexResult[n];
+ RoaringNavigableMap64 unionCandidates = new RoaringNavigableMap64();
+ for (int i = 0; i < n; i++) {
+ candidates[i] = results[i].topK(searchLimit);
+ unionCandidates.or(candidates[i].results());
+ }
+
+ if (unionCandidates.isEmpty()) {
+ return candidates;
+ }
+
+ Map<Long, float[]> rawVectors = readRawVectors(unionCandidates, false);
+ String metric = rawSearchMetric(globalIndexer);
+ ScoredGlobalIndexResult[] reranked = new ScoredGlobalIndexResult[n];
for (int i = 0; i < n; i++) {
- merged[i] = maybeRerankIndexedResult(merged[i], indexType,
globalIndexer, vectors[i]);
+ reranked[i] =
+ scoreRawVectors(candidates[i].results(), rawVectors,
vectors[i], metric, limit);
}
- return merged;
+ return reranked;
}
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
index fdac8eef8a..fefe3bce60 100644
---
a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
@@ -48,6 +48,7 @@ import org.apache.paimon.predicate.Transform;
import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.schema.Schema;
import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.SpecialFields;
import org.apache.paimon.table.TableTestBase;
import org.apache.paimon.table.sink.BatchTableCommit;
import org.apache.paimon.table.sink.BatchTableWrite;
@@ -59,6 +60,7 @@ import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.Range;
+import org.apache.paimon.utils.RoaringNavigableMap64;
import org.junit.jupiter.api.Test;
@@ -69,7 +71,9 @@ import java.io.ObjectOutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import static
org.apache.paimon.table.source.DeletionVectorTestUtils.commitDeletionVectors;
import static org.assertj.core.api.Assertions.assertThat;
@@ -584,6 +588,56 @@ public class VectorSearchBuilderTest extends TableTestBase
{
assertThat(batchRefined.get(1).results()).containsExactly(2L);
}
+ @Test
+ public void testRawRefineReadTypeContainsOnlyVectorAndRowId() throws
Exception {
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+ Predicate idFilter = new
PredicateBuilder(table.rowType()).greaterOrEqual(0, 1);
+
+ ExposingVectorRead read = new ExposingVectorRead(table, idFilter);
+
+ assertThat(read.rawReadType(false).getFieldNames())
+ .containsExactly(VECTOR_FIELD_NAME,
SpecialFields.ROW_ID.name());
+ assertThat(read.rawReadType(true).getFieldNames())
+ .containsExactly(VECTOR_FIELD_NAME, "id",
SpecialFields.ROW_ID.name());
+ }
+
+ @Test
+ public void testRawCandidateSearchScoresOnlyCandidateBitmap() throws
Exception {
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+ writeVectors(table, new float[][] {{0.0f, 0.0f}, {10.0f, 0.0f},
{20.0f, 0.0f}});
+
+ RoaringNavigableMap64 candidates = new RoaringNavigableMap64();
+ candidates.add(1L);
+
+ ExposingVectorRead read = new ExposingVectorRead(table, null);
+ ScoredGlobalIndexResult result =
+ read.rawCandidateSearch(
+ Collections.singletonList(new Range(0, 2)),
+ candidates,
+ new float[] {0.0f, 0.0f});
+
+ assertThat(result.results()).containsExactly(1L);
+ }
+
+ @Test
+ public void testBatchRefineReadsUnionCandidatesOnceAndScoresPerQuery() {
+ RecordingBatchVectorRead read = new RecordingBatchVectorRead();
+
+ ScoredGlobalIndexResult[] reranked =
+ read.rerank(
+ new ScoredGlobalIndexResult[] {
+ scoredResult(1.0f, 0L, 2L), scoredResult(1.0f, 1L,
2L)
+ });
+
+ assertThat(read.rawReadCount).isEqualTo(1);
+ assertThat(read.rawReadRanges).containsExactly(new Range(0, 2));
+ assertThat(read.rawCandidates).containsExactly(0L, 1L, 2L);
+ assertThat(reranked[0].results()).containsExactly(2L);
+ assertThat(reranked[1].results()).containsExactly(2L);
+ }
+
@Test
public void testVectorSearchRefineFactorValidation() throws Exception {
createTableDefault();
@@ -1355,6 +1409,85 @@ public class VectorSearchBuilderTest extends
TableTestBase {
buildAndCommitIndex(table, VECTOR_FIELD_NAME, vectors);
}
+ private static class ExposingVectorRead extends VectorReadImpl {
+
+ private ExposingVectorRead(FileStoreTable table, Predicate filter) {
+ super(
+ table,
+ null,
+ filter,
+ 1,
+ table.rowType().getField(VECTOR_FIELD_NAME),
+ new float[] {0.0f, 0.0f},
+ null);
+ }
+
+ private RowType rawReadType(boolean includeFilter) {
+ return rawSearchReadType(includeFilter);
+ }
+
+ private ScoredGlobalIndexResult rawCandidateSearch(
+ List<Range> rawRowRanges, RoaringNavigableMap64 candidates,
float[] queryVector) {
+ return readRawCandidateSearch(rawRowRanges, candidates, "l2",
queryVector, false);
+ }
+ }
+
+ private static class RecordingBatchVectorRead extends BatchVectorReadImpl {
+
+ private int rawReadCount;
+ private List<Range> rawReadRanges;
+ private RoaringNavigableMap64 rawCandidates;
+ private final Map<Long, float[]> rawVectors = new HashMap<>();
+
+ private RecordingBatchVectorRead() {
+ super(
+ null,
+ null,
+ null,
+ 1,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[][] {{0.0f}, {10.0f}},
+ refineOptions());
+ rawVectors.put(0L, new float[] {10.0f});
+ rawVectors.put(1L, new float[] {0.0f});
+ rawVectors.put(2L, new float[] {5.0f});
+ }
+
+ private ScoredGlobalIndexResult[] rerank(ScoredGlobalIndexResult[]
results) {
+ return maybeRerankIndexedBatchResults(results, "ivf-pq", null);
+ }
+
+ @Override
+ protected Map<Long, float[]> readRawVectors(
+ List<Range> rawRowRanges, RoaringNavigableMap64 candidates,
boolean includeFilter) {
+ rawReadCount++;
+ rawReadRanges = rawRowRanges;
+ rawCandidates = candidates;
+ assertThat(includeFilter).isFalse();
+
+ Map<Long, float[]> result = new HashMap<>();
+ for (long rowId : candidates) {
+ result.put(rowId, rawVectors.get(rowId));
+ }
+ return result;
+ }
+ }
+
+ private static Map<String, String> refineOptions() {
+ Map<String, String> options = new HashMap<>();
+ options.put("refine_factor", "2");
+ options.put("test.vector.metric", "l2");
+ return options;
+ }
+
+ private static ScoredGlobalIndexResult scoredResult(float score, long...
rowIds) {
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ for (long rowId : rowIds) {
+ rows.add(rowId);
+ }
+ return ScoredGlobalIndexResult.create(rows, rowId -> score);
+ }
+
private void buildAndCommitIndex(FileStoreTable table, String fieldName,
float[][] vectors)
throws Exception {
Options options = table.coreOptions().toConfiguration();
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index bcb142e4f1..b7bde67ec9 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -26,6 +26,7 @@ from pypaimon.globalindex.global_index_result import
GlobalIndexResult
from pypaimon.globalindex.offset_global_index_reader import
OffsetGlobalIndexReader
from pypaimon.globalindex.vector_search import VectorSearch
from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
+from pypaimon.table.special_fields import SpecialFields
from pypaimon.table.source.vector_search_split import (
IndexVectorSearchSplit,
RawVectorSearchSplit,
@@ -243,48 +244,120 @@ class AbstractVectorSearchReadImpl:
future.add_done_callback(lambda _: reader.close())
return future
- def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector,
index_type=None):
- raw_row_ranges = Range.sort_and_merge_overlap(raw_row_ranges, True)
- if pre_filter is not None:
- raw_row_ranges = Range.and_(
- raw_row_ranges,
- Range.sort_and_merge_overlap(pre_filter.to_range_list(), True),
- )
+ def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector,
+ index_type=None, include_filter=True,
+ score_candidates=None):
+ raw_row_ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter)
if not raw_row_ranges:
return DictBasedScoredIndexResult({})
- read_builder = self._table.new_read_builder()
- if self._partition_filter is not None:
- read_builder = read_builder.with_partition_filter(
- self._partition_filter)
- if self._filter is not None:
- read_builder = read_builder.with_filter(self._filter)
- from pypaimon.table.special_fields import SpecialFields
- projection = [f.name for f in self._table.fields]
- if SpecialFields.ROW_ID.name not in projection:
- projection.append(SpecialFields.ROW_ID.name)
- read_builder = read_builder.with_projection(projection)
- plan = read_builder.new_scan().with_global_index_result(
- GlobalIndexResult.from_ranges(raw_row_ranges)).plan()
- table = read_builder.new_read().to_arrow(plan.splits())
+ table = self._read_raw_arrow(raw_row_ranges, include_filter)
if table is None or table.num_rows == 0:
return DictBasedScoredIndexResult({})
- row_ids = table.column(SpecialFields.ROW_ID.name).to_pylist()
- vectors = table.column(self._vector_column.name).to_pylist()
+ top_k_heap = []
metric = _raw_search_metric(
self._table, self._vector_column, self._options, index_type)
- scores = {}
+ row_ids = table.column(SpecialFields.ROW_ID.name).to_pylist()
+ vectors = table.column(self._vector_column.name).to_pylist()
for row_id, stored in zip(row_ids, vectors):
+ if score_candidates is not None and row_id not in score_candidates:
+ continue
if stored is None:
continue
stored_vector = _to_vector_list(stored)
- if len(stored_vector) != len(query_vector):
- raise ValueError(
- "Query vector dimension mismatch: expected %d, got %d"
- % (len(stored_vector), len(query_vector)))
- scores[row_id] = _compute_score(query_vector, stored_vector,
metric)
- return DictBasedScoredIndexResult(scores).top_k(self._limit)
+ _check_vector_dimension(query_vector, stored_vector)
+ _offer_score(
+ top_k_heap,
+ self._limit,
+ row_id,
+ _compute_score(query_vector, stored_vector, metric),
+ )
+ return _scored_result(top_k_heap)
+
+ def _read_raw_vectors(self, candidates, include_filter=True):
+ return self._read_raw_candidate_vectors(
+ candidates.to_range_list(), candidates, include_filter)
+
+ def _read_raw_candidate_vectors(self, raw_row_ranges, candidates,
+ include_filter=True):
+ raw_row_ranges = _filtered_raw_row_ranges(raw_row_ranges, None)
+ if not raw_row_ranges:
+ return {}
+
+ table = self._read_raw_arrow(raw_row_ranges, include_filter)
+ if table is None or table.num_rows == 0:
+ return {}
+
+ row_ids = table.column(SpecialFields.ROW_ID.name).to_pylist()
+ vectors = table.column(self._vector_column.name).to_pylist()
+ raw_vectors = {}
+ for row_id, stored in zip(row_ids, vectors):
+ if candidates is not None and row_id not in candidates:
+ continue
+ if stored is None:
+ continue
+ raw_vectors[row_id] = _to_vector_list(stored)
+ return raw_vectors
+
+ def _read_raw_arrow(self, raw_row_ranges, include_filter):
+ read_builder = self._table.new_read_builder()
+ if self._partition_filter is not None:
+ read_builder = read_builder.with_partition_filter(
+ self._partition_filter)
+ if include_filter and self._filter is not None:
+ read_builder = read_builder.with_filter(self._filter)
+ read_builder = read_builder.with_projection(
+ self._raw_search_projection(include_filter))
+ plan = read_builder.new_scan().with_global_index_result(
+ GlobalIndexResult.from_ranges(raw_row_ranges)).plan()
+ return read_builder.new_read().to_arrow(plan.splits())
+
+ def _score_raw_vectors(self, candidates, raw_vectors, query_vector,
metric, top_k):
+ top_k_heap = []
+ for row_id in candidates:
+ stored_vector = raw_vectors.get(row_id)
+ if stored_vector is None:
+ continue
+ _check_vector_dimension(query_vector, stored_vector)
+ _offer_score(
+ top_k_heap,
+ top_k,
+ row_id,
+ _compute_score(query_vector, stored_vector, metric),
+ )
+ return _scored_result(top_k_heap)
+
+ def _read_raw_refine_search(self, candidates, query_vector,
index_type=None):
+ return self._read_raw_candidate_search(
+ candidates.to_range_list(),
+ candidates,
+ query_vector,
+ index_type,
+ include_filter=False,
+ )
+
+ def _read_raw_candidate_search(self, raw_row_ranges, candidates,
query_vector,
+ index_type=None, include_filter=False):
+ return self._read_raw_search(
+ raw_row_ranges,
+ None,
+ query_vector,
+ index_type,
+ include_filter=include_filter,
+ score_candidates=candidates,
+ )
+
+ def _raw_search_projection(self, include_filter):
+ projection = [self._vector_column.name]
+ if include_filter and self._filter is not None:
+ filter_fields = _predicate_field_names(self._filter)
+ for field in self._table.fields:
+ if field.name in filter_fields and field.name not in
projection:
+ projection.append(field.name)
+ if SpecialFields.ROW_ID.name not in projection:
+ projection.append(SpecialFields.ROW_ID.name)
+ return projection
def _eval_batch(self, row_range_start, row_range_end, vector_index_files,
query_vectors, search_limit, include_row_ids):
@@ -319,13 +392,39 @@ class AbstractVectorSearchReadImpl:
result.results().is_empty()):
return result
candidates = result.top_k(self._indexed_search_limit(index_type))
- return self._read_raw_search(
- candidates.results().to_range_list(),
+ return self._read_raw_refine_search(
candidates.results(),
query_vector,
index_type,
)
+ def _maybe_rerank_indexed_results(self, results, index_type,
query_vectors):
+ if self._configured_refine_factor(index_type) == 0:
+ return results
+
+ search_limit = self._indexed_search_limit(index_type)
+ candidates = [result.top_k(search_limit) for result in results]
+ union_candidates = RoaringBitmap64()
+ for result in candidates:
+ union_candidates = RoaringBitmap64.or_(
+ union_candidates, result.results())
+ if union_candidates.is_empty():
+ return candidates
+
+ raw_vectors = self._read_raw_vectors(union_candidates,
include_filter=False)
+ metric = _raw_search_metric(
+ self._table, self._vector_column, self._options, index_type)
+ return [
+ self._score_raw_vectors(
+ candidates[i].results(),
+ raw_vectors,
+ query_vectors[i],
+ metric,
+ self._limit,
+ )
+ for i in range(len(candidates))
+ ]
+
def _configured_refine_factor(self, index_type):
value = _configured_refine_factor(
self._options, self._vector_column.name, index_type)
@@ -454,18 +553,22 @@ class
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
if row_id not in merged_scores[i]:
merged_scores[i][row_id] = score_getter(row_id)
+ indexed_results = [
+ DictBasedScoredIndexResult(merged_scores[i]).top_k(search_limit)
+ for i in range(n)
+ ]
+ indexed_results = self._maybe_rerank_indexed_results(
+ indexed_results, index_type, self._query_vectors)
+
# Each query: merge indexed results with the raw (brute-force)
fallback.
raw_pre_filter = self._raw_pre_filter(raw_splits)
raw_ranges = _raw_row_ranges(raw_splits)
raw_index_type = _raw_search_index_type(raw_splits)
results = []
for i in range(n):
- indexed =
DictBasedScoredIndexResult(merged_scores[i]).top_k(search_limit)
- indexed = self._maybe_rerank_indexed_result(
- indexed, index_type, self._query_vectors[i])
raw = self._read_raw_search(
raw_ranges, raw_pre_filter, self._query_vectors[i],
raw_index_type)
- results.append(indexed.or_(raw).top_k(self._limit))
+ results.append(indexed_results[i].or_(raw).top_k(self._limit))
return results
@@ -508,6 +611,16 @@ def _raw_row_ranges(raw_splits):
return Range.sort_and_merge_overlap(ranges, True)
+def _filtered_raw_row_ranges(raw_row_ranges, pre_filter):
+ raw_row_ranges = Range.sort_and_merge_overlap(raw_row_ranges, True)
+ if pre_filter is None:
+ return raw_row_ranges
+ return Range.and_(
+ raw_row_ranges,
+ Range.sort_and_merge_overlap(pre_filter.to_range_list(), True),
+ )
+
+
def _raw_search_index_type(raw_splits):
for split in raw_splits:
if split.index_type is not None:
@@ -515,6 +628,17 @@ def _raw_search_index_type(raw_splits):
return None
+def _predicate_field_names(predicate):
+ if predicate is None:
+ return set()
+ if predicate.method in ("and", "or"):
+ names = set()
+ for child in predicate.literals or []:
+ names.update(_predicate_field_names(child))
+ return names
+ return {predicate.field} if predicate.field is not None else set()
+
+
def _vector_index_type(index_splits):
for split in index_splits:
if split.vector_index_files:
@@ -547,6 +671,29 @@ def _to_vector_list(value):
return list(value)
+def _offer_score(heap, top_k, row_id, score):
+ if top_k <= 0:
+ return
+ import heapq
+
+ entry = (score, -row_id, row_id)
+ if len(heap) < top_k:
+ heapq.heappush(heap, entry)
+ elif entry[:2] > heap[0][:2]:
+ heapq.heapreplace(heap, entry)
+
+
+def _scored_result(heap):
+ return DictBasedScoredIndexResult({row_id: score for score, _, row_id in
heap})
+
+
+def _check_vector_dimension(query_vector, stored_vector):
+ if len(stored_vector) != len(query_vector):
+ raise ValueError(
+ "Query vector dimension mismatch: expected %d, got %d"
+ % (len(stored_vector), len(query_vector)))
+
+
def _configured_refine_factor(options, vector_column_name, index_type):
prefixes = []
field_prefix = "fields.%s." % vector_column_name
diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py
b/paimon-python/pypaimon/tests/vector_search_filter_test.py
index 3648317f1b..6eb700cf4c 100644
--- a/paimon-python/pypaimon/tests/vector_search_filter_test.py
+++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py
@@ -178,6 +178,7 @@ def _install_raw_vector_read_builder(table,
vector_column_name, row_id_to_vector
def with_global_index_result(self, result):
ranges = result.results().to_range_list()
+ calls["raw_read_count"] = calls.get("raw_read_count", 0) + 1
calls["global_index_ranges"] = ranges
self._row_ids = [
row_id
@@ -1141,6 +1142,7 @@ class VectorSearchFilterTest(unittest.TestCase):
self.assertEqual([3], captured_limits)
self.assertEqual([Range(0, 2)], raw_calls["global_index_ranges"])
self.assertEqual([0, 1, 2], raw_calls["candidate_ids"])
+ self.assertEqual(["embedding", "_ROW_ID"], raw_calls["projection"])
self.assertEqual([0], sorted(list(result.results())))
def test_refine_factor_one_reranks_without_expanding_candidates(self):
@@ -1186,9 +1188,33 @@ class VectorSearchFilterTest(unittest.TestCase):
self.assertEqual([1], captured_limits)
self.assertEqual([Range(2, 2)], raw_calls["global_index_ranges"])
self.assertEqual([2], raw_calls["candidate_ids"])
+ self.assertEqual(["embedding", "_ROW_ID"], raw_calls["projection"])
self.assertEqual([2], sorted(list(result.results())))
self.assertLess(result.score_getter()(2), 1.0)
+ def test_raw_candidate_search_scores_only_candidate_bitmap(self):
+ from pypaimon.table.source.vector_search_read import
VectorSearchReadImpl
+
+ table = _StubTable(fields=[self.id_field, self.embedding_field],
+ entries=[])
+ raw_calls = _install_raw_vector_read_builder(
+ table, "embedding", {0: [0.0], 1: [10.0], 2: [20.0]})
+
+ reader = VectorSearchReadImpl(
+ table,
+ limit=1,
+ vector_column=self.embedding_field,
+ query_vector=[0.0],
+ filter_=None,
+ )
+
+ result = reader._read_raw_candidate_search(
+ [Range(0, 2)], _bitmap(1), [0.0], "ivf-pq")
+
+ self.assertEqual([Range(0, 2)], raw_calls["global_index_ranges"])
+ self.assertEqual([0, 1, 2], raw_calls["candidate_ids"])
+ self.assertEqual([1], sorted(list(result.results())))
+
def test_refine_factor_query_options_override_table_options(self):
from pypaimon.common.options.options import Options
from pypaimon.globalindex.vector_search_result import (
@@ -2432,7 +2458,7 @@ class VectorSearchManySplitsTest(unittest.TestCase):
self.assertIs(partition_filter, calls["partition_filter"])
self.assertIs(filter_pred, calls["filter"])
self.assertEqual([Range(5, 6)], calls["global_index_ranges"])
- self.assertIn("_ROW_ID", calls["projection"])
+ self.assertEqual(["embedding", "id", "_ROW_ID"], calls["projection"])
self.assertEqual(["split"], calls["splits"])
self.assertEqual([5], sorted(list(result.results())))
@@ -2798,8 +2824,8 @@ class BatchVectorSearchTest(unittest.TestCase):
row_range_start=0, row_range_end=2)
table = _StubTable(fields=[embedding_field], entries=[entry])
_patch_snapshot(self, [entry])
- _install_raw_vector_read_builder(
- table, "embedding", {0: [0.0], 1: [10.0], 2: [20.0]})
+ raw_calls = _install_raw_vector_read_builder(
+ table, "embedding", {0: [10.0], 1: [0.0], 2: [5.0]})
captured_limits = []
def _fake_create(index_type, file_io, index_path,
@@ -2807,11 +2833,9 @@ class BatchVectorSearchTest(unittest.TestCase):
class _FakeReader(GlobalIndexReader):
def visit_batch_vector_search(self_inner, bvs):
captured_limits.append(bvs.limit)
- approximate_scores = [(2, 100.0), (1, 50.0), (0, 1.0)]
return _completed_future([
- DictBasedScoredIndexResult(
- dict(approximate_scores[:bvs.limit]))
- for _ in range(bvs.vector_count)
+ DictBasedScoredIndexResult({0: 100.0, 2: 50.0}),
+ DictBasedScoredIndexResult({1: 100.0, 2: 50.0}),
])
def close(self_inner):
@@ -2827,12 +2851,16 @@ class BatchVectorSearchTest(unittest.TestCase):
.with_vector_column("embedding")
.with_query_vectors([[0.0], [20.0]])
.with_limit(1)
- .with_option("ivf.refine_factor", "3")
+ .with_option("ivf.refine_factor", "2")
.execute_batch_local()
)
- self.assertEqual([3], captured_limits)
- self.assertEqual([0], sorted(list(results[0].results())))
+ self.assertEqual([2], captured_limits)
+ self.assertEqual(1, raw_calls["raw_read_count"])
+ self.assertEqual([Range(0, 2)], raw_calls["global_index_ranges"])
+ self.assertEqual([0, 1, 2], raw_calls["candidate_ids"])
+ self.assertEqual(["embedding", "_ROW_ID"], raw_calls["projection"])
+ self.assertEqual([2], sorted(list(results[0].results())))
self.assertEqual([2], sorted(list(results[1].results())))
def test_batch_empty_splits_returns_empty_per_query(self):
diff --git
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkVectorReadImplTest.java
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkVectorReadImplTest.java
index 16b998ced6..a88d2d9cbf 100644
---
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkVectorReadImplTest.java
+++
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkVectorReadImplTest.java
@@ -195,22 +195,20 @@ public class SparkVectorReadImplTest {
}
@Override
- protected ScoredGlobalIndexResult readRawSearch(
- List<Range> rawRowRanges,
- @Nullable RoaringNavigableMap64 preFilter,
+ protected ScoredGlobalIndexResult readRawRefineSearch(
+ RoaringNavigableMap64 candidates,
@Nullable GlobalIndexer globalIndexer,
float[] queryVector) {
assertThat(globalIndexer).isInstanceOf(VectorGlobalIndexer.class);
assertThat(((VectorGlobalIndexer)
globalIndexer).metric()).isEqualTo("l2");
assertThat(queryVector).containsExactly(0.0f);
- assertThat(preFilter).isNotNull();
rawSearchCandidateRows = new ArrayList<>();
- for (long rowId : preFilter) {
+ for (long rowId : candidates) {
rawSearchCandidateRows.add(rowId);
}
RoaringNavigableMap64 rows = new RoaringNavigableMap64();
- for (long rowId : preFilter) {
+ for (long rowId : candidates) {
rows.add(rowId);
}
return ScoredGlobalIndexResult.create(
