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 b1daaa197a [vector] Add refine factor rerank for vector search (#8352)
b1daaa197a is described below
commit b1daaa197aed88b532cb5f487eb997ce158e99f3
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Jun 25 20:44:52 2026 +0800
[vector] Add refine factor rerank for vector search (#8352)
Add a LanceDB-style refine factor for vector search so approximate IVF
candidates can be reranked with the original vectors stored in the
Paimon table. This is especially useful for compressed vector indexes
such as IVF-PQ, where index scores may differ from exact raw-vector
scores.
---
docs/docs/multimodal-table/global-index/vector.mdx | 5 +
.../testvector/TestVectorGlobalIndexReader.java | 16 +-
.../testvector/TestVectorGlobalIndexer.java | 12 +-
.../paimon/table/source/AbstractVectorRead.java | 110 ++++++++-
.../paimon/table/source/BatchVectorReadImpl.java | 6 +
.../apache/paimon/table/source/VectorReadImpl.java | 5 +-
.../table/source/VectorSearchBuilderTest.java | 68 +++++
.../pypaimon/table/source/vector_search_read.py | 106 +++++++-
.../pypaimon/tests/vector_search_filter_test.py | 274 +++++++++++++++++++++
.../paimon/spark/read/SparkVectorReadImpl.java | 20 +-
.../paimon/spark/read/SparkVectorReadImplTest.java | 133 ++++++++++
11 files changed, 735 insertions(+), 20 deletions(-)
diff --git a/docs/docs/multimodal-table/global-index/vector.mdx
b/docs/docs/multimodal-table/global-index/vector.mdx
index 16d2642a4a..1fa560645e 100644
--- a/docs/docs/multimodal-table/global-index/vector.mdx
+++ b/docs/docs/multimodal-table/global-index/vector.mdx
@@ -156,6 +156,7 @@ Search-time options are passed with each vector search
request:
| Option | Default | Description |
|---|---|---|
| `ivf.nprobe` | `16` | Number of IVF clusters to probe during search. Higher
values usually improve recall but increase latency. |
+| `ivf.refine_factor` | Disabled | Retrieves `top_k * refine_factor` IVF
candidates and reranks them with the original vectors stored in the Paimon
table. It is disabled by default for every IVF variant and is most useful for
compressed indexes such as `ivf-pq` and `ivf-hnsw-sq` when recall is more
important than latency. |
| `hnsw.ef_search` | `0` | HNSW search width during search. Higher values
usually improve recall but increase latency. `0` uses the native library
default. |
| `diskann.search.list_size` | `max(1.5x top_k, 16)` | Lumina DiskANN search
list size. Higher values usually improve recall but increase latency. |
| `diskann.search.beam_width` | `4` | Lumina DiskANN search beam width. |
@@ -166,6 +167,10 @@ so you can use a larger `ivf.nprobe` or `hnsw.ef_search`
for higher recall queri
value for latency-sensitive queries. Lumina query-time options use the native
keys shown above; when
the same options are configured as table or index options, use the `lumina.`
prefix.
+`ivf.refine_factor` can also be configured with `refine_factor`,
`rerank_factor`, and hyphenated
+spellings such as `ivf.refine-factor`. Setting `ivf.refine_factor=1` still
performs the raw-vector
+rerank for the indexed candidates; leaving it unset skips the rerank stage.
+
<Tabs groupId="vector-search">
<TabItem value="spark-sql" label="Spark SQL">
diff --git
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
index 1e8cc432f2..f073052dd6 100644
---
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
+++
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
@@ -56,6 +56,7 @@ public class TestVectorGlobalIndexReader implements
GlobalIndexReader {
private final GlobalIndexFileReader fileReader;
private final GlobalIndexIOMeta ioMeta;
private final String metric;
+ private final boolean reverseScore;
private final String requiredOptionKey;
private final String requiredOptionValue;
@@ -66,18 +67,20 @@ public class TestVectorGlobalIndexReader implements
GlobalIndexReader {
public TestVectorGlobalIndexReader(
GlobalIndexFileReader fileReader, GlobalIndexIOMeta ioMeta, String
metric) {
- this(fileReader, ioMeta, metric, null, null);
+ this(fileReader, ioMeta, metric, false, null, null);
}
public TestVectorGlobalIndexReader(
GlobalIndexFileReader fileReader,
GlobalIndexIOMeta ioMeta,
String metric,
+ boolean reverseScore,
String requiredOptionKey,
String requiredOptionValue) {
this.fileReader = fileReader;
this.ioMeta = ioMeta;
this.metric = metric;
+ this.reverseScore = reverseScore;
this.requiredOptionKey = requiredOptionKey;
this.requiredOptionValue = requiredOptionValue;
}
@@ -148,16 +151,21 @@ public class TestVectorGlobalIndexReader implements
GlobalIndexReader {
}
private float computeScore(float[] query, float[] stored) {
+ float score;
switch (metric) {
case "l2":
- return computeL2Score(query, stored);
+ score = computeL2Score(query, stored);
+ break;
case "cosine":
- return computeCosineScore(query, stored);
+ score = computeCosineScore(query, stored);
+ break;
case "inner_product":
- return computeInnerProductScore(query, stored);
+ score = computeInnerProductScore(query, stored);
+ break;
default:
throw new IllegalArgumentException("Unknown metric: " +
metric);
}
+ return reverseScore ? -score : score;
}
private static float computeL2Score(float[] a, float[] b) {
diff --git
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
index f652719834..0b6d8d9fd6 100644
---
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
+++
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
@@ -56,6 +56,9 @@ public class TestVectorGlobalIndexer implements
VectorGlobalIndexer {
/** Option key for distance metric. */
public static final String OPT_METRIC = "test.vector.metric";
+ /** Option key to reverse scores for testing refine/rerank behavior. */
+ public static final String OPT_REVERSE_SCORE = "test.vector.reverse-score";
+
public static final String OPT_REQUIRED_OPTION_KEY =
"test.vector.required-option.key";
public static final String OPT_REQUIRED_OPTION_VALUE =
"test.vector.required-option.value";
@@ -65,6 +68,7 @@ public class TestVectorGlobalIndexer implements
VectorGlobalIndexer {
private final DataType fieldType;
private final int dimension;
private final String metric;
+ private final boolean reverseScore;
private final String requiredOptionKey;
private final String requiredOptionValue;
@@ -76,6 +80,7 @@ public class TestVectorGlobalIndexer implements
VectorGlobalIndexer {
this.fieldType = fieldType;
this.dimension = options.getInteger(OPT_DIMENSION, 0);
this.metric = options.getString(OPT_METRIC, "l2");
+ this.reverseScore = options.getBoolean(OPT_REVERSE_SCORE, false);
this.requiredOptionKey = options.getString(OPT_REQUIRED_OPTION_KEY,
null);
this.requiredOptionValue =
options.getString(OPT_REQUIRED_OPTION_VALUE, null);
}
@@ -92,7 +97,12 @@ public class TestVectorGlobalIndexer implements
VectorGlobalIndexer {
ExecutorService executor) {
checkArgument(files.size() == 1, "Expected exactly one index file per
shard");
return new TestVectorGlobalIndexReader(
- fileReader, files.get(0), metric, requiredOptionKey,
requiredOptionValue);
+ fileReader,
+ files.get(0),
+ metric,
+ reverseScore,
+ requiredOptionKey,
+ requiredOptionValue);
}
public int dimension() {
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 bb783415d9..a1996f9203 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
@@ -213,6 +213,7 @@ public abstract class AbstractVectorRead implements
Serializable {
long rowRangeEnd,
List<IndexFileMeta> vectorIndexFiles,
float[] vector,
+ int searchLimit,
@Nullable RoaringNavigableMap64 includeRowIds,
ExecutorService executor) {
if (vectorIndexFiles.isEmpty()) {
@@ -227,7 +228,7 @@ public abstract class AbstractVectorRead implements
Serializable {
GlobalIndexReader reader =
globalIndexer.createReader(indexFileReader, indexIOMetaList,
executor);
VectorSearch vectorSearch =
- new VectorSearch(vector, limit, vectorColumn.name(), options)
+ new VectorSearch(vector, searchLimit, vectorColumn.name(),
options)
.withIncludeRowIds(includeRowIds);
return new OffsetGlobalIndexReader(reader, rowRangeStart, rowRangeEnd)
.visitVectorSearch(vectorSearch)
@@ -241,6 +242,7 @@ public abstract class AbstractVectorRead implements
Serializable {
long rowRangeEnd,
List<IndexFileMeta> vectorIndexFiles,
float[][] vectors,
+ int searchLimit,
@Nullable RoaringNavigableMap64 includeRowIds,
ExecutorService executor) {
if (vectorIndexFiles.isEmpty()) {
@@ -255,7 +257,7 @@ public abstract class AbstractVectorRead implements
Serializable {
GlobalIndexReader reader =
globalIndexer.createReader(indexFileReader, indexIOMetaList,
executor);
BatchVectorSearch batchVectorSearch =
- new BatchVectorSearch(vectors, limit, vectorColumn.name(),
options)
+ new BatchVectorSearch(vectors, searchLimit,
vectorColumn.name(), options)
.withIncludeRowIds(includeRowIds);
return new OffsetGlobalIndexReader(reader, rowRangeStart, rowRangeEnd)
.visitBatchVectorSearch(batchVectorSearch)
@@ -282,6 +284,40 @@ public abstract class AbstractVectorRead implements
Serializable {
return result.or(rawResult).topK(limit);
}
+ protected int indexedSearchLimit(String indexType) {
+ int refineFactor = configuredRefineFactor(indexType);
+ if (refineFactor == 0) {
+ return limit;
+ }
+ if (limit > Integer.MAX_VALUE / refineFactor) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Vector search limit overflow: limit=%d, refine
factor=%d",
+ limit, refineFactor));
+ }
+ return limit * refineFactor;
+ }
+
+ protected ScoredGlobalIndexResult maybeRerankIndexedResult(
+ ScoredGlobalIndexResult result,
+ String indexType,
+ @Nullable GlobalIndexer globalIndexer,
+ float[] queryVector) {
+ if (configuredRefineFactor(indexType) == 0 ||
result.results().isEmpty()) {
+ return result;
+ }
+ ScoredGlobalIndexResult candidates =
result.topK(indexedSearchLimit(indexType));
+ return readRawSearch(
+ candidates.results().toRangeList(),
+ candidates.results(),
+ globalIndexer,
+ queryVector);
+ }
+
+ protected String vectorIndexType(List<IndexVectorSearchSplit> splits) {
+ return firstVectorIndexFile(splits).indexType();
+ }
+
protected ScoredGlobalIndexResult[] emptyScoredResults(int n) {
ScoredGlobalIndexResult[] results = new ScoredGlobalIndexResult[n];
for (int i = 0; i < n; i++) {
@@ -509,6 +545,76 @@ public abstract class AbstractVectorRead implements
Serializable {
return metric.toLowerCase().replace('-', '_');
}
+ private int configuredRefineFactor(String indexType) {
+ String value = configuredRefineFactor(options, indexType);
+ if (value == null) {
+ value = configuredRefineFactor(table.options(), indexType);
+ }
+ if (value == null) {
+ return 0;
+ }
+ try {
+ int factor = Integer.parseInt(value);
+ if (factor <= 0) {
+ throw new IllegalArgumentException(
+ "Vector refine factor must be positive, got: " +
value);
+ }
+ return factor;
+ } catch (NumberFormatException e) {
+ throw new IllegalArgumentException(
+ "Invalid vector refine factor: " + value + ". Must be an
integer.", e);
+ }
+ }
+
+ @Nullable
+ private String configuredRefineFactor(Map<String, String> options, String
indexType) {
+ List<String> prefixes = new ArrayList<>();
+ String fieldPrefix = "fields." + vectorColumn.name() + ".";
+ addRefinePrefixes(prefixes, fieldPrefix, indexType);
+ addRefinePrefixes(prefixes, "", indexType);
+
+ for (String prefix : prefixes) {
+ String value = refineFactorOption(options, prefix +
"refine_factor");
+ if (value == null) {
+ value = refineFactorOption(options, prefix + "refine-factor");
+ }
+ if (value == null) {
+ value = refineFactorOption(options, prefix + "rerank_factor");
+ }
+ if (value == null) {
+ value = refineFactorOption(options, prefix + "rerank-factor");
+ }
+ if (value != null) {
+ return value;
+ }
+ }
+ return null;
+ }
+
+ private static void addRefinePrefixes(List<String> prefixes, String base,
String indexType) {
+ if (indexType != null && !indexType.isEmpty()) {
+ prefixes.add(base + indexType + ".");
+ String normalizedIndexType = normalizeIndexType(indexType);
+ if (!normalizedIndexType.equals(indexType)) {
+ prefixes.add(base + normalizedIndexType + ".");
+ }
+ if (normalizedIndexType.startsWith("ivf")) {
+ prefixes.add(base + "ivf.");
+ }
+ }
+ prefixes.add(base);
+ }
+
+ @Nullable
+ private static String refineFactorOption(Map<String, String> options,
String key) {
+ String value = options.get(key);
+ return value == null ? null : value.trim();
+ }
+
+ private static String normalizeIndexType(String indexType) {
+ return indexType.toLowerCase().replace('-', '_');
+ }
+
private static IndexFileMeta
firstVectorIndexFile(List<IndexVectorSearchSplit> splits) {
for (IndexVectorSearchSplit split : splits) {
if (!split.vectorIndexFiles().isEmpty()) {
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 2a6fb07f7c..739b16260d 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
@@ -98,6 +98,8 @@ public class BatchVectorReadImpl extends AbstractVectorRead
implements BatchVect
List<IndexVectorSearchSplit> splits, GlobalIndexer globalIndexer) {
int n = vectors.length;
List<RoaringNavigableMap64> preFilters = preFilters(splits);
+ String indexType = vectorIndexType(splits);
+ int searchLimit = indexedSearchLimit(indexType);
IndexPathFactory indexPathFactory =
table.store().pathFactory().globalIndexFileFactory();
@@ -116,6 +118,7 @@ public class BatchVectorReadImpl extends AbstractVectorRead
implements BatchVect
split.rowRangeEnd(),
split.vectorIndexFiles(),
vectors,
+ searchLimit,
preFilters.isEmpty() ? null : preFilters.get(i),
executor));
}
@@ -135,6 +138,9 @@ public class BatchVectorReadImpl extends AbstractVectorRead
implements BatchVect
}
}
}
+ for (int i = 0; i < n; i++) {
+ merged[i] = maybeRerankIndexedResult(merged[i], indexType,
globalIndexer, vectors[i]);
+ }
return merged;
}
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
index f04a03c064..9d8bd1541e 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
@@ -84,6 +84,8 @@ public class VectorReadImpl extends AbstractVectorRead
implements VectorRead {
protected ScoredGlobalIndexResult readIndexed(
List<IndexVectorSearchSplit> splits, GlobalIndexer globalIndexer) {
List<RoaringNavigableMap64> preFilters = preFilters(splits);
+ String indexType = vectorIndexType(splits);
+ int searchLimit = indexedSearchLimit(indexType);
IndexPathFactory indexPathFactory =
table.store().pathFactory().globalIndexFileFactory();
@@ -102,6 +104,7 @@ public class VectorReadImpl extends AbstractVectorRead
implements VectorRead {
split.rowRangeEnd(),
split.vectorIndexFiles(),
vector,
+ searchLimit,
preFilters.isEmpty() ? null : preFilters.get(i),
executor));
}
@@ -115,6 +118,6 @@ public class VectorReadImpl extends AbstractVectorRead
implements VectorRead {
merged = merged.or(splitResult.get());
}
}
- return merged;
+ return maybeRerankIndexedResult(merged, indexType, globalIndexer,
vector);
}
}
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 e776dd8314..b861ea29a4 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
@@ -72,6 +72,7 @@ import java.util.Collections;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link VectorSearchBuilder} using test-only brute-force vector
index. */
public class VectorSearchBuilderTest extends TableTestBase {
@@ -479,6 +480,73 @@ public class VectorSearchBuilderTest extends TableTestBase
{
assertThat(result.results().isEmpty()).isFalse();
}
+ @Test
+ public void testVectorSearchRefineFactorReranksIndexCandidates() throws
Exception {
+ catalog.createTable(
+ identifier("refine_factor_table"),
+ vectorSchemaBuilder(VECTOR_FIELD_NAME)
+ .option(TestVectorGlobalIndexer.OPT_REVERSE_SCORE,
"true")
+ .build(),
+ false);
+ FileStoreTable table = getTable(identifier("refine_factor_table"));
+
+ float[][] vectors = {{0.0f, 0.0f}, {10.0f, 0.0f}, {20.0f, 0.0f}};
+ writeVectors(table, vectors);
+ buildAndCommitIndex(table, vectors);
+
+ GlobalIndexResult approximate =
+ table.newVectorSearchBuilder()
+ .withVector(new float[] {0.0f, 0.0f})
+ .withLimit(1)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .executeLocal();
+ assertThat(approximate.results()).containsExactly(2L);
+
+ GlobalIndexResult refined =
+ table.newVectorSearchBuilder()
+ .withVector(new float[] {0.0f, 0.0f})
+ .withLimit(1)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .withOption(
+ TestVectorGlobalIndexerFactory.IDENTIFIER +
".refine_factor", "3")
+ .executeLocal();
+ assertThat(refined.results()).containsExactly(0L);
+ assertThat(readIds(table, refined)).containsExactly(0);
+
+ List<GlobalIndexResult> batchRefined =
+ table.newBatchVectorSearchBuilder()
+ .withVectors(new float[][] {{0.0f, 0.0f}, {20.0f,
0.0f}})
+ .withLimit(1)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .withOption(
+ TestVectorGlobalIndexerFactory.IDENTIFIER +
".refine_factor", "3")
+ .executeBatchLocal();
+ assertThat(batchRefined).hasSize(2);
+ assertThat(batchRefined.get(0).results()).containsExactly(0L);
+ assertThat(batchRefined.get(1).results()).containsExactly(2L);
+ }
+
+ @Test
+ public void testVectorSearchRefineFactorValidation() throws Exception {
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+
+ float[][] vectors = {{0.0f, 0.0f}, {1.0f, 0.0f}};
+ writeVectors(table, vectors);
+ buildAndCommitIndex(table, vectors);
+
+ assertThatThrownBy(
+ () ->
+ table.newVectorSearchBuilder()
+ .withVector(new float[] {0.0f, 0.0f})
+ .withLimit(1)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .withOption("refine_factor", "0")
+ .executeLocal())
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("refine factor must be positive");
+ }
+
@Test
public void testVectorSearchWithMultipleIndexFiles() throws Exception {
createTableDefault();
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index faca1b71b9..b0a229f043 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -197,7 +197,7 @@ class AbstractVectorSearchReadImpl:
return reader, OffsetGlobalIndexReader(reader, row_range_start,
row_range_end)
def _eval(self, row_range_start, row_range_end, vector_index_files,
- query_vector, include_row_ids):
+ query_vector, search_limit, include_row_ids):
from pypaimon.globalindex.global_index_reader import _completed_future
if not vector_index_files:
@@ -205,7 +205,7 @@ class AbstractVectorSearchReadImpl:
vector_search = VectorSearch(
vector=query_vector,
- limit=self._limit,
+ limit=search_limit,
field_name=self._vector_column.name,
options=self._options,
)
@@ -262,7 +262,7 @@ class AbstractVectorSearchReadImpl:
return DictBasedScoredIndexResult(scores).top_k(self._limit)
def _eval_batch(self, row_range_start, row_range_end, vector_index_files,
- query_vectors, include_row_ids):
+ query_vectors, search_limit, include_row_ids):
from pypaimon.globalindex.global_index_reader import _completed_future
if not vector_index_files:
@@ -270,7 +270,7 @@ class AbstractVectorSearchReadImpl:
batch_vector_search = BatchVectorSearch(
vectors=query_vectors,
- limit=self._limit,
+ limit=search_limit,
field_name=self._vector_column.name,
options=self._options,
)
@@ -283,6 +283,42 @@ class AbstractVectorSearchReadImpl:
future.add_done_callback(lambda _: reader.close())
return future
+ def _indexed_search_limit(self, index_type):
+ refine_factor = self._configured_refine_factor(index_type)
+ if refine_factor == 0:
+ return self._limit
+ return self._limit * refine_factor
+
+ def _maybe_rerank_indexed_result(self, result, index_type, query_vector):
+ if (self._configured_refine_factor(index_type) == 0 or
+ 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(),
+ candidates.results(),
+ query_vector,
+ index_type,
+ )
+
+ def _configured_refine_factor(self, index_type):
+ value = _configured_refine_factor(
+ self._options, self._vector_column.name, index_type)
+ if value is None:
+ value = _configured_refine_factor(
+ _table_options_map(self._table), self._vector_column.name,
index_type)
+ if value is None:
+ return 0
+ try:
+ factor = int(value)
+ except ValueError as e:
+ raise ValueError(
+ "Invalid vector refine factor: %s. Must be an integer." % value
+ ) from e
+ if factor <= 0:
+ raise ValueError("Vector refine factor must be positive, got: %s"
% value)
+ return factor
+
class VectorSearchReadImpl(AbstractVectorSearchReadImpl, VectorSearchRead):
"""Implementation for VectorSearchRead."""
@@ -315,12 +351,15 @@ class VectorSearchReadImpl(AbstractVectorSearchReadImpl,
VectorSearchRead):
return indexed.or_(raw_result).top_k(self._limit)
def _read_indexed(self, splits, query_vector):
+ index_type = _vector_index_type(splits)
+ search_limit = self._indexed_search_limit(index_type)
pre_filters = self._pre_filters(splits)
futures = [
self._eval(
split.row_range_start, split.row_range_end,
split.vector_index_files,
query_vector,
+ search_limit,
None if not pre_filters else pre_filters[i]
)
for i, split in enumerate(splits)
@@ -337,7 +376,8 @@ class VectorSearchReadImpl(AbstractVectorSearchReadImpl,
VectorSearchRead):
if row_id not in merged_scores:
merged_scores[row_id] = score_getter(row_id)
- return DictBasedScoredIndexResult(merged_scores).top_k(self._limit)
+ indexed = DictBasedScoredIndexResult(merged_scores).top_k(search_limit)
+ return self._maybe_rerank_indexed_result(indexed, index_type,
query_vector)
class BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
@@ -361,11 +401,14 @@ class
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
# One native batch call per INDEX split (all query vectors at once),
# passing that split's pre-filter. Each future returns n per-query
results.
+ index_type = _vector_index_type(index_splits)
+ search_limit = self._indexed_search_limit(index_type)
pre_filters = self._pre_filters(index_splits)
futures = [
self._eval_batch(
split.row_range_start, split.row_range_end,
split.vector_index_files, self._query_vectors,
+ search_limit,
None if not pre_filters else pre_filters[i],
)
for i, split in enumerate(index_splits)
@@ -392,7 +435,9 @@ class
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
raw_index_type = _raw_search_index_type(raw_splits)
results = []
for i in range(n):
- indexed = DictBasedScoredIndexResult(merged_scores[i])
+ 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))
@@ -445,6 +490,13 @@ def _raw_search_index_type(raw_splits):
return None
+def _vector_index_type(index_splits):
+ for split in index_splits:
+ if split.vector_index_files:
+ return split.vector_index_files[0].index_type
+ return None
+
+
def _empty_bitmaps(size):
return [RoaringBitmap64() for _ in range(size)]
@@ -470,6 +522,45 @@ def _to_vector_list(value):
return list(value)
+def _configured_refine_factor(options, vector_column_name, index_type):
+ prefixes = []
+ field_prefix = "fields.%s." % vector_column_name
+ _add_refine_prefixes(prefixes, field_prefix, index_type)
+ _add_refine_prefixes(prefixes, "", index_type)
+
+ for prefix in prefixes:
+ for suffix in (
+ "refine_factor",
+ "refine-factor",
+ "rerank_factor",
+ "rerank-factor",
+ ):
+ value = options.get(prefix + suffix)
+ if value is not None:
+ return str(value).strip()
+ return None
+
+
+def _add_refine_prefixes(prefixes, base, index_type):
+ if index_type:
+ prefixes.append(base + index_type + ".")
+ normalized = _normalize_index_type(index_type)
+ if normalized != index_type:
+ prefixes.append(base + normalized + ".")
+ if normalized.startswith("ivf"):
+ prefixes.append(base + "ivf.")
+ prefixes.append(base)
+
+
+def _normalize_index_type(index_type):
+ return str(index_type).lower().replace("-", "_")
+
+
+def _table_options_map(table):
+ table_options = getattr(getattr(table, "options", None), "options", None)
+ return table_options.to_map() if table_options is not None else {}
+
+
def _raw_search_metric(table, vector_column, options, index_type=None):
candidates = []
field_prefix = "fields.%s." % vector_column.name
@@ -488,8 +579,7 @@ def _raw_search_metric(table, vector_column, options,
index_type=None):
]:
if key in options:
candidates.append(options[key])
- table_options = getattr(getattr(table, "options", None), "options", None)
- table_map = table_options.to_map() if table_options is not None else {}
+ table_map = _table_options_map(table)
for key in [
field_prefix + "distance.metric",
field_prefix + "metric",
diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py
b/paimon-python/pypaimon/tests/vector_search_filter_test.py
index d176d4f091..58f2fdc26e 100644
--- a/paimon-python/pypaimon/tests/vector_search_filter_test.py
+++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py
@@ -121,6 +121,70 @@ def _entry(partition_row, field_id, index_type, file_name,
index_file=index_file)
+def _install_raw_vector_read_builder(table, vector_column_name,
row_id_to_vector,
+ calls=None):
+ """Install a fake raw read builder which honors GlobalIndexResult
ranges."""
+ import pyarrow as pa
+
+ calls = calls if calls is not None else {}
+
+ class _Plan:
+ def __init__(self, row_ids):
+ self._row_ids = row_ids
+
+ def splits(self):
+ return list(self._row_ids)
+
+ class _Scan:
+ def __init__(self):
+ self._row_ids = []
+
+ def with_global_index_result(self, result):
+ ranges = result.results().to_range_list()
+ calls["global_index_ranges"] = ranges
+ self._row_ids = [
+ row_id
+ for row_id in sorted(row_id_to_vector)
+ if any(r.contains(row_id) for r in ranges)
+ ]
+ calls["candidate_ids"] = list(self._row_ids)
+ return self
+
+ def plan(self):
+ return _Plan(self._row_ids)
+
+ class _Read:
+ def to_arrow(self, splits):
+ row_ids = list(splits)
+ return pa.table({
+ vector_column_name: pa.array(
+ [row_id_to_vector[row_id] for row_id in row_ids]),
+ "_ROW_ID": pa.array(row_ids, type=pa.int64()),
+ })
+
+ class _Builder:
+ def with_partition_filter(self, predicate):
+ calls["partition_filter"] = predicate
+ return self
+
+ def with_filter(self, predicate):
+ calls["filter"] = predicate
+ return self
+
+ def with_projection(self, projection):
+ calls["projection"] = list(projection)
+ return self
+
+ def new_scan(self):
+ return _Scan()
+
+ def new_read(self):
+ return _Read()
+
+ table.new_read_builder = lambda: _Builder()
+ return calls
+
+
def _patch_snapshot(testcase, entries, snapshot=None):
"""Stub IndexFileHandler.scan + snapshot resolution."""
@@ -1300,6 +1364,164 @@ class VectorSearchFilterTest(unittest.TestCase):
search.options,
)
+ def test_refine_factor_reranks_index_candidates_with_raw_vectors(self):
+ from pypaimon.globalindex.vector_search_result import (
+ DictBasedScoredIndexResult,
+ )
+
+ entry = _entry(None, field_id=1, index_type="ivf-pq",
+ file_name="vec.index", row_range_start=0,
+ row_range_end=2)
+ table = _StubTable(fields=[self.id_field, self.embedding_field],
+ entries=[entry])
+ _patch_snapshot(self, [entry])
+ raw_calls = _install_raw_vector_read_builder(
+ table, "embedding", {0: [0.0], 1: [10.0], 2: [20.0]})
+ captured_limits = []
+
+ def _fake_create(index_type, file_io, index_path,
+ index_io_meta_list, options=None):
+ class _FakeReader:
+ def visit_vector_search(self_inner, vs):
+ captured_limits.append(vs.limit)
+ approximate_scores = [(2, 100.0), (1, 50.0), (0, 1.0)]
+ return _completed_future(
+
DictBasedScoredIndexResult(dict(approximate_scores[:vs.limit])))
+
+ def close(self_inner):
+ pass
+
+ return _FakeReader()
+
+ with mock.patch(
+
"pypaimon.table.source.vector_search_read._create_vector_reader",
+ side_effect=_fake_create):
+ result = (
+ VectorSearchBuilderImpl(table)
+ .with_vector_column("embedding")
+ .with_query_vector([0.0])
+ .with_limit(1)
+ .with_option("ivf.refine_factor", "3")
+ .execute_local()
+ )
+
+ 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([0], sorted(list(result.results())))
+
+ def test_refine_factor_one_reranks_without_expanding_candidates(self):
+ from pypaimon.globalindex.vector_search_result import (
+ DictBasedScoredIndexResult,
+ )
+
+ entry = _entry(None, field_id=1, index_type="ivf-pq",
+ file_name="vec.index", row_range_start=0,
+ row_range_end=2)
+ table = _StubTable(fields=[self.id_field, self.embedding_field],
+ entries=[entry])
+ _patch_snapshot(self, [entry])
+ raw_calls = _install_raw_vector_read_builder(
+ table, "embedding", {0: [0.0], 1: [10.0], 2: [20.0]})
+ captured_limits = []
+
+ def _fake_create(index_type, file_io, index_path,
+ index_io_meta_list, options=None):
+ class _FakeReader:
+ def visit_vector_search(self_inner, vs):
+ captured_limits.append(vs.limit)
+ return _completed_future(
+ DictBasedScoredIndexResult({2: 100.0}))
+
+ def close(self_inner):
+ pass
+
+ return _FakeReader()
+
+ with mock.patch(
+
"pypaimon.table.source.vector_search_read._create_vector_reader",
+ side_effect=_fake_create):
+ result = (
+ VectorSearchBuilderImpl(table)
+ .with_vector_column("embedding")
+ .with_query_vector([0.0])
+ .with_limit(1)
+ .with_option("ivf.refine_factor", "1")
+ .execute_local()
+ )
+
+ self.assertEqual([1], captured_limits)
+ self.assertEqual([Range(2, 2)], raw_calls["global_index_ranges"])
+ self.assertEqual([2], raw_calls["candidate_ids"])
+ self.assertEqual([2], sorted(list(result.results())))
+ self.assertLess(result.score_getter()(2), 1.0)
+
+ def test_refine_factor_query_options_override_table_options(self):
+ from pypaimon.common.options.options import Options
+ from pypaimon.globalindex.vector_search_result import (
+ DictBasedScoredIndexResult,
+ )
+
+ class _Options:
+ options = Options({"ivf.refine_factor": "2"})
+
+ entry = _entry(None, field_id=1, index_type="ivf-pq",
+ file_name="vec.index", row_range_start=0,
+ row_range_end=9)
+ table = _StubTable(fields=[self.id_field, self.embedding_field],
+ entries=[entry])
+ table.options = _Options()
+ _patch_snapshot(self, [entry])
+ _install_raw_vector_read_builder(
+ table, "embedding", {i: [float(i)] for i in range(10)})
+ captured_limits = []
+
+ def _fake_create(index_type, file_io, index_path,
+ index_io_meta_list, options=None):
+ class _FakeReader:
+ def visit_vector_search(self_inner, vs):
+ captured_limits.append(vs.limit)
+ return _completed_future(
+ DictBasedScoredIndexResult(
+ {i: float(i) for i in range(vs.limit)}))
+
+ def close(self_inner):
+ pass
+
+ return _FakeReader()
+
+ with mock.patch(
+
"pypaimon.table.source.vector_search_read._create_vector_reader",
+ side_effect=_fake_create):
+ (
+ VectorSearchBuilderImpl(table)
+ .with_vector_column("embedding")
+ .with_query_vector([0.0])
+ .with_limit(1)
+ .with_option("ivf.refine_factor", "3")
+ .execute_local()
+ )
+
+ self.assertEqual([3], captured_limits)
+
+ def test_refine_factor_validation(self):
+ entry = _entry(None, field_id=1, index_type="ivf-pq",
+ file_name="vec.index", row_range_start=0,
+ row_range_end=2)
+ table = _StubTable(fields=[self.id_field, self.embedding_field],
+ entries=[entry])
+ _patch_snapshot(self, [entry])
+
+ with self.assertRaisesRegex(ValueError, "refine factor must be
positive"):
+ (
+ VectorSearchBuilderImpl(table)
+ .with_vector_column("embedding")
+ .with_query_vector([0.0])
+ .with_limit(1)
+ .with_option("refine_factor", "0")
+ .execute_local()
+ )
+
def test_scanner_threads_external_path_to_btree_reader(self):
"""GlobalIndexScanner (backing _pre_filter) must thread external_path
onto the GlobalIndexIOMeta handed to the btree reader factory."""
@@ -2520,6 +2742,58 @@ class BatchVectorSearchTest(unittest.TestCase):
for i, query_vector in enumerate(query_vectors):
self.assertTrue(results[i].results().contains(int(query_vector[0])))
+ def test_batch_refine_factor_reranks_each_query(self):
+ from pypaimon.globalindex.global_index_reader import GlobalIndexReader
+ from pypaimon.globalindex.vector_search_result import (
+ DictBasedScoredIndexResult,
+ )
+ from pypaimon.table.source.batch_vector_search_builder import (
+ BatchVectorSearchBuilderImpl,
+ )
+
+ embedding_field = _field(1, "embedding", "FLOAT")
+ entry = _entry(None, field_id=1, index_type="ivf-pq",
+ file_name="vec.index",
+ 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]})
+ captured_limits = []
+
+ def _fake_create(index_type, file_io, index_path,
+ index_io_meta_list, options=None):
+ 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)
+ ])
+
+ def close(self_inner):
+ pass
+
+ return _FakeReader()
+
+ with mock.patch(
+
"pypaimon.table.source.vector_search_read._create_vector_reader",
+ side_effect=_fake_create):
+ results = (
+ BatchVectorSearchBuilderImpl(table)
+ .with_vector_column("embedding")
+ .with_query_vectors([[0.0], [20.0]])
+ .with_limit(1)
+ .with_option("ivf.refine_factor", "3")
+ .execute_batch_local()
+ )
+
+ self.assertEqual([3], captured_limits)
+ self.assertEqual([0], sorted(list(results[0].results())))
+ self.assertEqual([2], sorted(list(results[1].results())))
+
def test_batch_empty_splits_returns_empty_per_query(self):
from pypaimon.table.source.batch_vector_search_builder import (
BatchVectorSearchBuilderImpl,
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
index 90b1eef331..1b51e9cce1 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
@@ -103,7 +103,8 @@ public class SparkVectorReadImpl extends VectorReadImpl {
}
List<RoaringNavigableMap64> preFilters = preFilters(splits);
- String indexType = splits.get(0).vectorIndexFiles().get(0).indexType();
+ String indexType = vectorIndexType(splits);
+ int searchLimit = indexedSearchLimit(indexType);
List<SerializedSplit> serializedSplits = new
ArrayList<>(splits.size());
for (int i = 0; i < splits.size(); i++) {
try {
@@ -143,6 +144,7 @@ public class SparkVectorReadImpl extends VectorReadImpl {
split.rowRangeEnd(),
split.vectorIndexFiles(),
vector,
+ searchLimit,
deserializePreFilter(serializedSplit.preFilter),
executor));
}
@@ -154,7 +156,7 @@ public class SparkVectorReadImpl extends VectorReadImpl {
result = result.or(next.get());
}
}
- result = result.topK(limit);
+ result = result.topK(searchLimit);
if (result.results().isEmpty()) {
return null;
}
@@ -168,7 +170,13 @@ public class SparkVectorReadImpl extends VectorReadImpl {
List<byte[]> remoteResults = mapInSpark(splitGroups, task,
splitGroups.size());
- return mergeRemoteResults(remoteResults);
+ GlobalIndexer rerankGlobalIndexer =
+ globalIndexer == null ? createGlobalIndexer(splits) :
globalIndexer;
+ return maybeRerankIndexedResult(
+ mergeRemoteResults(remoteResults, searchLimit),
+ indexType,
+ rerankGlobalIndexer,
+ vector);
}
protected ScoredGlobalIndexResult readRawSplitsInSpark(
@@ -290,6 +298,10 @@ public class SparkVectorReadImpl extends VectorReadImpl {
}
private ScoredGlobalIndexResult mergeRemoteResults(List<byte[]>
remoteResults) {
+ return mergeRemoteResults(remoteResults, limit);
+ }
+
+ private ScoredGlobalIndexResult mergeRemoteResults(List<byte[]>
remoteResults, int topK) {
ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
GlobalIndexResultSerializer serializer = new
GlobalIndexResultSerializer();
for (byte[] bytes : remoteResults) {
@@ -301,7 +313,7 @@ public class SparkVectorReadImpl extends VectorReadImpl {
}
}
}
- return result.topK(limit);
+ return result.topK(topK);
}
private List<List<Range>> rangeGroups(List<Range> ranges, int parallelism)
{
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 c9a3e330a0..16b998ced6 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
@@ -18,9 +18,18 @@
package org.apache.paimon.spark.read;
+import org.apache.paimon.globalindex.GlobalIndexIOMeta;
+import org.apache.paimon.globalindex.GlobalIndexReader;
import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.globalindex.GlobalIndexResultSerializer;
+import org.apache.paimon.globalindex.GlobalIndexWriter;
import org.apache.paimon.globalindex.GlobalIndexer;
import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
+import org.apache.paimon.globalindex.VectorGlobalIndexer;
+import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
+import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
+import org.apache.paimon.index.GlobalIndexMeta;
+import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.table.source.IndexVectorSearchSplit;
import org.apache.paimon.table.source.RawVectorSearchSplit;
import org.apache.paimon.table.source.VectorScan;
@@ -30,13 +39,18 @@ import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.utils.Range;
import org.apache.paimon.utils.RoaringNavigableMap64;
+import org.apache.paimon.utils.SerializableFunction;
import org.junit.jupiter.api.Test;
import javax.annotation.Nullable;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
+import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
@@ -76,6 +90,32 @@ public class SparkVectorReadImplTest {
assertThat(result.results().getLongCardinality()).isEqualTo(64);
}
+ @Test
+ public void testDistributedIndexRefinesAfterGlobalMerge() {
+ DistributedRefineSparkVectorRead read = new
DistributedRefineSparkVectorRead();
+
+ ScoredGlobalIndexResult result =
+ read.readIndexSplitsInSpark(indexSplits("test-vector-ann", 4),
new L2Indexer());
+
+ assertThat(read.sparkParallelism).isEqualTo(2);
+ assertThat(read.rawSearchCandidateRows).containsExactly(0L, 2L);
+ assertThat(result.results().getLongCardinality()).isEqualTo(1);
+ assertThat(result.results().contains(0L)).isTrue();
+ }
+
+ private static List<IndexVectorSearchSplit> indexSplits(String indexType,
int count) {
+ List<IndexVectorSearchSplit> splits = new ArrayList<>();
+ for (int i = 0; i < count; i++) {
+ GlobalIndexMeta globalIndexMeta = new GlobalIndexMeta(i, i, 0,
null, new byte[0]);
+ IndexFileMeta indexFile =
+ new IndexFileMeta(indexType, "index-" + i, 1L, 1L,
globalIndexMeta, null);
+ splits.add(
+ new IndexVectorSearchSplit(
+ i, i, Collections.singletonList(indexFile),
Collections.emptyList()));
+ }
+ return splits;
+ }
+
private static class TestingSparkVectorRead extends SparkVectorReadImpl {
private boolean rawSparkPathUsed;
@@ -118,6 +158,78 @@ public class SparkVectorReadImplTest {
}
}
+ private static class DistributedRefineSparkVectorRead extends
SparkVectorReadImpl {
+
+ private int sparkParallelism;
+ private List<Long> rawSearchCandidateRows = Collections.emptyList();
+
+ private DistributedRefineSparkVectorRead() {
+ super(
+ null,
+ null,
+ null,
+ 1,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[] {0.0f},
+ Collections.singletonMap("refine_factor", "2"));
+ }
+
+ @Override
+ protected int sparkParallelism() {
+ return 2;
+ }
+
+ @Override
+ protected <I, O> List<O> mapInSpark(
+ List<I> data, SerializableFunction<I, O> func, int
parallelism) {
+ sparkParallelism = parallelism;
+ assertThat(data).hasSize(2);
+ try {
+ GlobalIndexResultSerializer serializer = new
GlobalIndexResultSerializer();
+ return Arrays.asList(
+ uncheckedCast(serializer.serialize(scoredResult(2L,
100.0f))),
+ uncheckedCast(serializer.serialize(scoredResult(0L,
1.0f))));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ protected ScoredGlobalIndexResult readRawSearch(
+ List<Range> rawRowRanges,
+ @Nullable RoaringNavigableMap64 preFilter,
+ @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) {
+ rawSearchCandidateRows.add(rowId);
+ }
+
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ for (long rowId : preFilter) {
+ rows.add(rowId);
+ }
+ return ScoredGlobalIndexResult.create(
+ rows, rowId -> rowId == 0L ? 1.0f : 1.0f / (1.0f +
rowId * rowId))
+ .topK(1);
+ }
+
+ @SuppressWarnings("unchecked")
+ private <O> O uncheckedCast(byte[] value) {
+ return (O) value;
+ }
+
+ private static ScoredGlobalIndexResult scoredResult(long rowId, float
score) {
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ rows.add(rowId);
+ return ScoredGlobalIndexResult.create(rows, candidate -> score);
+ }
+ }
+
private static class RecordingSparkVectorRead extends SparkVectorReadImpl {
private final AtomicInteger nextTask = new AtomicInteger();
@@ -169,4 +281,25 @@ public class SparkVectorReadImplTest {
return ScoredGlobalIndexResult.create(rows, rowId -> scoreBase +
(float) rowId);
}
}
+
+ private static class L2Indexer implements VectorGlobalIndexer {
+
+ @Override
+ public GlobalIndexWriter createWriter(GlobalIndexFileWriter
fileWriter) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public GlobalIndexReader createReader(
+ GlobalIndexFileReader fileReader,
+ List<GlobalIndexIOMeta> files,
+ ExecutorService executor) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public String metric() {
+ return "l2";
+ }
+ }
}
