This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 7f2895b3776 [Vector Upsert 1/5] Support allowed-document filtering in
exact vector scan (#19297)
7f2895b3776 is described below
commit 7f2895b377668b2ee7722f69f7d2d052e3722926
Author: Xiang Fu <[email protected]>
AuthorDate: Wed Aug 26 12:34:59 2026 -0700
[Vector Upsert 1/5] Support allowed-document filtering in exact vector scan
(#19297)
Extract a shared computeExactMatches helper that unifies the exact top-K and
threshold scan paths and can restrict scoring to an allowed-document bitmap.
Introduce VectorCandidateScope, an immutable set of document IDs a vector
predicate is allowed to consider as candidates. Vector top-K is not
monotonic,
so unlike an ordinary predicate -- which is correct to intersect with the
result
afterwards -- a document set that defines what the query may see has to be
applied before candidate generation. The type carries that contract and says
nothing about where the restriction came from; callers decide that.
ExactVectorScanFilterOperator accepts one via a new constructor overload,
and no
production caller passes one yet.
Also let the operator record its own exact-scan search and fallback metrics,
which were previously invisible to VectorSearchMetrics, and report the
applied
candidate-filter cardinality in explain output.
Add filtered-reader coverage for HnswVectorIndexReader.
---
.../filter/ExactVectorScanFilterOperator.java | 188 ++++++++++++---------
.../org/apache/pinot/core/plan/FilterPlanNode.java | 2 +-
.../filter/ExactVectorScanFilterOperatorTest.java | 152 ++++++++++++++++-
.../index/creator/HnswVectorIndexCreatorTest.java | 16 ++
4 files changed, 273 insertions(+), 85 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperator.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperator.java
index eaf0367ba18..958e72a7287 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperator.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperator.java
@@ -19,6 +19,7 @@
package org.apache.pinot.core.operator.filter;
import com.google.common.base.CaseFormat;
+import com.google.common.base.Preconditions;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;
@@ -36,6 +37,7 @@ import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.trace.FilterType;
import org.apache.pinot.spi.trace.InvocationRecording;
import org.apache.pinot.spi.trace.Tracing;
+import org.roaringbitmap.IntIterator;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.slf4j.Logger;
@@ -46,11 +48,11 @@ import org.slf4j.LoggerFactory;
///
/// This operator is used when no ANN vector index exists on a segment for the
target column
/// (e.g., the segment was built before the vector index was added, or the
index type is not
-/// supported). It reads all vectors from the forward index, computes exact
distances to the
-/// query vector, and returns the top-K closest document IDs.
+/// supported). It reads vectors from the forward index, computes exact
distances to the query vector, and returns the
+/// top-K closest document IDs. A mandatory candidate bitmap restricts the
scan when present.
///
-/// The distance computation uses L2 (Euclidean) squared distance. For COSINE
similarity,
-/// vectors should be pre-normalized. This matches the behavior of Lucene's
HNSW implementation.
+/// Distance computation uses the function configured for the vector index, or
Pinot's default when no index
+/// configuration is available.
///
/// This operator is intentionally simple and correct rather than fast -- it
is a safety net.
/// A warning is logged when this operator is used because it scans all
documents in the segment.
@@ -66,6 +68,10 @@ public class ExactVectorScanFilterOperator extends
BaseFilterOperator {
private final VectorExplainContext _vectorExplainContext;
private final boolean _hasDistanceThreshold;
private final float _distanceThreshold;
+ /// Documents this scan is allowed to consider, or null to scan the whole
segment. Held by reference: the query
+ /// owns this bitmap and does not modify it while the plan executes.
+ @Nullable
+ private final ImmutableRoaringBitmap _requiredDocIds;
private ImmutableRoaringBitmap _matches;
/// Creates an exact scan operator.
@@ -74,28 +80,21 @@ public class ExactVectorScanFilterOperator extends
BaseFilterOperator {
/// @param predicate the vector similarity predicate containing query vector
and top-K
/// @param column the column name (for logging and explain)
/// @param numDocs the total number of documents in the segment
- public ExactVectorScanFilterOperator(ForwardIndexReader<?>
forwardIndexReader,
- VectorSimilarityPredicate predicate, String column, int numDocs) {
- this(forwardIndexReader, predicate, column, numDocs, null,
"vector_index_missing",
- VectorSearchParams.DEFAULT);
- }
-
- public ExactVectorScanFilterOperator(ForwardIndexReader<?>
forwardIndexReader,
- VectorSimilarityPredicate predicate, String column, int numDocs,
@Nullable VectorIndexConfig vectorIndexConfig,
- String fallbackReason) {
- this(forwardIndexReader, predicate, column, numDocs, vectorIndexConfig,
fallbackReason,
- VectorSearchParams.DEFAULT);
- }
-
- public ExactVectorScanFilterOperator(ForwardIndexReader<?>
forwardIndexReader,
- VectorSimilarityPredicate predicate, String column, int numDocs,
@Nullable VectorIndexConfig vectorIndexConfig,
- String fallbackReason, VectorSearchParams searchParams) {
+ /// @param vectorIndexConfig vector index configuration, used to resolve the
distance function (may be null)
+ /// @param fallbackReason why this scan runs instead of an ANN search,
reported in explain output
+ /// @param searchParams vector search parameters from query options
+ /// @param requiredDocIds documents this scan may consider, or null to scan
the whole segment. The caller must not
+ /// modify the bitmap afterwards.
+ public ExactVectorScanFilterOperator(ForwardIndexReader<?>
forwardIndexReader, VectorSimilarityPredicate predicate,
+ String column, int numDocs, @Nullable VectorIndexConfig
vectorIndexConfig, String fallbackReason,
+ VectorSearchParams searchParams, @Nullable ImmutableRoaringBitmap
requiredDocIds) {
super(numDocs, false);
_forwardIndexReader = forwardIndexReader;
_predicate = predicate;
_column = column;
_hasDistanceThreshold = searchParams.hasDistanceThreshold();
_distanceThreshold = searchParams.getDistanceThreshold();
+ _requiredDocIds = requiredDocIds;
float effectiveThreshold = _hasDistanceThreshold ? _distanceThreshold :
-1f;
_vectorExplainContext = new
VectorExplainContext(VectorDistanceUtils.resolveBackendType(vectorIndexConfig),
VectorDistanceUtils.resolveDistanceFunction(vectorIndexConfig),
VectorExecutionMode.EXACT_SCAN,
@@ -149,6 +148,9 @@ public class ExactVectorScanFilterOperator extends
BaseFilterOperator {
+ ", vector literal:" + Arrays.toString(_predicate.getValue())
+ ", topK to search:" + _predicate.getTopK()
+ ", fallbackReason:" + _vectorExplainContext.getFallbackReason()
+ + (_requiredDocIds != null
+ ? ", requiredDocIdFilterApplied:true,
requiredDocIdFilterCardinality:" + _requiredDocIds.getCardinality()
+ : "")
+ ')';
}
@@ -169,55 +171,42 @@ public class ExactVectorScanFilterOperator extends
BaseFilterOperator {
attributeBuilder.putString("vectorLiteral",
Arrays.toString(_predicate.getValue()));
attributeBuilder.putString("fallbackReason",
_vectorExplainContext.getFallbackReason());
attributeBuilder.putLongIdempotent("topKtoSearch", _predicate.getTopK());
+ if (_requiredDocIds != null) {
+ attributeBuilder.putBool("requiredDocIdFilterApplied", true);
+ attributeBuilder.putLong("requiredDocIdFilterCardinality",
_requiredDocIds.getCardinality());
+ }
}
/// Performs brute-force exact search over all documents in the segment.
/// When a distance threshold is set, returns all vectors within the
threshold.
/// Otherwise uses a max-heap to maintain the top-K closest vectors.
- @SuppressWarnings("unchecked")
private ImmutableRoaringBitmap computeExactTopK() {
- LOGGER.warn("Performing exact vector scan fallback on column: {} for
segment with {} docs. "
- + "reason={}, distanceFunction={}, hasThreshold={}. "
- + "This is expensive -- consider adding a vector index.",
- _column, _numDocs, _vectorExplainContext.getFallbackReason(),
- _vectorExplainContext.getDistanceFunction(), _hasDistanceThreshold);
-
- float[] queryVector = _predicate.getValue();
-
- if (_hasDistanceThreshold) {
- return computeExactThreshold(queryVector);
+ ImmutableRoaringBitmap allowedDocIds = _requiredDocIds;
+ VectorSearchMetrics.getInstance().recordSearch(VectorSearchMode.EXACT_SCAN,
+ _vectorExplainContext.getBackendType());
+
VectorSearchMetrics.getInstance().recordFallback(_vectorExplainContext.getFallbackReason());
+ if (allowedDocIds != null && allowedDocIds.isEmpty()) {
+ return new MutableRoaringBitmap();
}
-
- int topK = _predicate.getTopK();
-
- // Max-heap: entry with largest distance is at the top so we can
efficiently evict it
- PriorityQueue<DocDistance> maxHeap = new PriorityQueue<>(topK + 1,
- (a, b) -> Float.compare(b._distance, a._distance));
-
- ForwardIndexReader rawReader = _forwardIndexReader;
- try (ForwardIndexReaderContext context = rawReader.createContext()) {
- for (int docId = 0; docId < _numDocs; docId++) {
- float[] docVector = rawReader.getFloatMV(docId, context);
- if (docVector == null || docVector.length == 0) {
- continue;
- }
- float distance = VectorDistanceUtils.computeDistance(queryVector,
docVector,
- _vectorExplainContext.getDistanceFunction());
- if (maxHeap.size() < topK) {
- maxHeap.add(new DocDistance(docId, distance));
- } else if (distance < maxHeap.peek()._distance) {
- maxHeap.poll();
- maxHeap.add(new DocDistance(docId, distance));
- }
- }
- } catch (Exception e) {
- throw new RuntimeException("Error during exact vector scan on column: "
+ _column, e);
+ if (allowedDocIds != null) {
+ // Scanning a restricted candidate set is the steady-state path whenever
the segment's vector index cannot do
+ // filtered search, so it is expected rather than anomalous. The
recorded fallback metric carries the signal.
+ LOGGER.debug("Performing exact vector scan over {} allowed docs on
column: {}. reason={}, distanceFunction={}, "
+ + "hasThreshold={}.",
+ allowedDocIds.getCardinality(), _column,
_vectorExplainContext.getFallbackReason(),
+ _vectorExplainContext.getDistanceFunction(), _hasDistanceThreshold);
+ } else {
+ LOGGER.warn("Performing exact vector scan fallback on column: {} for
segment with {} docs. "
+ + "reason={}, distanceFunction={}, hasThreshold={}. "
+ + "This is expensive -- consider adding a vector index.",
+ _column, _numDocs, _vectorExplainContext.getFallbackReason(),
+ _vectorExplainContext.getDistanceFunction(), _hasDistanceThreshold);
}
- MutableRoaringBitmap result = new MutableRoaringBitmap();
- for (DocDistance dd : maxHeap) {
- result.add(dd._docId);
- }
+ float[] queryVector = _predicate.getValue();
+ Float threshold = _hasDistanceThreshold ? _distanceThreshold : null;
+ ImmutableRoaringBitmap result = computeExactMatches(_forwardIndexReader,
queryVector, _predicate.getTopK(),
+ _numDocs, _vectorExplainContext.getDistanceFunction(), threshold,
allowedDocIds, _column);
LOGGER.debug("Exact vector scan on column: {} returned {} results from {}
docs",
_column, result.getCardinality(), _numDocs);
@@ -225,33 +214,80 @@ public class ExactVectorScanFilterOperator extends
BaseFilterOperator {
return result;
}
- /// Performs brute-force threshold scan: returns all vectors within the
distance threshold.
- @SuppressWarnings("unchecked")
- private ImmutableRoaringBitmap computeExactThreshold(float[] queryVector) {
+ /// Performs an exact top-K or threshold search over all documents or the
supplied allowed-document bitmap.
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ private static ImmutableRoaringBitmap
computeExactMatches(ForwardIndexReader<?> forwardIndexReader,
+ float[] queryVector, int topK, int numDocs,
+ VectorIndexConfig.VectorDistanceFunction distanceFunction, @Nullable
Float distanceThreshold,
+ @Nullable ImmutableRoaringBitmap allowedDocIds, String column) {
+ // Threshold search ignores topK; top-K search rejects non-positive values
the same way the IVF readers do, so
+ // one query cannot fail on an indexed segment and quietly return nothing
on a segment that falls back here.
+ if (distanceThreshold == null) {
+ Preconditions.checkArgument(topK > 0, "topK must be positive, got: %s",
topK);
+ }
+ if (allowedDocIds != null && allowedDocIds.isEmpty()) {
+ return new MutableRoaringBitmap();
+ }
+
+ PriorityQueue<DocDistance> maxHeap = distanceThreshold == null
+ ? new PriorityQueue<>(topK + 1, (a, b) -> Float.compare(b._distance,
a._distance)) : null;
MutableRoaringBitmap result = new MutableRoaringBitmap();
- ForwardIndexReader rawReader = _forwardIndexReader;
+ ForwardIndexReader rawReader = forwardIndexReader;
try (ForwardIndexReaderContext context = rawReader.createContext()) {
- for (int docId = 0; docId < _numDocs; docId++) {
- float[] docVector = rawReader.getFloatMV(docId, context);
- if (docVector == null || docVector.length == 0) {
- continue;
+ if (allowedDocIds == null) {
+ for (int docId = 0; docId < numDocs; docId++) {
+ scoreDocument(rawReader, context, docId, queryVector, topK,
distanceFunction, distanceThreshold,
+ maxHeap, result);
}
- float distance = VectorDistanceUtils.computeDistance(queryVector,
docVector,
- _vectorExplainContext.getDistanceFunction());
- if (distance <= _distanceThreshold) {
- result.add(docId);
+ } else {
+ IntIterator iterator = allowedDocIds.getIntIterator();
+ while (iterator.hasNext()) {
+ int docId = iterator.next();
+ // Query planning can race with later segment changes. Never read
beyond the numDocs watermark captured by
+ // FilterPlanNode, even if a supplied bitmap contains larger
document IDs.
+ if (docId < 0 || docId >= numDocs) {
+ break;
+ }
+ scoreDocument(rawReader, context, docId, queryVector, topK,
distanceFunction, distanceThreshold,
+ maxHeap, result);
}
}
} catch (Exception e) {
- throw new RuntimeException("Error during exact threshold scan on column:
" + _column, e);
+ throw new RuntimeException("Error during exact vector scan on column: "
+ column, e);
}
- LOGGER.debug("Exact threshold scan on column: {} returned {} results from
{} docs (threshold={})",
- _column, result.getCardinality(), _numDocs, _distanceThreshold);
-
+ if (maxHeap != null) {
+ for (DocDistance docDistance : maxHeap) {
+ result.add(docDistance._docId);
+ }
+ }
return result.toImmutableRoaringBitmap();
}
+ private static void scoreDocument(ForwardIndexReader rawReader,
ForwardIndexReaderContext context, int docId,
+ float[] queryVector, int topK, VectorIndexConfig.VectorDistanceFunction
distanceFunction,
+ @Nullable Float distanceThreshold, @Nullable PriorityQueue<DocDistance>
maxHeap,
+ MutableRoaringBitmap thresholdMatches) {
+ float[] docVector = rawReader.getFloatMV(docId, context);
+ if (docVector == null || docVector.length == 0) {
+ return;
+ }
+ float distance = VectorDistanceUtils.computeDistance(queryVector,
docVector, distanceFunction);
+ if (distanceThreshold != null) {
+ if (distance <= distanceThreshold) {
+ thresholdMatches.add(docId);
+ }
+ return;
+ }
+
+ if (maxHeap.size() < topK) {
+ maxHeap.add(new DocDistance(docId, distance));
+ } else if (distance < maxHeap.peek()._distance) {
+ maxHeap.poll();
+ maxHeap.add(new DocDistance(docId, distance));
+ }
+ }
+
/// Computes the squared L2 (Euclidean) distance between two vectors.
/// Delegates to [VectorFunctions#euclideanDistance(float[], float[])] which
returns
/// the sum of squared differences (no sqrt), sufficient for ranking.
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java
b/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java
index dea76ec01de..a1a2024c909 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java
@@ -384,7 +384,7 @@ public class FilterPlanNode implements PlanNode {
Preconditions.checkState(forwardIndexReader != null,
"Cannot apply VECTOR_SIMILARITY on column: %s -- no vector index and
no forward index available", column);
return new ExactVectorScanFilterOperator(forwardIndexReader, predicate,
column, numDocs, vectorIndexConfig,
- getVectorFallbackReason(vectorIndexConfig, isMutableSegment),
searchParams);
+ getVectorFallbackReason(vectorIndexConfig, isMutableSegment),
searchParams, null);
}
/// Constructs a vector operator for a VECTOR_SIMILARITY predicate that is
part of an AND
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperatorTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperatorTest.java
index 6bd92ba0f32..8bfc40a9b80 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperatorTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperatorTest.java
@@ -27,10 +27,14 @@ import
org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.mockito.Mockito;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.testng.Assert;
import org.testng.annotations.Test;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
@@ -56,7 +60,7 @@ public class ExactVectorScanFilterOperatorTest {
VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(lhs,
queryVector, 2);
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
- "embedding", numDocs);
+ "embedding", numDocs, null, "vector_index_missing",
VectorSearchParams.DEFAULT, null);
// Should return doc 0 (distance=0) and doc 4 (distance=0.02)
ImmutableRoaringBitmap result = operator.getBitmaps().reduce();
@@ -81,7 +85,7 @@ public class ExactVectorScanFilterOperatorTest {
VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(lhs,
queryVector, 10);
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
- "embedding", numDocs);
+ "embedding", numDocs, null, "vector_index_missing",
VectorSearchParams.DEFAULT, null);
ImmutableRoaringBitmap result = operator.getBitmaps().reduce();
Assert.assertEquals(result.getCardinality(), 3);
@@ -90,6 +94,132 @@ public class ExactVectorScanFilterOperatorTest {
Assert.assertTrue(result.contains(2));
}
+ @Test
+ public void testExactSearchOnlyScoresAllowedDocuments() {
+ float[][] vectors = {
+ {1.0f, 0.0f},
+ {1.0f, 0.0f},
+ {0.0f, 1.0f},
+ {0.0f, -1.0f}
+ };
+ ForwardIndexReader<?> mockReader = createMockForwardIndexReader(vectors);
+ VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), new float[]{1.0f, 0.0f},
2);
+ ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
+ "embedding", 4, createVectorIndexConfig("HNSW",
VectorIndexConfig.VectorDistanceFunction.EUCLIDEAN),
+ "mutable_vector_index_not_filter_aware", VectorSearchParams.DEFAULT,
+ bitmapOf(2, 3));
+
+ Assert.assertEquals(operator.getBitmaps().reduce(), bitmapOf(2, 3));
+ ForwardIndexReader rawReader = mockReader;
+ verify(rawReader, never()).getFloatMV(Mockito.eq(0), Mockito.any());
+ verify(rawReader, never()).getFloatMV(Mockito.eq(1), Mockito.any());
+ verify(rawReader).getFloatMV(Mockito.eq(2), Mockito.any());
+ verify(rawReader).getFloatMV(Mockito.eq(3), Mockito.any());
+ }
+
+ @Test
+ public void testExactSearchWithEmptyAllowedDocumentsSkipsForwardIndex() {
+ ForwardIndexReader<?> mockReader = mock(ForwardIndexReader.class);
+ VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), new float[]{1.0f, 0.0f},
2);
+ ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
+ "embedding", 4, null, "vector_index_missing",
VectorSearchParams.DEFAULT,
+ new MutableRoaringBitmap());
+
+ Assert.assertTrue(operator.getBitmaps().reduce().isEmpty());
+ verifyNoInteractions(mockReader);
+ }
+
+ @Test
+ public void testExactSearchIgnoresAllowedDocumentsBeyondNumDocs() {
+ ForwardIndexReader<?> mockReader = createMockForwardIndexReader(new
float[][]{
+ {1.0f, 0.0f}, {0.5f, 0.5f}, {0.0f, 1.0f}
+ });
+ VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), new float[]{1.0f, 0.0f},
2);
+ ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
+ "embedding", 3, null, "mandatory_scope", VectorSearchParams.DEFAULT,
+ bitmapOf(2, 99));
+
+ Assert.assertEquals(operator.getBitmaps().reduce(), bitmapOf(2));
+ ForwardIndexReader rawReader = mockReader;
+ verify(rawReader).getFloatMV(Mockito.eq(2), Mockito.any());
+ verify(rawReader, never()).getFloatMV(Mockito.eq(99), Mockito.any());
+ }
+
+ /// topK comes straight from the query literal with no parse-time
validation, so a segment falling back to this
+ /// operator must reject non-positive values exactly like the IVF readers
do. Otherwise the same query errors on an
+ /// indexed segment and quietly returns nothing here.
+ @Test
+ public void testExactSearchRejectsNonPositiveTopK() {
+ ForwardIndexReader<?> mockReader = mock(ForwardIndexReader.class);
+
+ for (int topK : new int[]{0, -1}) {
+ VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), new float[]{1.0f,
0.0f}, topK);
+ ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
+ "embedding", 4, null, "no_vector_index", VectorSearchParams.DEFAULT,
null);
+ IllegalArgumentException exception =
+ Assert.expectThrows(IllegalArgumentException.class,
operator::getBitmaps);
+ Assert.assertEquals(exception.getMessage(), "topK must be positive, got:
" + topK);
+ }
+ verifyNoInteractions(mockReader);
+ }
+
+ /// Threshold search ignores topK entirely, so it must keep working when
topK is not meaningful.
+ @Test
+ public void testExactThresholdSearchIgnoresNonPositiveTopK() {
+ ForwardIndexReader<?> mockReader = createMockForwardIndexReader(new
float[][]{
+ {1.0f, 0.0f}, {0.0f, 1.0f}
+ });
+ VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), new float[]{1.0f, 0.0f},
0);
+ VectorSearchParams searchParams = new VectorSearchParams(null, null, null,
0.5f, null, null, null);
+ ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
+ "embedding", 2, null, "vector_index_missing", searchParams, null);
+
+ Assert.assertEquals(operator.getBitmaps().reduce(), bitmapOf(0));
+ }
+
+ @Test
+ public void testExactSearchTopKLargerThanAllowedCardinality() {
+ ForwardIndexReader<?> mockReader = createMockForwardIndexReader(new
float[][]{
+ {1.0f, 0.0f},
+ {0.5f, 0.5f},
+ {0.0f, 1.0f},
+ {0.0f, -1.0f}
+ });
+ VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), new float[]{1.0f, 0.0f},
10);
+ ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
+ "embedding", 4, null, "vector_index_missing",
VectorSearchParams.DEFAULT,
+ bitmapOf(2, 3));
+
+ Assert.assertEquals(operator.getBitmaps().reduce(), bitmapOf(2, 3));
+ }
+
+ @Test
+ public void testExactThresholdSearchOnlyScoresAllowedDocuments() {
+ ForwardIndexReader<?> mockReader = createMockForwardIndexReader(new
float[][]{
+ {1.0f, 0.0f},
+ {0.9f, 0.1f},
+ {0.0f, 1.0f},
+ {-1.0f, 0.0f}
+ });
+ VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), new float[]{1.0f, 0.0f},
10);
+ VectorSearchParams searchParams = new VectorSearchParams(null, null, null,
2.0f, null, null, null);
+ ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
+ "embedding", 4, null, "vector_index_missing", searchParams,
+ bitmapOf(2, 3));
+
+ Assert.assertEquals(operator.getBitmaps().reduce(), bitmapOf(2));
+ ForwardIndexReader rawReader = mockReader;
+ verify(rawReader, never()).getFloatMV(Mockito.eq(0), Mockito.any());
+ verify(rawReader, never()).getFloatMV(Mockito.eq(1), Mockito.any());
+ }
+
@Test
public void testL2SquaredDistance() {
float[] a = {1.0f, 2.0f, 3.0f};
@@ -122,7 +252,7 @@ public class ExactVectorScanFilterOperatorTest {
VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(lhs,
queryVector, 2);
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
- "embedding", numDocs);
+ "embedding", numDocs, null, "vector_index_missing",
VectorSearchParams.DEFAULT, null);
Assert.assertEquals(operator.getNumMatchingDocs(), 2);
}
@@ -133,7 +263,7 @@ public class ExactVectorScanFilterOperatorTest {
ExpressionContext lhs = ExpressionContext.forIdentifier("embedding");
VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(lhs,
new float[]{1.0f}, 1);
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
- "embedding", 1);
+ "embedding", 1, null, "vector_index_missing",
VectorSearchParams.DEFAULT, null);
Assert.assertTrue(operator.canProduceBitmaps());
}
@@ -144,7 +274,7 @@ public class ExactVectorScanFilterOperatorTest {
VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(lhs,
new float[]{1.0f, 2.0f}, 5);
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
"embedding", 1, createVectorIndexConfig("IVF_PQ",
VectorIndexConfig.VectorDistanceFunction.COSINE),
- "ivf_pq_index_unavailable");
+ "ivf_pq_index_unavailable", VectorSearchParams.DEFAULT, null);
String explain = operator.toExplainString();
Assert.assertTrue(explain.contains("exact_scan"));
Assert.assertTrue(explain.contains("embedding"));
@@ -167,7 +297,7 @@ public class ExactVectorScanFilterOperatorTest {
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
"embedding", 2, createVectorIndexConfig("IVF_PQ",
VectorIndexConfig.VectorDistanceFunction.COSINE),
- "ivf_pq_index_unavailable");
+ "ivf_pq_index_unavailable", VectorSearchParams.DEFAULT, null);
ImmutableRoaringBitmap result = operator.getBitmaps().reduce();
Assert.assertEquals(result.getCardinality(), 1);
@@ -188,7 +318,7 @@ public class ExactVectorScanFilterOperatorTest {
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
"embedding", 2, createVectorIndexConfig("IVF_PQ",
VectorIndexConfig.VectorDistanceFunction.INNER_PRODUCT),
- "ivf_pq_index_unavailable");
+ "ivf_pq_index_unavailable", VectorSearchParams.DEFAULT, null);
ImmutableRoaringBitmap result = operator.getBitmaps().reduce();
Assert.assertEquals(result.getCardinality(), 1);
@@ -208,7 +338,7 @@ public class ExactVectorScanFilterOperatorTest {
VectorSimilarityPredicate predicate = new VectorSimilarityPredicate(lhs,
queryVector, 1);
ExactVectorScanFilterOperator operator = new
ExactVectorScanFilterOperator(mockReader, predicate,
- "embedding", 2);
+ "embedding", 2, null, "vector_index_missing",
VectorSearchParams.DEFAULT, null);
ImmutableRoaringBitmap result = operator.getBitmaps().reduce();
Assert.assertEquals(result.getCardinality(), 1);
@@ -239,4 +369,10 @@ public class ExactVectorScanFilterOperatorTest {
return new VectorIndexConfig(false, backendType, 2, 1, distanceFunction,
Map.of("nlist", "4", "pqM", "2", "pqNbits", "8", "trainSampleSize",
"16"));
}
+
+ private static MutableRoaringBitmap bitmapOf(int... docIds) {
+ MutableRoaringBitmap bitmap = new MutableRoaringBitmap();
+ bitmap.add(docIds);
+ return bitmap;
+ }
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/HnswVectorIndexCreatorTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/HnswVectorIndexCreatorTest.java
index aacdb4e3980..0c42ec0f9e6 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/HnswVectorIndexCreatorTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/HnswVectorIndexCreatorTest.java
@@ -26,6 +26,7 @@ import org.apache.commons.io.FileUtils;
import
org.apache.pinot.segment.local.segment.creator.impl.vector.HnswVectorIndexCreator;
import
org.apache.pinot.segment.local.segment.index.readers.vector.HnswVectorIndexReader;
import org.apache.pinot.segment.spi.index.creator.VectorIndexConfig;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
@@ -93,6 +94,21 @@ public class HnswVectorIndexCreatorTest {
}
}
+ @Test
+ public void testFilteredReaderReturnsKAllowedDocuments()
+ throws IOException {
+ MutableRoaringBitmap allowedDocIds = new MutableRoaringBitmap();
+ allowedDocIds.add(1);
+ allowedDocIds.add(3);
+ try (HnswVectorIndexReader reader = new HnswVectorIndexReader("foo",
INDEX_DIR, 4, _config)) {
+ float[] queryVector = {5.0F, 42.0F, 54.33333F, 42.24F, 1001.045F};
+ Assert.assertEquals(reader.getDocIds(queryVector, 1).toArray(), new
int[]{0},
+ "The nearest physical document should be outside the allowed set");
+ int[] matchedDocIds = reader.getDocIds(queryVector, 2,
allowedDocIds).toArray();
+ Assert.assertEquals(matchedDocIds, new int[]{1, 3});
+ }
+ }
+
@Test
public void testEfSearchChangesRuntimeSearchBehavior()
throws IOException {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
