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 20e3397efb9 Enable filtered vector search on mutable HNSW and share
the doc-id filter query (#19303)
20e3397efb9 is described below
commit 20e3397efb9b26a6fb3c8940c512ddd48028e300
Author: Xiang Fu <[email protected]>
AuthorDate: Fri Sep 4 03:30:02 2026 -0700
Enable filtered vector search on mutable HNSW and share the doc-id filter
query (#19303)
FULL-upsert vector search must apply the visible-document scope during
candidate generation so obsolete rows cannot consume top-K slots. Make mutable
HNSW honor that contract while consuming-segment rows are still uncommitted.
Store Pinot document IDs as numeric doc values for filter membership and
hit translation, and share the Lucene bitmap-filter query with immutable HNSW.
Use a shared near-real-time SearcherManager and coalesce concurrent
refreshes by exact writer generation, including out-of-order Pinot document IDs.
Choose an exact forward-index scan for sparse mutable required scopes when
available; retain filtered ANN for dense scopes and when the forward index is
disabled.
Avoid eager strategy explanation formatting for mode-only planning, return
empty required scopes before search work, and preserve interruption as the
IOException cause.
Cover planner routing, refresh coalescing, live-writer visibility,
shared-query behavior, realtime filtered HNSW, and FULL-upsert integration.
---
.../core/operator/filter/VectorSearchStrategy.java | 48 +-
.../org/apache/pinot/core/plan/FilterPlanNode.java | 27 +-
.../operator/filter/VectorSearchStrategyTest.java | 28 ++
.../apache/pinot/core/plan/FilterPlanNodeTest.java | 128 ++++++
.../tests/custom/HnswVectorRealtimeTest.java | 236 ++++++++++
.../tests/custom/VectorUpsertTableTest.java | 10 +-
.../realtime/impl/vector/MutableVectorIndex.java | 309 +++++++++++--
.../index/readers/vector/BaseFilterQuery.java | 149 ++++++
.../readers/vector/HnswVectorIndexReader.java | 88 +---
.../impl/vector/MutableVectorIndexTest.java | 501 ++++++++++++++++++++-
.../index/readers/vector/BaseFilterQueryTest.java | 130 ++++++
11 files changed, 1530 insertions(+), 124 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSearchStrategy.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSearchStrategy.java
index cc1d3a2905c..f7e0edc03c8 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSearchStrategy.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSearchStrategy.java
@@ -71,17 +71,32 @@ public final class VectorSearchStrategy {
public static Decision decide(int numDocs, int estimatedFilteredDocs,
boolean hasVectorIndex,
boolean indexSupportsPreFilter, boolean isMutableSegment, @Nullable
VectorBackendType backendType,
@Nullable VectorSearchParams searchParams) {
+ return decide(numDocs, estimatedFilteredDocs, hasVectorIndex,
indexSupportsPreFilter, isMutableSegment,
+ backendType, searchParams, true);
+ }
+
+ /// Decides only the search mode without constructing the explanation
carried by [Decision].
+ public static VectorSearchMode decideMode(int numDocs, int
estimatedFilteredDocs, boolean hasVectorIndex,
+ boolean indexSupportsPreFilter, boolean isMutableSegment, @Nullable
VectorBackendType backendType,
+ @Nullable VectorSearchParams searchParams) {
+ return decide(numDocs, estimatedFilteredDocs, hasVectorIndex,
indexSupportsPreFilter, isMutableSegment,
+ backendType, searchParams, false).getMode();
+ }
+
+ private static Decision decide(int numDocs, int estimatedFilteredDocs,
boolean hasVectorIndex,
+ boolean indexSupportsPreFilter, boolean isMutableSegment, @Nullable
VectorBackendType backendType,
+ @Nullable VectorSearchParams searchParams, boolean includeReason) {
// No index → always exact scan
if (!hasVectorIndex) {
return new Decision(VectorSearchMode.EXACT_SCAN, -1.0,
- "no_vector_index", numDocs, estimatedFilteredDocs);
+ includeReason ? "no_vector_index" : "", numDocs,
estimatedFilteredDocs);
}
// Very small segment → exact scan is cheaper
if (numDocs < MIN_ANN_SEGMENT_SIZE) {
return new Decision(VectorSearchMode.EXACT_SCAN, -1.0,
- "segment_too_small (numDocs=" + numDocs + " < " +
MIN_ANN_SEGMENT_SIZE + ")",
+ includeReason ? "segment_too_small (numDocs=" + numDocs + " < " +
MIN_ANN_SEGMENT_SIZE + ")" : "",
numDocs, estimatedFilteredDocs);
}
@@ -89,7 +104,7 @@ public final class VectorSearchStrategy {
boolean hasFilter = estimatedFilteredDocs < numDocs;
if (!hasFilter) {
return new Decision(VectorSearchMode.POST_FILTER_ANN, 1.0,
- "no_filter", numDocs, estimatedFilteredDocs);
+ includeReason ? "no_filter" : "", numDocs, estimatedFilteredDocs);
}
double selectivity = numDocs > 0 ? (double) estimatedFilteredDocs /
numDocs : 1.0;
@@ -97,7 +112,9 @@ public final class VectorSearchStrategy {
// Very few docs pass filter → exact scan on filtered set
if (estimatedFilteredDocs < EXACT_SCAN_THRESHOLD) {
return new Decision(VectorSearchMode.EXACT_SCAN, selectivity,
- "filter_too_selective (filteredDocs=" + estimatedFilteredDocs + " <
" + EXACT_SCAN_THRESHOLD + ")",
+ includeReason
+ ? "filter_too_selective (filteredDocs=" + estimatedFilteredDocs
+ " < " + EXACT_SCAN_THRESHOLD + ")"
+ : "",
numDocs, estimatedFilteredDocs);
}
@@ -105,21 +122,26 @@ public final class VectorSearchStrategy {
if (selectivity < HIGH_SELECTIVITY_THRESHOLD) {
if (indexSupportsPreFilter && !isMutableSegment) {
return new Decision(VectorSearchMode.FILTER_THEN_ANN, selectivity,
- "high_selectivity (ratio=" + String.format("%.4f", selectivity) +
" < "
- + HIGH_SELECTIVITY_THRESHOLD + ")",
+ includeReason
+ ? "high_selectivity (ratio=" + String.format("%.4f",
selectivity) + " < "
+ + HIGH_SELECTIVITY_THRESHOLD + ")"
+ : "",
numDocs, estimatedFilteredDocs);
}
// Pre-filter not supported → fall back to exact scan on filtered set
return new Decision(VectorSearchMode.EXACT_SCAN, selectivity,
- "high_selectivity_no_prefilter_support (ratio=" +
String.format("%.4f", selectivity) + ")",
+ includeReason
+ ? "high_selectivity_no_prefilter_support (ratio=" +
String.format("%.4f", selectivity) + ")" : "",
numDocs, estimatedFilteredDocs);
}
// Low selectivity → post-filter ANN
if (selectivity > LOW_SELECTIVITY_THRESHOLD) {
return new Decision(VectorSearchMode.POST_FILTER_ANN, selectivity,
- "low_selectivity (ratio=" + String.format("%.4f", selectivity) + " >
"
- + LOW_SELECTIVITY_THRESHOLD + ")",
+ includeReason
+ ? "low_selectivity (ratio=" + String.format("%.4f", selectivity)
+ " > "
+ + LOW_SELECTIVITY_THRESHOLD + ")"
+ : "",
numDocs, estimatedFilteredDocs);
}
@@ -133,14 +155,16 @@ public final class VectorSearchStrategy {
double midpoint = (HIGH_SELECTIVITY_THRESHOLD +
LOW_SELECTIVITY_THRESHOLD) / 2;
if (selectivity < midpoint) {
return new Decision(VectorSearchMode.FILTER_THEN_ANN, selectivity,
- "cost_model_prefilter (ratio=" + String.format("%.4f",
selectivity) + " < midpoint="
- + String.format("%.4f", midpoint) + ")",
+ includeReason
+ ? "cost_model_prefilter (ratio=" + String.format("%.4f",
selectivity) + " < midpoint="
+ + String.format("%.4f", midpoint) + ")"
+ : "",
numDocs, estimatedFilteredDocs);
}
}
return new Decision(VectorSearchMode.POST_FILTER_ANN, selectivity,
- "cost_model_postfilter (ratio=" + String.format("%.4f", selectivity) +
")",
+ includeReason ? "cost_model_postfilter (ratio=" +
String.format("%.4f", selectivity) + ")" : "",
numDocs, estimatedFilteredDocs);
}
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 013774f3c03..74bdb9f2a38 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
@@ -391,6 +391,7 @@ public class FilterPlanNode implements PlanNode {
}
if (vectorIndex != null) {
+ VectorBackendType backendType =
VectorDistanceUtils.resolveBackendType(vectorIndexConfig);
// A required candidate scope can only be honored by a reader that does
filtered search. When it cannot, the
// ANN index is unusable for this query and an exact scan over the
allowed documents is the correct plan --
// decided here, where both the reader capability and the forward index
availability are known.
@@ -403,9 +404,29 @@ public class FilterPlanNode implements PlanNode {
getFilteredSearchUnsupportedReason(isMutableSegment),
searchParams, _requiredVectorDocIds);
}
+ // Mutable HNSW filtered search materializes the full NumericDocValues
view before applying the bitmap. For a
+ // selective required scope, scanning only the visible forward-index
rows is both exact and substantially
+ // cheaper. Preserve correctness for the other strategy outcomes by
continuing to pass the required scope to
+ // candidate generation -- a required scope must never be applied as a
post-filter. If the forward index is
+ // disabled, retain filtered ANN so this optimization does not make a
previously executable query fail.
+ if (_requiredVectorDocIds != null && isMutableSegment) {
+ VectorSearchMode mode = VectorSearchStrategy.decideMode(numDocs,
+ _requiredVectorDocIds.getCardinality(),
+ /* hasVectorIndex= */ true,
+ /* indexSupportsPreFilter= */ true,
+ /* isMutableSegment= */ true,
+ backendType, searchParams);
+ if (mode == VectorSearchMode.EXACT_SCAN) {
+ ForwardIndexReader<?> exactScanReader = dataSource.getForwardIndex();
+ if (exactScanReader != null) {
+ return new ExactVectorScanFilterOperator(exactScanReader,
predicate, column, numDocs, vectorIndexConfig,
+ "required_doc_ids_strategy_exact_scan", searchParams,
_requiredVectorDocIds);
+ }
+ }
+ }
+
// ANN index path: pass the forward index when rerank or threshold needs
exact distances.
ForwardIndexReader<?> forwardIndexReader = null;
- VectorBackendType backendType =
VectorDistanceUtils.resolveBackendType(vectorIndexConfig);
if (searchParams.isExactRerank(backendType) ||
searchParams.hasDistanceThreshold()) {
forwardIndexReader = dataSource.getForwardIndex();
Preconditions.checkState(!searchParams.hasDistanceThreshold() ||
forwardIndexReader != null,
@@ -575,7 +596,7 @@ public class FilterPlanNode implements PlanNode {
// stage we are deciding whether to activate pre-filtering at all, not
per-backend tuning.
// The strategy currently uses only selectivity (numDocs,
estimatedFilteredDocs) for this
// decision. Per-backend and per-query-option tuning is handled later
inside the operator.
- VectorSearchStrategy.Decision decision = VectorSearchStrategy.decide(
+ VectorSearchMode mode = VectorSearchStrategy.decideMode(
numDocs, estimatedFilteredDocs,
/* hasVectorIndex= */ true,
/* indexSupportsPreFilter= */ true,
@@ -583,7 +604,7 @@ public class FilterPlanNode implements PlanNode {
/* backendType= */ null,
/* searchParams= */ null);
- if (decision.getMode() != VectorSearchMode.FILTER_THEN_ANN) {
+ if (mode != VectorSearchMode.FILTER_THEN_ANN) {
return;
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/VectorSearchStrategyTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/VectorSearchStrategyTest.java
index 0f81a4e0c9b..06f528b7863 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/VectorSearchStrategyTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/VectorSearchStrategyTest.java
@@ -19,6 +19,7 @@
package org.apache.pinot.core.operator.filter;
import org.apache.pinot.segment.spi.index.creator.VectorBackendType;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
@@ -28,6 +29,33 @@ import static org.testng.Assert.assertTrue;
/// Tests for [VectorSearchStrategy] adaptive planner.
public class VectorSearchStrategyTest {
+ @DataProvider(name = "strategyCases")
+ public Object[][] strategyCases() {
+ return new Object[][]{
+ {100000, 100000, false, false, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {100, 100, true, true, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {100000, 100000, true, true, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {100000, 50, true, true, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {1000000, 5000, true, true, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {1000000, 5000, true, false, false, VectorBackendType.IVF_FLAT,
VectorSearchParams.DEFAULT},
+ {100000, 50000, true, true, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {1000000, 5000, true, true, true, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {100000, 5000, true, true, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT},
+ {100000, 15000, true, true, false, VectorBackendType.HNSW,
VectorSearchParams.DEFAULT}
+ };
+ }
+
+ @Test(dataProvider = "strategyCases")
+ public void testModeOnlyDecisionMatchesDetailedDecision(int numDocs, int
estimatedFilteredDocs,
+ boolean hasVectorIndex, boolean indexSupportsPreFilter, boolean
isMutableSegment,
+ VectorBackendType backendType, VectorSearchParams searchParams) {
+ assertEquals(
+ VectorSearchStrategy.decideMode(numDocs, estimatedFilteredDocs,
hasVectorIndex, indexSupportsPreFilter,
+ isMutableSegment, backendType, searchParams),
+ VectorSearchStrategy.decide(numDocs, estimatedFilteredDocs,
hasVectorIndex, indexSupportsPreFilter,
+ isMutableSegment, backendType, searchParams).getMode());
+ }
+
@Test
public void testNoVectorIndex() {
VectorSearchStrategy.Decision decision = VectorSearchStrategy.decide(
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeTest.java
index a74e72d28b5..896e9b24588 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeTest.java
@@ -35,6 +35,7 @@ import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.operator.blocks.FilterBlock;
import org.apache.pinot.core.operator.filter.BaseFilterOperator;
import org.apache.pinot.core.operator.filter.ExactVectorScanFilterOperator;
+import org.apache.pinot.core.operator.filter.VectorSimilarityFilterOperator;
import
org.apache.pinot.core.operator.filter.predicate.BaseDictIdBasedRegexpLikePredicateEvaluator;
import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator;
import org.apache.pinot.core.query.request.context.QueryContext;
@@ -390,6 +391,133 @@ public class FilterPlanNodeTest {
Assert.assertTrue(explain.contains("requiredDocIdFilterCardinality:2"),
explain);
}
+ @Test
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ public void testMutableSparseRequiredScopeSelectsExactScan() {
+ int numDocs = 100_000;
+ float[] queryVector = {1.0f, 0.0f};
+ FilterAwareVectorIndexReader vectorReader =
mock(FilterAwareVectorIndexReader.class);
+ when(vectorReader.supportsPreFilter()).thenReturn(true);
+ ForwardIndexReader forwardIndexReader = mock(ForwardIndexReader.class);
+ ForwardIndexReaderContext context = mock(ForwardIndexReaderContext.class);
+ when(forwardIndexReader.createContext()).thenReturn(context);
+ when(forwardIndexReader.getFloatMV(Mockito.anyInt(),
any())).thenAnswer(invocation -> {
+ int docId = invocation.getArgument(0);
+ if (docId == 50_003) {
+ return new float[]{1.0f, 0.0f};
+ }
+ if (docId == 99_991) {
+ return new float[]{0.9f, 0.0f};
+ }
+ return new float[]{0.0f, 1.0f};
+ });
+ DataSource dataSource = mock(DataSource.class);
+ when(dataSource.getVectorIndex()).thenReturn(vectorReader);
+ Mockito.doReturn(forwardIndexReader).when(dataSource).getForwardIndex();
+ IndexSegment segment = mock(IndexSegment.class);
+ SegmentMetadata metadata = mock(SegmentMetadata.class);
+ when(metadata.getTotalDocs()).thenReturn(numDocs);
+ when(metadata.isMutableSegment()).thenReturn(true);
+ when(segment.getSegmentMetadata()).thenReturn(metadata);
+ when(segment.getDataSource("embedding", null)).thenReturn(dataSource);
+ QueryContext queryContext = mock(QueryContext.class);
+ when(queryContext.getFilter()).thenReturn(FilterContext.forPredicate(new
VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), queryVector, 2)));
+ SegmentContext segmentContext = new SegmentContext(segment);
+ segmentContext.setDocIdsSnapshot(bitmapOf(17, 50_003, 99_991));
+
+ BaseFilterOperator operator = new FilterPlanNode(segmentContext,
queryContext).run();
+ Assert.assertEquals(getFilteredDocIds(operator), bitmapOf(50_003, 99_991));
+ verify(vectorReader, never()).getDocIds(any(float[].class),
Mockito.anyInt());
+ verify(vectorReader, never()).getDocIds(any(float[].class),
Mockito.anyInt(),
+ any(ImmutableRoaringBitmap.class));
+ verify(forwardIndexReader).getFloatMV(eq(17), any());
+ verify(forwardIndexReader).getFloatMV(eq(50_003), any());
+ verify(forwardIndexReader).getFloatMV(eq(99_991), any());
+ verify(forwardIndexReader, never()).getFloatMV(eq(0), any());
+ ExactVectorScanFilterOperator exactScan = findOperator(operator,
ExactVectorScanFilterOperator.class);
+ Assert.assertNotNull(exactScan);
+ String explain = exactScan.toExplainString();
+
Assert.assertTrue(explain.contains("fallbackReason:required_doc_ids_strategy_exact_scan"),
explain);
+ Assert.assertTrue(explain.contains("requiredDocIdFilterCardinality:3"),
explain);
+ }
+
+ @Test
+ public void
testMutableHighCardinalityRequiredScopeUsesPreFilterInsteadOfPostFilter() {
+ int numDocs = 10_000;
+ float[] queryVector = {1.0f, 0.0f};
+ FilterAwareVectorIndexReader vectorReader =
mock(FilterAwareVectorIndexReader.class);
+ when(vectorReader.supportsPreFilter()).thenReturn(true);
+ when(vectorReader.getDocIds(eq(queryVector), eq(2),
any(ImmutableRoaringBitmap.class)))
+ .thenReturn(bitmapOf(4_998, 4_999));
+ ForwardIndexReader<?> forwardIndexReader = mock(ForwardIndexReader.class);
+ DataSource dataSource = mock(DataSource.class);
+ when(dataSource.getVectorIndex()).thenReturn(vectorReader);
+ Mockito.doReturn(forwardIndexReader).when(dataSource).getForwardIndex();
+ IndexSegment segment = mock(IndexSegment.class);
+ SegmentMetadata metadata = mock(SegmentMetadata.class);
+ when(metadata.getTotalDocs()).thenReturn(numDocs);
+ when(metadata.isMutableSegment()).thenReturn(true);
+ when(segment.getSegmentMetadata()).thenReturn(metadata);
+ when(segment.getDataSource("embedding", null)).thenReturn(dataSource);
+ QueryContext queryContext = mock(QueryContext.class);
+ when(queryContext.getFilter()).thenReturn(FilterContext.forPredicate(new
VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), queryVector, 2)));
+ MutableRoaringBitmap requiredDocIds = new MutableRoaringBitmap();
+ requiredDocIds.add(0L, 5_000L);
+ SegmentContext segmentContext = new SegmentContext(segment);
+ segmentContext.setDocIdsSnapshot(requiredDocIds);
+
+ BaseFilterOperator operator = new FilterPlanNode(segmentContext,
queryContext).run();
+ Assert.assertEquals(getFilteredDocIds(operator), bitmapOf(4_998, 4_999));
+ ArgumentCaptor<ImmutableRoaringBitmap> requiredScope =
+ ArgumentCaptor.forClass(ImmutableRoaringBitmap.class);
+ verify(vectorReader).getDocIds(eq(queryVector), eq(2),
requiredScope.capture());
+ Assert.assertEquals(requiredScope.getValue().getCardinality(), 5_000);
+ Assert.assertTrue(requiredScope.getValue().contains(4_999));
+ verify(vectorReader, never()).getDocIds(queryVector, 2);
+ verifyNoInteractions(forwardIndexReader);
+ Assert.assertNull(findOperator(operator,
ExactVectorScanFilterOperator.class));
+ VectorSimilarityFilterOperator filteredAnn = findOperator(operator,
VectorSimilarityFilterOperator.class);
+ Assert.assertNotNull(filteredAnn);
+
Assert.assertTrue(filteredAnn.toExplainString().contains("searchMode:FILTER_THEN_ANN"));
+ }
+
+ @Test
+ public void
testMutableSparseRequiredScopeWithoutForwardIndexRetainsPreFilter() {
+ int numDocs = 100_000;
+ float[] queryVector = {1.0f, 0.0f};
+ FilterAwareVectorIndexReader vectorReader =
mock(FilterAwareVectorIndexReader.class);
+ when(vectorReader.supportsPreFilter()).thenReturn(true);
+ when(vectorReader.getDocIds(eq(queryVector), eq(1),
any(ImmutableRoaringBitmap.class)))
+ .thenReturn(bitmapOf(50_003));
+ DataSource dataSource = mock(DataSource.class);
+ when(dataSource.getVectorIndex()).thenReturn(vectorReader);
+ IndexSegment segment = mock(IndexSegment.class);
+ SegmentMetadata metadata = mock(SegmentMetadata.class);
+ when(metadata.getTotalDocs()).thenReturn(numDocs);
+ when(metadata.isMutableSegment()).thenReturn(true);
+ when(segment.getSegmentMetadata()).thenReturn(metadata);
+ when(segment.getDataSource("embedding", null)).thenReturn(dataSource);
+ QueryContext queryContext = mock(QueryContext.class);
+ when(queryContext.getFilter()).thenReturn(FilterContext.forPredicate(new
VectorSimilarityPredicate(
+ ExpressionContext.forIdentifier("embedding"), queryVector, 1)));
+ SegmentContext segmentContext = new SegmentContext(segment);
+ segmentContext.setDocIdsSnapshot(bitmapOf(17, 50_003, 99_991));
+
+ BaseFilterOperator operator = new FilterPlanNode(segmentContext,
queryContext).run();
+ Assert.assertEquals(getFilteredDocIds(operator), bitmapOf(50_003));
+ ArgumentCaptor<ImmutableRoaringBitmap> requiredScope =
+ ArgumentCaptor.forClass(ImmutableRoaringBitmap.class);
+ verify(vectorReader).getDocIds(eq(queryVector), eq(1),
requiredScope.capture());
+ Assert.assertEquals(requiredScope.getValue(), bitmapOf(17, 50_003,
99_991));
+ verify(vectorReader, never()).getDocIds(queryVector, 1);
+ Assert.assertNull(findOperator(operator,
ExactVectorScanFilterOperator.class));
+ VectorSimilarityFilterOperator filteredAnn = findOperator(operator,
VectorSimilarityFilterOperator.class);
+ Assert.assertNotNull(filteredAnn);
+
Assert.assertTrue(filteredAnn.toExplainString().contains("searchMode:FILTER_THEN_ANN"));
+ }
+
@Test(expectedExceptions = IllegalStateException.class,
expectedExceptionsMessageRegExp = ".*required candidate doc IDs.*no
forward index.*")
public void testPlannerFailsWhenRequiredFilterCannotBeHonored() {
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/HnswVectorRealtimeTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/HnswVectorRealtimeTest.java
new file mode 100644
index 00000000000..e2e4b6142a9
--- /dev/null
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/HnswVectorRealtimeTest.java
@@ -0,0 +1,236 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.integration.tests.custom;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.File;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.config.table.FieldConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// Realtime integration coverage for filtered HNSW search on a **consuming**
segment, where the vector index is
+/// `MutableVectorIndex` rather than the offline reader.
+///
+/// This pins a behavior change that reaches plain realtime tables, not only
upsert ones: because the consuming
+/// segment now advertises `supportsPreFilter()`, a `vectorSimilarity`
predicate combined with a metadata filter
+/// plans as `FILTER_THEN_ANN` instead of running an unfiltered ANN and
intersecting afterwards. Nothing else
+/// covers this -- the other HNSW integration tests are offline-only, and
[IvfPqVectorRealtimeTest] asserts the
+/// IVF_PQ exact-scan fallback.
+///
+/// ## Why the fixture is sized the way it is
+///
+/// `VectorSearchStrategy.decide` only wires the optional pre-filter when the
filter matches at least
+/// `EXACT_SCAN_THRESHOLD` (1000) documents *and* selectivity falls below the
mid-range cutoff of 0.105; anything
+/// more selective is cheaper as an exact scan, anything less selective is
left to post-filtering. A single
+/// consuming segment of [#getCountStarResult] rows over [#NUM_CATEGORIES]
categories yields 1250 matches at
+/// selectivity 0.083, which clears both bounds with margin. One Kafka
partition and a flush size above the row
+/// count keep every row in one consuming segment, so the per-segment counts
the planner sees are the ones
+/// computed here. The category column carries an inverted index because the
wiring additionally requires every
+/// non-vector filter to produce a bitmap -- a scan-based predicate silently
leaves the vector operator
+/// post-filtering.
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class HnswVectorRealtimeTest extends
CustomDataQueryClusterIntegrationTest {
+ private static final String DEFAULT_TABLE_NAME = "HnswVectorRealtimeTest";
+ private static final String VECTOR_COL = "embedding";
+ private static final String CATEGORY = "category";
+ private static final int VECTOR_DIM_SIZE = 32;
+ private static final int NUM_CATEGORIES = 12;
+ private static final int NUM_ROWS = 15000;
+ private static final String TARGET_CATEGORY = "cat_3";
+
+ @Override
+ protected long getCountStarResult() {
+ return NUM_ROWS;
+ }
+
+ @Override
+ public String getTableName() {
+ return DEFAULT_TABLE_NAME;
+ }
+
+ @Override
+ public boolean isRealtimeTable() {
+ return true;
+ }
+
+ /// Keep every row in a single consuming segment: the planner reasons about
per-segment counts, so a mid-stream
+ /// commit would shrink them below the pre-filter thresholds this test
depends on.
+ @Override
+ protected int getRealtimeSegmentFlushSize() {
+ return NUM_ROWS * 10;
+ }
+
+ @Override
+ protected int getNumKafkaPartitions() {
+ return 1;
+ }
+
+ @Override
+ public Schema createSchema() {
+ return new Schema.SchemaBuilder().setSchemaName(getTableName())
+ .addMultiValueDimension(VECTOR_COL, DataType.FLOAT)
+ .addSingleValueDimension(CATEGORY, DataType.STRING)
+ .addDateTimeField(getTimeColumnName(), DataType.LONG,
"1:MILLISECONDS:EPOCH", "1:MILLISECONDS")
+ .build();
+ }
+
+ /// The inherited realtime builder attaches no field configs, so the HNSW
index is declared here -- without it the
+ /// consuming segment would hold no vector index and these tests would
silently exercise a scan instead.
+ @Override
+ protected TableConfig createRealtimeTableConfig(File sampleAvroFile) {
+ AvroFileSchemaKafkaAvroMessageDecoder._avroFile = sampleAvroFile;
+ return getTableConfigBuilder(TableType.REALTIME)
+ // The pre-filter is only wired when every non-vector filter can
produce a bitmap; without an index the
+ // category predicate plans as a full scan and the vector operator
never receives the bitmap.
+ .setInvertedIndexColumns(List.of(CATEGORY))
+ .setFieldConfigList(List.of(
+ new FieldConfig.Builder(VECTOR_COL)
+ .withIndexTypes(List.of(FieldConfig.IndexType.VECTOR))
+ .withEncodingType(FieldConfig.EncodingType.RAW)
+ .withProperties(Map.of(
+ "vectorIndexType", "HNSW",
+ "vectorDimension", String.valueOf(VECTOR_DIM_SIZE),
+ "vectorDistanceFunction", "COSINE",
+ "version", "1"))
+ .build()
+ ))
+ .build();
+ }
+
+ @Override
+ public List<File> createAvroFiles()
+ throws Exception {
+ org.apache.avro.Schema avroSchema =
org.apache.avro.Schema.createRecord("myRecord", null, null, false);
+ org.apache.avro.Schema floatArraySchema =
+
org.apache.avro.Schema.createArray(org.apache.avro.Schema.create(org.apache.avro.Schema.Type.FLOAT));
+ avroSchema.setFields(List.of(
+ new org.apache.avro.Schema.Field(VECTOR_COL, floatArraySchema, null,
null),
+ new org.apache.avro.Schema.Field(CATEGORY,
+ org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING),
null, null),
+ new org.apache.avro.Schema.Field(getTimeColumnName(),
+ org.apache.avro.Schema.create(org.apache.avro.Schema.Type.LONG),
null, null)
+ ));
+
+ try (AvroFilesAndWriters avroFilesAndWriters =
createAvroFilesAndWriters(avroSchema)) {
+ List<DataFileWriter<GenericData.Record>> writers =
avroFilesAndWriters.getWriters();
+ Random random = new Random(42);
+ long baseTimestamp = System.currentTimeMillis();
+ for (int i = 0; i < NUM_ROWS; i++) {
+ GenericData.Record record = new GenericData.Record(avroSchema);
+ Collection<Float> vector = new ArrayList<>(VECTOR_DIM_SIZE);
+ for (int d = 0; d < VECTOR_DIM_SIZE; d++) {
+ vector.add(random.nextFloat());
+ }
+ record.put(VECTOR_COL, vector);
+ record.put(CATEGORY, "cat_" + (i % NUM_CATEGORIES));
+ record.put(getTimeColumnName(), baseTimestamp + i);
+ writers.get(i % getNumAvroFiles()).append(record);
+ }
+ return avroFilesAndWriters.getAvroFiles();
+ }
+ }
+
+ /// Pins that the planner actually hands the metadata bitmap to the
consuming segment's vector operator, so a
+ /// regression that stops wiring the pre-filter fails here rather than
surfacing only as a recall difference.
+ ///
+ /// The operator's `searchMode` is deliberately not asserted: it is
initialized to `POST_FILTER_ANN` and only
+ /// reassigned to `FILTER_THEN_ANN` while the search executes, and `EXPLAIN`
plans without executing, so it
+ /// always reports the initial value here. `filterSelectivity` is derived
from the bitmap the operator received,
+ /// which makes it the field that actually distinguishes a wired pre-filter
from an unwired one.
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testExplainShowsPreFilterReachesConsumingSegment(boolean
useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ String explainQuery = String.format(
+ "set explainAskingServers=true; EXPLAIN PLAN FOR "
+ + "SELECT cosineDistance(%s, %s) AS dist FROM %s "
+ + "WHERE vectorSimilarity(%s, %s, %d) AND %s = '%s' ORDER BY dist
ASC LIMIT %d",
+ VECTOR_COL, queryVector(), getTableName(), VECTOR_COL, queryVector(),
10, CATEGORY, TARGET_CATEGORY, 10);
+
+ String explain = postQuery(explainQuery).get("resultTable").toString();
+ assertTrue(explain.contains("backend"), "Explain should describe the
vector index: " + explain);
+ // NUM_ROWS / NUM_CATEGORIES matches out of NUM_ROWS -- the bitmap the
operator received, not an estimate.
+ assertTrue(explain.contains("filterSelectivity"),
+ "The metadata bitmap should be wired into the consuming segment's
vector operator: " + explain);
+ assertTrue(explain.contains("0.0833"),
+ "Pre-filter selectivity should reflect the " + (NUM_ROWS /
NUM_CATEGORIES) + " matching rows: " + explain);
+ }
+
+ /// Every row a filtered vector search returns must satisfy the predicate. A
disallowed row here would mean
+ /// candidate generation ignored the pre-filter bitmap.
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testFilteredAnnReturnsOnlyMatchingRows(boolean
useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ int topK = 10;
+ String query = String.format(
+ "SELECT cosineDistance(%s, %s) AS dist, %s FROM %s "
+ + "WHERE vectorSimilarity(%s, %s, %d) AND %s = '%s' "
+ + "ORDER BY dist ASC LIMIT %d",
+ VECTOR_COL, queryVector(), CATEGORY, getTableName(),
+ VECTOR_COL, queryVector(), topK, CATEGORY, TARGET_CATEGORY, topK);
+
+ JsonNode rows = postQuery(query).get("resultTable").get("rows");
+ assertEquals(rows.size(), topK, "Filtered ANN on a consuming segment must
return a full topK");
+ double prevDist = -1;
+ for (int i = 0; i < rows.size(); i++) {
+ assertEquals(rows.get(i).get(1).asText(), TARGET_CATEGORY, "All results
must match the filter");
+ double dist = rows.get(i).get(0).asDouble();
+ assertTrue(dist >= prevDist, "Results must be ordered by distance");
+ prevDist = dist;
+ }
+ }
+
+ /// Candidate generation constrained by the filter yields a full topK drawn
from the matching rows. Under the
+ /// previous unfiltered-ANN-then-intersect behavior only the matching share
of an unfiltered top-K survived,
+ /// which for this fixture would be roughly `topK / NUM_CATEGORIES` rows.
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testFilteredAnnDrawsTopKFromTheFilteredSet(boolean
useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ int topK = 50;
+ String filteredQuery = String.format(
+ "SELECT count(*) FROM %s WHERE vectorSimilarity(%s, %s, %d) AND %s =
'%s'",
+ getTableName(), VECTOR_COL, queryVector(), topK, CATEGORY,
TARGET_CATEGORY);
+
+ long filteredCount =
postQuery(filteredQuery).get("resultTable").get("rows").get(0).get(0).asLong();
+ assertTrue(filteredCount >= topK,
+ "Filtered ANN must draw a full topK from the matching rows, got " +
filteredCount + " for topK " + topK
+ + "; post-intersection would yield roughly " + (topK /
NUM_CATEGORIES));
+ }
+
+ private static String queryVector() {
+ return "ARRAY[0.5" + StringUtils.repeat(", 0.5", VECTOR_DIM_SIZE - 1) +
"]";
+ }
+}
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VectorUpsertTableTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VectorUpsertTableTest.java
index f7679b15431..1ab4107b123 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VectorUpsertTableTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VectorUpsertTableTest.java
@@ -267,10 +267,12 @@ public class VectorUpsertTableTest extends
CustomDataQueryClusterIntegrationTest
String explain = GroupByOptionsTest.toExplainStr(explainResponse,
useMultiStageQueryEngine);
assertExplainContains(explain, "requiredDocIdFilterApplied", true);
if (tableName.equals(CONSUMING_TABLE_NAME)) {
- // The mutable vector index cannot restrict its search, so the planner
selects the exact scan operator.
- assertTrue(explain.contains("VECTOR_SIMILARITY_EXACT_SCAN") ||
explain.contains("VectorSimilarityExactScan"),
- "Consuming segments must fall back to the exact scan operator: " +
explain);
- assertExplainContains(explain, "fallbackReason",
"mutable_vector_index_not_filter_aware");
+ // The consuming segment is intentionally small, so scanning only the
required upsert-visible rows is cheaper
+ // and exact. Larger required scopes retain filtered ANN, as covered by
FilterPlanNodeTest.
+ assertTrue(explain.contains("VECTOR_SIMILARITY_EXACT_SCAN")
+ || explain.contains("VectorSimilarityExactScan"),
+ "Consuming segments with a sparse required scope must use an exact
scan: " + explain);
+ assertExplainContains(explain, "fallbackReason",
"required_doc_ids_strategy_exact_scan");
} else {
assertTrue(explain.contains("VECTOR_SIMILARITY_INDEX") ||
explain.contains("VectorSimilarityIndex"),
"Sealed segments must use their vector index: " + explain);
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndex.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndex.java
index 27d77ad37cf..d7ecb1b5e04 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndex.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndex.java
@@ -18,27 +18,39 @@
*/
package org.apache.pinot.segment.local.realtime.impl.vector;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
import java.io.File;
import java.io.IOException;
import java.util.Arrays;
+import java.util.Comparator;
import java.util.LinkedHashMap;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import javax.annotation.Nullable;
import org.apache.commons.io.FileUtils;
import org.apache.lucene.document.Document;
-import org.apache.lucene.document.StoredField;
+import org.apache.lucene.document.NumericDocValuesField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.NumericDocValues;
+import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnFloatVectorQuery;
-import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.SearcherManager;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.IOUtils;
import
org.apache.pinot.segment.local.realtime.impl.invertedindex.RealtimeLuceneTextIndexSearcherPool;
import
org.apache.pinot.segment.local.segment.creator.impl.vector.XKnnFloatVectorField;
+import
org.apache.pinot.segment.local.segment.index.readers.vector.BaseFilterQuery;
import
org.apache.pinot.segment.local.segment.index.readers.vector.LuceneHnswRuntimeControlUtils;
import org.apache.pinot.segment.local.segment.store.VectorIndexUtils;
import org.apache.pinot.segment.spi.V1Constants;
@@ -46,18 +58,42 @@ import
org.apache.pinot.segment.spi.index.VectorIndexConfigProvider;
import org.apache.pinot.segment.spi.index.creator.VectorIndexConfig;
import org.apache.pinot.segment.spi.index.mutable.MutableIndex;
import org.apache.pinot.segment.spi.index.reader.EfSearchAware;
-import org.apache.pinot.segment.spi.index.reader.VectorIndexReader;
+import org.apache.pinot.segment.spi.index.reader.FilterAwareVectorIndexReader;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-/// A Vector index reader for the real-time Vector index values on the fly.
-/// Since there is no good mutable vector index implementation for topK
search, we just do brute force search.
+/// A mutable HNSW vector index for real-time (consuming) segments, backed by
a Lucene [IndexWriter] over a
+/// temp-dir [FSDirectory].
+///
+/// Every added document stores the supplied Pinot doc id as a
[NumericDocValuesField]. The same doc values
+/// drive both filtered traversal and the translation of search hits, so no
assumption is made that
+/// `ScoreDoc.doc == Pinot docId` (Lucene may renumber on merges).
+
+///
+/// Filtered search ([#getDocIds(float[], int, ImmutableRoaringBitmap)])
restricts HNSW candidate generation
+/// to the given Pinot doc ids (used to enforce the upsert doc-ids snapshot).
It searches a near-real-time
+/// reader obtained from the writer, so uncommitted rows are visible --
required for upsert correctness,
+/// where the newest version of a record is the most recently added and may
not be committed yet. The
+/// unfiltered path keeps searching the last committed generation (cheaper;
commit cadence is controlled by
+/// `commitIntervalMs` / `commitDocs`).
+///
+/// **Cost of that freshness.** A refresh flushes the writer's RAM buffer,
which for a vector field writes out
+/// the HNSW graph for the pending rows as a new Lucene segment. While a
segment is actively consuming, nearly
+/// every filtered query finds new rows and therefore triggers one, so segment
creation is driven by query rate
+/// rather than ingestion rate, and the resulting merges rebuild graphs.
Concurrent callers coalesce onto a
+/// single refresh rather than each forcing their own. Bounding this properly
means reopening on a background
+/// thread and having queries wait on a sequence number instead of driving the
flush themselves -- see
+/// `RealtimeLuceneIndexRefreshManager`, which does that for the text indexes,
and Lucene's
+/// `ControlledRealTimeReopenThread`. Until then, correctness is preserved at
the cost of that flush.
///
/// This class is thread-safe for single writer multiple readers.
-public class MutableVectorIndex implements VectorIndexReader, MutableIndex,
VectorIndexConfigProvider, EfSearchAware {
+public class MutableVectorIndex
+ implements FilterAwareVectorIndexReader, MutableIndex,
VectorIndexConfigProvider, EfSearchAware {
private static final Logger LOGGER =
LoggerFactory.getLogger(MutableVectorIndex.class);
+ private static final Comparator<ScoreDoc> LUCENE_DOC_ID_ORDER =
Comparator.comparingInt(scoreDoc -> scoreDoc.doc);
private static final RealtimeLuceneTextIndexSearcherPool SEARCHER_POOL =
RealtimeLuceneTextIndexSearcherPool.getInstance();
public static final String VECTOR_INDEX_DOC_ID_COLUMN_NAME = "DocID";
@@ -73,7 +109,25 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
private final File _indexDir;
private final FSDirectory _indexDirectory;
private final IndexWriter _indexWriter;
- private int _nextDocId;
+ // Near-real-time searcher over the writer, used by the filtered search path
(upsert doc-ids snapshot
+ // enforcement) so uncommitted rows are visible; refreshed on demand, reused
across queries
+ private final SearcherManager _searcherManager;
+ // Coordinates callers that need the same near-real-time generation. One
caller refreshes while the others wait
+ // for its published watermark instead of serializing on
maybeRefreshBlocking or reopening the same generation.
+ private final Object _searcherRefreshMonitor = new Object();
+ // Guarded by _searcherRefreshMonitor.
+ private boolean _searcherRefreshInProgress;
+ private volatile long _searcherRefreshCount;
+ // Number of documents added so far; used only for the commit cadence, never
as a doc id. Written by the indexing
+ // thread only; read cross-thread for debug output, where staleness is
acceptable.
+ private volatile int _numDocsAdded;
+ /// Sequence number of the newest row handed to the writer, and the sequence
number the shared searcher was
+ /// last refreshed through. Together they let a filtered search skip the
refresh when nothing has been added
+ /// since. Doc ids cannot be used for this: [MutableIndex#add] allows rows
in arbitrary doc-id order, so a
+ /// doc-id watermark would skip the refresh a lower-numbered but newer row
needs. Writer sequence numbers are
+ /// monotonic by construction. Written by the single indexing thread, read
by query threads.
+ private volatile long _lastAddedSequenceNumber = -1;
+ private volatile long _searcherRefreshedThroughSequenceNumber = -1;
private long _lastCommitTime;
private final ThreadLocal<Integer> _efSearchOverride = new ThreadLocal<>();
@@ -97,19 +151,25 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
FSDirectory indexDirectory = null;
IndexWriter indexWriter = null;
+ SearcherManager searcherManager = null;
try {
// segment generation is always in V1 and later we convert (as part of
post creation processing)
// to V3 if segmentVersion is set to V3 in SegmentGeneratorConfig.
indexDirectory = FSDirectory.open(_indexDir.toPath());
LOGGER.info("Creating mutable HNSW index for segment: {}, column: {} at
path: {} with {}", segmentName,
vectorColumn, _indexDir.getAbsolutePath(),
vectorIndexConfig.getProperties());
- indexWriter = new IndexWriter(indexDirectory,
VectorIndexUtils.getIndexWriterConfig(vectorIndexConfig));
+ // Always start empty. The directory is temp-scoped but its name is
stable across restarts, so an unclean
+ // shutdown can leave documents written by an older build. The index is
rebuilt from the stream, so appending
+ // those stale rows would be both unnecessary and incorrect.
+ indexWriter = new IndexWriter(indexDirectory,
+
VectorIndexUtils.getIndexWriterConfig(vectorIndexConfig).setOpenMode(IndexWriterConfig.OpenMode.CREATE));
indexWriter.commit();
+ searcherManager = new SearcherManager(indexWriter, false, false, null);
_lastCommitTime = System.currentTimeMillis();
} catch (Exception e) {
// IndexWriter does not close the Directory passed to it, so both need
to be closed.
try {
- IOUtils.close(indexWriter, indexDirectory);
+ IOUtils.close(searcherManager, indexWriter, indexDirectory);
} catch (Exception closeEx) {
e.addSuppressed(closeEx);
}
@@ -119,6 +179,7 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
}
_indexDirectory = indexDirectory;
_indexWriter = indexWriter;
+ _searcherManager = searcherManager;
}
@Override
@@ -136,10 +197,13 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
XKnnFloatVectorField xKnnFloatVectorField =
new XKnnFloatVectorField(_vectorColumn, floatValues,
_vectorSimilarityFunction);
docToIndex.add(xKnnFloatVectorField);
- docToIndex.add(new StoredField(VECTOR_INDEX_DOC_ID_COLUMN_NAME,
_nextDocId++));
+ // Store the SUPPLIED Pinot doc id (not an internal counter): the doc
value translates search hits
+ // back to Pinot doc ids, and the doc-values field lets filtered search
test bitmap membership per doc
+ docToIndex.add(new NumericDocValuesField(VECTOR_INDEX_DOC_ID_COLUMN_NAME,
docId));
try {
- _indexWriter.addDocument(docToIndex);
- if ((_lastCommitTime + _commitIntervalMs < System.currentTimeMillis())
|| (_nextDocId % _commitDocs == 0)) {
+ _lastAddedSequenceNumber = _indexWriter.addDocument(docToIndex);
+ _numDocsAdded++;
+ if ((_lastCommitTime + _commitIntervalMs < System.currentTimeMillis())
|| (_numDocsAdded % _commitDocs == 0)) {
_indexWriter.commit();
_lastCommitTime = System.currentTimeMillis();
LOGGER.debug("Committed index for column: {}, segment: {}",
_vectorColumn, _segmentName);
@@ -152,6 +216,22 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
@Override
public MutableRoaringBitmap getDocIds(float[] vector, int topK) {
+ return submitSearch(vector, topK, null);
+ }
+
+ @Override
+ public ImmutableRoaringBitmap getDocIds(float[] vector, int topK,
ImmutableRoaringBitmap preFilterBitmap) {
+ // A null bitmap would fall through to an unfiltered search, returning doc
ids outside any filter -- the
+ // silent degradation this reader's contract forbids, so fail loudly
instead of deep inside Lucene.
+ Preconditions.checkNotNull(preFilterBitmap, "Pre-filter bitmap must not be
null for filtered vector search");
+ if (preFilterBitmap.isEmpty()) {
+ return new MutableRoaringBitmap();
+ }
+ return submitSearch(vector, topK, preFilterBitmap);
+ }
+
+ private MutableRoaringBitmap submitSearch(float[] vector, int topK,
+ @Nullable ImmutableRoaringBitmap preFilterBitmap) {
int effectiveEfSearch = getEffectiveEfSearch();
boolean effectiveUseRelativeDistance = getEffectiveUseRelativeDistance();
boolean effectiveUseBoundedQueue = getEffectiveUseBoundedQueue();
@@ -159,9 +239,12 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
// This propagates QueryThreadContext for CPU/memory tracking without
registering the task for cancellation,
// preventing Thread.interrupt() during Lucene search which could corrupt
FSDirectory.
// See https://github.com/apache/lucene/issues/3315 and
https://github.com/apache/lucene/issues/9309
+ // The pre-filter bitmap is captured by the lambda; callers hand this
reader a bitmap that is never
+ // mutated after submission (the planner passes a query-scoped defensive
copy), so the transfer into the
+ // executor thread is safe.
Future<MutableRoaringBitmap> searchFuture =
SEARCHER_POOL.getExecutorService().submit(
() -> executeVectorSearch(vector, topK, effectiveEfSearch,
effectiveUseRelativeDistance,
- effectiveUseBoundedQueue));
+ effectiveUseBoundedQueue, preFilterBitmap));
try {
return searchFuture.get();
} catch (InterruptedException e) {
@@ -230,7 +313,7 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
info.put("effectiveEfSearch", getEffectiveEfSearch());
info.put("effectiveHnswUseRelativeDistance",
getEffectiveUseRelativeDistance());
info.put("effectiveHnswUseBoundedQueue", getEffectiveUseBoundedQueue());
- info.put("supportsPreFilter", false);
+ info.put("supportsPreFilter", supportsPreFilter());
try (DirectoryReader directoryReader =
DirectoryReader.open(_indexDirectory)) {
info.put("numDocs", directoryReader.numDocs());
info.put("numDeletedDocs", directoryReader.numDeletedDocs());
@@ -238,7 +321,7 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
} catch (IOException e) {
LOGGER.warn("Failed to load mutable HNSW debug stats for segment: {},
column: {}", _segmentName, _vectorColumn,
e);
- info.put("numDocs", _nextDocId);
+ info.put("numDocs", _numDocsAdded);
info.put("numDeletedDocs", 0);
info.put("luceneSegments", 0);
}
@@ -246,17 +329,185 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
}
private MutableRoaringBitmap executeVectorSearch(float[] vector, int topK,
int efSearch,
- boolean useRelativeDistance, boolean useBoundedQueue) throws IOException
{
+ boolean useRelativeDistance, boolean useBoundedQueue, @Nullable
ImmutableRoaringBitmap preFilterBitmap)
+ throws IOException {
+ if (preFilterBitmap != null) {
+ // Filtered search enforces the query's visible-document set, so it must
see every row that set names --
+ // including rows still in the writer's RAM buffer. Refreshing flushes
the writer, so only refresh when this
+ // query can actually see past the last refresh, and coalesce callers
waiting for the same generation. The
+ // added-doc watermark is read BEFORE refreshing so rows arriving during
the refresh are not wrongly claimed
+ // as visible.
+ long lastAdded = _lastAddedSequenceNumber;
+ if (lastAdded > _searcherRefreshedThroughSequenceNumber) {
+ refreshSearcherThrough(lastAdded);
+ }
+ IndexSearcher indexSearcher = _searcherManager.acquire();
+ try {
+ return search(indexSearcher, vector, topK, efSearch,
useRelativeDistance, useBoundedQueue,
+ new NumericDocValuesBitmapFilterQuery(preFilterBitmap));
+ } finally {
+ _searcherManager.release(indexSearcher);
+ }
+ }
+ // The unfiltered path keeps the cheaper last-committed view (bounded by
commitIntervalMs / commitDocs)
try (DirectoryReader directoryReader =
DirectoryReader.open(_indexDirectory)) {
- IndexSearcher indexSearcher = new IndexSearcher(directoryReader);
- KnnFloatVectorQuery query =
- LuceneHnswRuntimeControlUtils.createQuery(_vectorColumn, vector,
topK, efSearch,
- useRelativeDistance, useBoundedQueue, null);
- MutableRoaringBitmap docIds = new MutableRoaringBitmap();
- TopDocs search = indexSearcher.search(query, topK);
- Arrays.stream(search.scoreDocs).map(scoreDoc ->
scoreDoc.doc).forEach(docIds::add);
+ return search(new IndexSearcher(directoryReader), vector, topK,
efSearch, useRelativeDistance,
+ useBoundedQueue, null);
+ }
+ }
+
+ /// Refreshes the shared near-real-time searcher through
`targetSequenceNumber`, coalescing concurrent callers.
+ ///
+ /// The winning caller publishes the exact generation it captured before
refreshing and wakes the waiters. A
+ /// waiter whose target is newer then performs the next refresh. This is
necessary because rows added during a
+ /// refresh are not guaranteed to be visible in the reopened reader.
+ private void refreshSearcherThrough(long targetSequenceNumber)
+ throws IOException {
+ synchronized (_searcherRefreshMonitor) {
+ while (_searcherRefreshInProgress
+ && _searcherRefreshedThroughSequenceNumber < targetSequenceNumber) {
+ onSearcherRefreshWait();
+ try {
+ _searcherRefreshMonitor.wait();
+ } catch (InterruptedException e) {
+ throw new IOException("Interrupted while waiting to refresh vector
searcher for segment " + _segmentName
+ + " column " + _vectorColumn, e);
+ }
+ }
+ if (_searcherRefreshedThroughSequenceNumber >= targetSequenceNumber) {
+ return;
+ }
+ _searcherRefreshInProgress = true;
+ }
+
+ try {
+ beforeSearcherRefresh();
+ // A false return means another caller already owns SearcherManager's
refresh lock. This is benign: use the
+ // blocking form so this call still returns a searcher covering the
requested writer generation.
+ if (!_searcherManager.maybeRefresh()) {
+ _searcherManager.maybeRefreshBlocking();
+ }
+ synchronized (_searcherRefreshMonitor) {
+ _searcherRefreshedThroughSequenceNumber =
+ Math.max(_searcherRefreshedThroughSequenceNumber,
targetSequenceNumber);
+ _searcherRefreshCount++;
+ }
+ } finally {
+ synchronized (_searcherRefreshMonitor) {
+ _searcherRefreshInProgress = false;
+ _searcherRefreshMonitor.notifyAll();
+ }
+ }
+ }
+
+ /// Test hook invoked by the winning refresher after it publishes the
in-progress state.
+ @VisibleForTesting
+ void beforeSearcherRefresh()
+ throws IOException {
+ }
+
+ /// Test hook invoked when another caller is about to wait for the winning
refresher.
+ @VisibleForTesting
+ void onSearcherRefreshWait() {
+ }
+
+ @VisibleForTesting
+ long getSearcherRefreshCount() {
+ return _searcherRefreshCount;
+ }
+
+ private MutableRoaringBitmap search(IndexSearcher indexSearcher, float[]
vector, int topK, int efSearch,
+ boolean useRelativeDistance, boolean useBoundedQueue, @Nullable Query
filterQuery)
+ throws IOException {
+ KnnFloatVectorQuery query =
+ LuceneHnswRuntimeControlUtils.createQuery(_vectorColumn, vector, topK,
efSearch,
+ useRelativeDistance, useBoundedQueue, filterQuery);
+ MutableRoaringBitmap docIds = new MutableRoaringBitmap();
+ ScoreDoc[] scoreDocs = indexSearcher.search(query, topK).scoreDocs;
+ if (scoreDocs.length == 0) {
+ // Nothing matched, which is also the shape of a reader with no leaves:
an index that has not committed
+ // yet has no doc values to resolve, and asking for them would fail
rather than return no results.
return docIds;
}
+ // Translate Lucene doc ids to Pinot doc ids through the same doc values
that drive filtering. Lucene doc
+ // ids are NOT guaranteed to equal Pinot doc ids (merges can renumber), so
ScoreDoc.doc must never be used
+ // directly. Doc values are resolved per leaf rather than through a merged
view over every leaf, and hits
+ // are visited in Lucene doc id order because NumericDocValues only
advances forward.
+ Arrays.sort(scoreDocs, LUCENE_DOC_ID_ORDER);
+ List<LeafReaderContext> leaves = indexSearcher.getIndexReader().leaves();
+ int leafIndex = -1;
+ NumericDocValues pinotDocIds = null;
+ int leafDocBase = 0;
+ for (ScoreDoc scoreDoc : scoreDocs) {
+ int hitLeafIndex = ReaderUtil.subIndex(scoreDoc.doc, leaves);
+ if (hitLeafIndex != leafIndex) {
+ LeafReaderContext leaf = leaves.get(hitLeafIndex);
+ leafIndex = hitLeafIndex;
+ leafDocBase = leaf.docBase;
+ pinotDocIds =
leaf.reader().getNumericDocValues(VECTOR_INDEX_DOC_ID_COLUMN_NAME);
+ if (pinotDocIds == null) {
+ throw new IllegalStateException("Missing Pinot doc id doc values for
column: " + _vectorColumn);
+ }
+ }
+ if (!pinotDocIds.advanceExact(scoreDoc.doc - leafDocBase)) {
+ throw new IllegalStateException("Missing Pinot doc id for Lucene
document: " + scoreDoc.doc);
+ }
+ docIds.add(Math.toIntExact(pinotDocIds.longValue()));
+ }
+ return docIds;
+ }
+
+ /// A Lucene query accepting only documents whose stored Pinot doc id is
present in the given bitmap.
+ /// Membership is tested through the [NumericDocValues] written by
[#add(Object[], int[], int)], so it is
+ /// correct regardless of how Lucene numbers or renumbers its internal doc
ids.
+ private static class NumericDocValuesBitmapFilterQuery extends
BaseFilterQuery {
+
+ NumericDocValuesBitmapFilterQuery(ImmutableRoaringBitmap bitmap) {
+ super(bitmap);
+ }
+
+ @Override
+ protected DocIdSetIterator createLeafIterator(LeafReaderContext context)
+ throws IOException {
+ NumericDocValues docIdValues =
context.reader().getNumericDocValues(VECTOR_INDEX_DOC_ID_COLUMN_NAME);
+ if (docIdValues == null) {
+ // Every indexed document carries this field. Its absence indicates a
corrupt or incompatible index, not a
+ // leaf with no matches; returning null would silently discard the
entire leaf.
+ throw new IllegalStateException(
+ "Missing Pinot doc id doc values for column: " +
VECTOR_INDEX_DOC_ID_COLUMN_NAME);
+ }
+ return new DocIdSetIterator() {
+ @Override
+ public int docID() {
+ return docIdValues.docID();
+ }
+
+ @Override
+ public int nextDoc()
+ throws IOException {
+ return skipToAccepted(docIdValues.nextDoc());
+ }
+
+ @Override
+ public int advance(int target)
+ throws IOException {
+ return skipToAccepted(docIdValues.advance(target));
+ }
+
+ private int skipToAccepted(int doc)
+ throws IOException {
+ while (doc != NO_MORE_DOCS && !_docIds.contains((int)
docIdValues.longValue())) {
+ doc = docIdValues.nextDoc();
+ }
+ return doc;
+ }
+
+ @Override
+ public long cost() {
+ return _docIds.getLongCardinality();
+ }
+ };
+ }
}
private int getEffectiveEfSearch() {
@@ -279,11 +530,13 @@ public class MutableVectorIndex implements
VectorIndexReader, MutableIndex, Vect
try {
_indexWriter.commit();
// IndexWriter does not close the Directory passed to it, so both need
to be closed.
- IOUtils.close(_indexWriter, _indexDirectory);
- } catch (IOException e) {
- // Both close() implementations are idempotent, so this is a no-op for
whatever was already closed above.
- IOUtils.closeWhileHandlingException(_indexWriter, _indexDirectory);
- throw new RuntimeException(e);
+ IOUtils.close(_searcherManager, _indexWriter, _indexDirectory);
+ } catch (Exception e) {
+ // commit() can also fail unchecked (for example, after a tragic writer
event). Close every remaining resource
+ // so the SearcherManager cannot pin an NRT reader and its file handles
for the life of the process. All close()
+ // implementations are idempotent, so this is a no-op for resources that
were already closed above.
+ IOUtils.closeWhileHandlingException(_searcherManager, _indexWriter,
_indexDirectory);
+ throw e instanceof RuntimeException ? (RuntimeException) e : new
RuntimeException(e);
} finally {
deleteIndexDir();
}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/BaseFilterQuery.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/BaseFilterQuery.java
new file mode 100644
index 00000000000..144db5c174d
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/BaseFilterQuery.java
@@ -0,0 +1,149 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.segment.local.segment.index.readers.vector;
+
+import java.io.IOException;
+import javax.annotation.Nullable;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.ConstantScoreWeight;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.Weight;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+
+
+/// Base class for Lucene [Query] implementations that accept only documents
whose Pinot doc id is present
+/// in a [ImmutableRoaringBitmap]. Used to implement pre-filter ANN search by
restricting HNSW graph
+/// traversal to the filtered document set.
+///
+/// Because Lucene uses its own internal doc ids (which differ from Pinot doc
ids), subclasses supply the
+/// per-leaf iterator that maps Lucene doc ids to Pinot doc ids before testing
membership in the bitmap
+/// (via a doc-id translator, doc values, etc.). This class owns the
constant-score weight/scorer
+/// scaffolding, identity-based equality, and cache opt-out, so
filter-correctness fixes apply to every
+/// implementation at once.
+///
+/// Instances are single-use per search and must never be cached by Lucene
([Weight#isCacheable] returns
+/// false), since the accepted docs depend on the bitmap instance.
+///
+/// **Bitmap ownership.** The bitmap is retained by reference, not copied.
[ImmutableRoaringBitmap] only
+/// promises that *this* type exposes no mutators -- a caller may pass a
`MutableRoaringBitmap`, which is a
+/// subtype -- so the caller must not modify it once it has been handed over.
Mutating it during a search
+/// changes which documents are accepted midway through traversal and yields
results matching neither the old
+/// nor the new set. Callers that cannot promise that must pass a detached
copy.
+///
+/// Given that, instances are safe to share across the threads of a single
search: the bitmap is only read,
+/// and each leaf gets its own iterator.
+public abstract class BaseFilterQuery extends Query {
+ protected final ImmutableRoaringBitmap _docIds;
+
+ protected BaseFilterQuery(ImmutableRoaringBitmap docIds) {
+ _docIds = docIds;
+ }
+
+ /// Returns an iterator over the Lucene doc ids of the given leaf whose
corresponding Pinot doc ids are
+ /// in the bitmap, in increasing Lucene doc id order; or null when the leaf
cannot match any document.
+ @Nullable
+ protected abstract DocIdSetIterator createLeafIterator(LeafReaderContext
context)
+ throws IOException;
+
+ @Override
+ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode,
float boost) {
+ return new ConstantScoreWeight(this, boost) {
+ @Override
+ @Nullable
+ public Scorer scorer(LeafReaderContext context)
+ throws IOException {
+ DocIdSetIterator iterator = createLeafIterator(context);
+ if (iterator == null) {
+ return null;
+ }
+ float constScore = score();
+ return new Scorer(this) {
+ @Override
+ public DocIdSetIterator iterator() {
+ return iterator;
+ }
+
+ @Override
+ public float getMaxScore(int upTo) {
+ return constScore;
+ }
+
+ @Override
+ public float score() {
+ return constScore;
+ }
+
+ @Override
+ public int docID() {
+ return iterator.docID();
+ }
+ };
+ }
+
+ @Override
+ public boolean isCacheable(LeafReaderContext ctx) {
+ return false;
+ }
+ };
+ }
+
+ @Override
+ public String toString(String field) {
+ return getClass().getSimpleName() + "(cardinality=" +
_docIds.getCardinality() + ")";
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ // Identity semantics: the accepted docs depend on the bitmap instance,
and these queries are
+ // single-use per search (never cached), so structural equality is neither
needed nor meaningful
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ BaseFilterQuery that = (BaseFilterQuery) other;
+ return _docIds == that._docIds && equalsDocIdSource(that);
+ }
+
+ /// Whether `other` resolves Pinot doc ids the same way this query does.
Subclasses that read doc ids through
+ /// a collaborator must compare it here: two queries over one bitmap that
resolve doc ids differently accept
+ /// different documents, and Lucene requires unequal queries in that case.
+ ///
+ /// Called only from [#equals] after the concrete classes have been
confirmed identical, so `other` may be
+ /// cast to the subclass type without an `instanceof` check.
+ protected boolean equalsDocIdSource(BaseFilterQuery other) {
+ return true;
+ }
+
+ @Override
+ public int hashCode() {
+ return getClass().hashCode() * 31 + System.identityHashCode(_docIds);
+ }
+
+ @Override
+ public void visit(QueryVisitor visitor) {
+ visitor.visitLeaf(this);
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/HnswVectorIndexReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/HnswVectorIndexReader.java
index f60d6a13a4c..53f5744f9bb 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/HnswVectorIndexReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/HnswVectorIndexReader.java
@@ -30,16 +30,12 @@ import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.search.ConstantScoreWeight;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnFloatVectorQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
-import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.TopDocs;
-import org.apache.lucene.search.Weight;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import
org.apache.pinot.segment.local.segment.creator.impl.vector.HnswVectorIndexCreator;
@@ -200,7 +196,7 @@ public class HnswVectorIndexReader implements
FilterAwareVectorIndexReader, EfSe
@Override
public ImmutableRoaringBitmap getDocIds(float[] searchQuery, int topK,
ImmutableRoaringBitmap preFilterBitmap) {
try {
- Query filterQuery = new RoaringBitmapFilterQuery(preFilterBitmap,
_docIdTranslator, _indexReader.numDocs());
+ Query filterQuery = new RoaringBitmapFilterQuery(preFilterBitmap,
_docIdTranslator);
return translateTopDocs(search(searchQuery, topK, filterQuery));
} catch (RuntimeException e) {
throw e;
@@ -372,81 +368,23 @@ public class HnswVectorIndexReader implements
FilterAwareVectorIndexReader, EfSe
/// Because Lucene uses its own internal doc IDs (which differ from Pinot
doc IDs),
/// this query translates Lucene doc IDs to Pinot doc IDs using the
[DocIdTranslator]
/// before checking membership in the bitmap.
- static class RoaringBitmapFilterQuery extends Query {
- private final ImmutableRoaringBitmap _bitmap;
+ static class RoaringBitmapFilterQuery extends BaseFilterQuery {
private final DocIdTranslator _docIdTranslator;
- private final int _maxDoc;
-
- RoaringBitmapFilterQuery(ImmutableRoaringBitmap bitmap, DocIdTranslator
docIdTranslator, int maxDoc) {
- _bitmap = bitmap;
- _docIdTranslator = docIdTranslator;
- _maxDoc = maxDoc;
- }
-
- @Override
- public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode,
float boost) {
- return new ConstantScoreWeight(this, boost) {
- @Override
- public Scorer scorer(LeafReaderContext context) {
- int docBase = context.docBase;
- int maxDocInLeaf = context.reader().maxDoc();
- DocIdSetIterator iterator = new BitmapDocIdSetIterator(docBase,
maxDocInLeaf);
- float constScore = score();
- return new Scorer(this) {
- @Override
- public DocIdSetIterator iterator() {
- return iterator;
- }
-
- @Override
- public float getMaxScore(int upTo) {
- return constScore;
- }
-
- @Override
- public float score() {
- return constScore;
- }
-
- @Override
- public int docID() {
- return iterator.docID();
- }
- };
- }
-
- @Override
- public boolean isCacheable(LeafReaderContext ctx) {
- return false;
- }
- };
- }
-
- @Override
- public String toString(String field) {
- return "RoaringBitmapFilterQuery(cardinality=" +
_bitmap.getCardinality() + ")";
- }
@Override
- public boolean equals(Object other) {
- if (this == other) {
- return true;
- }
- if (!(other instanceof RoaringBitmapFilterQuery)) {
- return false;
- }
- RoaringBitmapFilterQuery that = (RoaringBitmapFilterQuery) other;
- return _bitmap == that._bitmap && _docIdTranslator ==
that._docIdTranslator;
+ protected boolean equalsDocIdSource(BaseFilterQuery other) {
+ // Same bitmap but a different translator accepts different documents.
+ return _docIdTranslator == ((RoaringBitmapFilterQuery)
other)._docIdTranslator;
}
- @Override
- public int hashCode() {
- return System.identityHashCode(_bitmap) * 31 +
System.identityHashCode(_docIdTranslator);
+ RoaringBitmapFilterQuery(ImmutableRoaringBitmap bitmap, DocIdTranslator
docIdTranslator) {
+ super(bitmap);
+ _docIdTranslator = docIdTranslator;
}
@Override
- public void visit(org.apache.lucene.search.QueryVisitor visitor) {
- visitor.visitLeaf(this);
+ protected DocIdSetIterator createLeafIterator(LeafReaderContext context) {
+ return new BitmapDocIdSetIterator(context.docBase,
context.reader().maxDoc());
}
/// Iterates over Lucene doc IDs whose corresponding Pinot doc IDs are in
the bitmap.
@@ -470,7 +408,7 @@ public class HnswVectorIndexReader implements
FilterAwareVectorIndexReader, EfSe
_doc++;
while (_doc < _maxDocInLeaf) {
int pinotDocId = _docIdTranslator.getPinotDocId(_docBase + _doc);
- if (_bitmap.contains(pinotDocId)) {
+ if (_docIds.contains(pinotDocId)) {
return _doc;
}
_doc++;
@@ -484,7 +422,7 @@ public class HnswVectorIndexReader implements
FilterAwareVectorIndexReader, EfSe
_doc = target;
while (_doc < _maxDocInLeaf) {
int pinotDocId = _docIdTranslator.getPinotDocId(_docBase + _doc);
- if (_bitmap.contains(pinotDocId)) {
+ if (_docIds.contains(pinotDocId)) {
return _doc;
}
_doc++;
@@ -495,7 +433,7 @@ public class HnswVectorIndexReader implements
FilterAwareVectorIndexReader, EfSe
@Override
public long cost() {
- return _bitmap.getLongCardinality();
+ return _docIds.getLongCardinality();
}
}
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndexTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndexTest.java
index 16b18b8509d..f077be3aff0 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndexTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/vector/MutableVectorIndexTest.java
@@ -18,10 +18,25 @@
*/
package org.apache.pinot.segment.local.realtime.impl.vector;
+import java.io.IOException;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.Phaser;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import
org.apache.pinot.segment.local.realtime.impl.invertedindex.RealtimeLuceneTextIndexSearcherPool;
import org.apache.pinot.segment.spi.index.creator.VectorIndexConfig;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.testng.Assert;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
@@ -30,10 +45,11 @@ import org.testng.annotations.Test;
public class MutableVectorIndexTest {
private static final String COLUMN_NAME = "embedding";
private static final String OTHER_COLUMN_NAME = "otherEmbedding";
+ private static final int SEARCHER_POOL_SIZE = 8;
@BeforeClass
public void setUpSearcherPool() {
- RealtimeLuceneTextIndexSearcherPool.init(1);
+ RealtimeLuceneTextIndexSearcherPool.init(SEARCHER_POOL_SIZE);
}
@Test
@@ -80,7 +96,271 @@ public class MutableVectorIndexTest {
Assert.assertEquals(debugInfo.get("effectiveEfSearch"), 6);
Assert.assertEquals(debugInfo.get("effectiveHnswUseRelativeDistance"),
Boolean.FALSE);
Assert.assertEquals(debugInfo.get("effectiveHnswUseBoundedQueue"),
Boolean.FALSE);
- Assert.assertEquals(debugInfo.get("supportsPreFilter"), Boolean.FALSE);
+ Assert.assertEquals(debugInfo.get("supportsPreFilter"), Boolean.TRUE);
+ Assert.assertTrue(index.supportsPreFilter(),
+ "The reader must advertise filtered search: that is what makes the
planner choose it over an exact scan");
+ } finally {
+ index.close();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Filtered search (upsert doc-ids snapshot enforcement)
+ // -----------------------------------------------------------------------
+
+ /// 2-D corpus with distinct distances from the query vector {1, 0}:
+ /// docs 0 and 1 are nearest (the "upsert-obsoleted" rows), docs 2 and 3 are
the valid rows.
+ private static MutableVectorIndex create2DIndex(long commitDocs, int
docIdOffset) {
+ Map<String, String> properties = new HashMap<>();
+ properties.put("commitDocs", String.valueOf(commitDocs));
+ properties.put("vectorIndexType", "HNSW");
+ properties.put("vectorDimension", "2");
+ VectorIndexConfig config = new VectorIndexConfig(false, "HNSW", 2, 1,
+ VectorIndexConfig.VectorDistanceFunction.EUCLIDEAN, properties);
+ MutableVectorIndex index =
+ new MutableVectorIndex("mutableVectorIndexFilterTest_" +
System.nanoTime(), COLUMN_NAME, config);
+ addVector(index, new float[]{1.0F, 0.0F}, docIdOffset);
+ addVector(index, new float[]{0.99F, 0.01F}, docIdOffset + 1);
+ addVector(index, new float[]{0.0F, 1.0F}, docIdOffset + 2);
+ addVector(index, new float[]{0.0F, -1.0F}, docIdOffset + 3);
+ return index;
+ }
+
+ @Test
+ public void testFilteredSearchExcludesNearestDisallowedDocs() {
+ // commitDocs=4 commits on the 4th add, so the unfiltered committed-view
sanity check below sees all rows
+ MutableVectorIndex index = create2DIndex(4, 0);
+ try {
+ // Sanity: unfiltered top-2 returns the physically nearest ("obsolete")
docs 0 and 1
+ ImmutableRoaringBitmap unfiltered = index.getDocIds(new float[]{1.0F,
0.0F}, 2);
+ Assert.assertEquals(unfiltered, ImmutableRoaringBitmap.bitmapOf(0, 1));
+
+ // Filtered top-2 restricted to docs 2 and 3 must return exactly those
docs. A post-intersection
+ // implementation would return empty here (the unfiltered top-2 has no
overlap with the allowed set),
+ // so this assertion genuinely discriminates filtered candidate
generation.
+ ImmutableRoaringBitmap filtered =
+ index.getDocIds(new float[]{1.0F, 0.0F}, 2,
ImmutableRoaringBitmap.bitmapOf(2, 3));
+ Assert.assertEquals(filtered, ImmutableRoaringBitmap.bitmapOf(2, 3),
+ "Filtered search must return the allowed docs, not the nearest
disallowed ones");
+ } finally {
+ index.close();
+ }
+ }
+
+ @Test(timeOut = 60_000)
+ public void testConcurrentFilteredSearchWithLiveWriter()
+ throws Exception {
+ // Single writer, concurrent reader: each phase publishes one new allowed
row, then makes the reader search for
+ // it while the writer continues adding disallowed rows. This proves both
NRT visibility and filtered-search
+ // correctness while ingestion and commits are active.
+ Map<String, String> properties = new HashMap<>();
+ properties.put("commitDocs", "7");
+ properties.put("vectorIndexType", "HNSW");
+ properties.put("vectorDimension", "2");
+ VectorIndexConfig config = new VectorIndexConfig(false, "HNSW", 2, 1,
+ VectorIndexConfig.VectorDistanceFunction.EUCLIDEAN, properties);
+ MutableVectorIndex index =
+ new MutableVectorIndex("mutableVectorIndexConcurrentTest_" +
System.nanoTime(), COLUMN_NAME, config);
+ int numPhases = 8;
+ int docsPerPhase = 25;
+ Phaser phaser = new Phaser(2);
+ CountDownLatch[] queryEnteredFilter = new CountDownLatch[numPhases];
+ CountDownLatch[] concurrentWritesFinished = new CountDownLatch[numPhases];
+ for (int phase = 0; phase < numPhases; phase++) {
+ queryEnteredFilter[phase] = new CountDownLatch(1);
+ concurrentWritesFinished[phase] = new CountDownLatch(1);
+ }
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+ Thread writer = null;
+ try {
+ addVector(index, new float[]{0.0F, 1.0F}, 0);
+ addVector(index, new float[]{0.0F, -1.0F}, 1);
+ addVector(index, new float[]{1.0F, 0.0F}, 2);
+ addVector(index, new float[]{0.99F, 0.01F}, 3);
+
+ writer = new Thread(() -> {
+ try {
+ for (int phase = 0; phase < numPhases; phase++) {
+ int firstDocId = 4 + phase * docsPerPhase;
+ // Publish the row this phase's filtered search must observe
before releasing the reader.
+ addVector(index, new float[]{1.0F, 0.0F}, firstDocId);
+ if (!advancePhase(phaser)) {
+ return;
+ }
+ // Wait until the searcher is evaluating the filter, then write
while that query is in progress.
+ // The test bitmap blocks the search at its first membership check
until these writes finish.
+ if (!queryEnteredFilter[phase].await(10, TimeUnit.SECONDS)) {
+ throw new TimeoutException("Filtered query did not enter its
bitmap in phase " + phase);
+ }
+ for (int docId = firstDocId + 1; docId < firstDocId +
docsPerPhase; docId++) {
+ addVector(index, new float[]{-1.0F, 0.0F}, docId);
+ }
+ concurrentWritesFinished[phase].countDown();
+ if (!advancePhase(phaser)) {
+ return;
+ }
+ }
+ } catch (Throwable t) {
+ failure.compareAndSet(null, t);
+ for (CountDownLatch writesFinished : concurrentWritesFinished) {
+ writesFinished.countDown();
+ }
+ phaser.forceTermination();
+ }
+ }, "mutable-vector-index-test-writer");
+ writer.start();
+
+ for (int phase = 0; phase < numPhases; phase++) {
+ Assert.assertTrue(advancePhase(phaser), "Writer terminated before
publishing phase " + phase);
+ Assert.assertNull(failure.get(), "Writer thread failed: " +
failure.get());
+ int allowedDocId = 4 + phase * docsPerPhase;
+ ImmutableRoaringBitmap allowed = new
ConcurrentWriteCoordinatingBitmap(allowedDocId,
+ queryEnteredFilter[phase], concurrentWritesFinished[phase]);
+ ImmutableRoaringBitmap filtered = index.getDocIds(new float[]{1.0F,
0.0F}, 1, allowed);
+ Assert.assertEquals(filtered, allowed,
+ "Filtered search must see the newly added allowed row while the
writer remains active");
+ Assert.assertTrue(advancePhase(phaser), "Writer terminated while
completing phase " + phase);
+ }
+
+ writer.join();
+ Assert.assertNull(failure.get(), "Writer thread failed: " +
failure.get());
+ } finally {
+ phaser.forceTermination();
+ if (writer != null) {
+ writer.interrupt();
+ writer.join(TimeUnit.SECONDS.toMillis(10));
+ Assert.assertFalse(writer.isAlive(), "Writer thread did not terminate
during test cleanup");
+ }
+ index.close();
+ }
+ }
+
+ @Test
+ public void testFilteredSearchSeesUncommittedDocs() {
+ // commitDocs is far larger than the number of added docs and
commitIntervalMs defaults to 10s, so no
+ // commit has happened: NRT visibility is what makes the filtered search
see the rows at all
+ MutableVectorIndex index = create2DIndex(1_000_000, 0);
+ try {
+ ImmutableRoaringBitmap filtered =
+ index.getDocIds(new float[]{1.0F, 0.0F}, 4,
ImmutableRoaringBitmap.bitmapOf(0, 1, 2, 3));
+ Assert.assertEquals(filtered.toArray(), new int[]{0, 1, 2, 3},
+ "Filtered search must see uncommitted rows through the NRT reader,
translated back to their Pinot "
+ + "doc ids");
+ } finally {
+ index.close();
+ }
+ }
+
+ @Test
+ public void testFilteredSearchTranslatesSuppliedPinotDocIds() {
+ // Pinot doc ids offset by 100: results must be the SUPPLIED doc ids,
proving no reliance on
+ // ScoreDoc.doc == Pinot docId
+ MutableVectorIndex index = create2DIndex(1000, 100);
+ try {
+ ImmutableRoaringBitmap filtered =
+ index.getDocIds(new float[]{1.0F, 0.0F}, 2,
ImmutableRoaringBitmap.bitmapOf(102, 103));
+ Assert.assertEquals(filtered, ImmutableRoaringBitmap.bitmapOf(102, 103));
+ } finally {
+ index.close();
+ }
+ }
+
+ @Test(timeOut = 60_000)
+ public void testFilteredSearchWithEmptyBitmapReturnsEmpty()
+ throws Exception {
+ MutableVectorIndex index = create2DIndex(1000, 0);
+ RealtimeLuceneTextIndexSearcherPool searcherPool =
RealtimeLuceneTextIndexSearcherPool.getInstance();
+ searcherPool.resize(1);
+ ExecutorService searcherExecutor = searcherPool.getExecutorService();
+ CountDownLatch blockerStarted = new CountDownLatch(1);
+ CountDownLatch releaseBlocker = new CountDownLatch(1);
+ ExecutorService caller = Executors.newSingleThreadExecutor();
+ Future<?> blocker = searcherExecutor.submit(() -> {
+ blockerStarted.countDown();
+ if (!releaseBlocker.await(30, TimeUnit.SECONDS)) {
+ throw new TimeoutException("Timed out waiting to release blocked
search worker");
+ }
+ return null;
+ });
+ try {
+ Assert.assertTrue(blockerStarted.await(10, TimeUnit.SECONDS), "Failed to
occupy the search worker");
+
+ Future<ImmutableRoaringBitmap> emptySearch = caller.submit(
+ () -> index.getDocIds(new float[]{1.0F, 0.0F}, 2,
ImmutableRoaringBitmap.bitmapOf()));
+ ImmutableRoaringBitmap filtered = emptySearch.get(10, TimeUnit.SECONDS);
+ Assert.assertEquals(filtered.getCardinality(), 0);
+ Assert.assertEquals(index.getSearcherRefreshCount(), 0L,
+ "An empty filter must return before search submission or NRT
refresh");
+ } finally {
+ releaseBlocker.countDown();
+ try {
+ blocker.get(10, TimeUnit.SECONDS);
+ } catch (Exception e) {
+ Assert.fail("Failed to clean up blocked search worker", e);
+ } finally {
+ searcherPool.resize(SEARCHER_POOL_SIZE);
+ }
+ caller.shutdownNow();
+ try {
+ Assert.assertTrue(caller.awaitTermination(10, TimeUnit.SECONDS),
"Empty-filter caller did not terminate");
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ Assert.fail("Interrupted while cleaning up empty-filter caller", e);
+ }
+ index.close();
+ }
+ }
+
+ @Test(timeOut = 60_000)
+ public void testConcurrentFilteredSearchesCoalesceRefresh()
+ throws Exception {
+ int numCallers = SEARCHER_POOL_SIZE;
+ ExecutorService callers = Executors.newFixedThreadPool(numCallers);
+ CoordinatedRefreshMutableVectorIndex index =
createCoordinatedIndexWithoutCommits(1);
+ try (index) {
+ float[] query = {1.0F, 0.0F, 0.0F, 0.0F, 0.0F};
+ Assert.assertEquals(index.getDocIds(query, 1, bitmapOf(0)).toArray(),
new int[]{0});
+ long initialRefreshCount = index.getSearcherRefreshCount();
+ addVector(index, query, 10);
+ index.coordinateNextRefresh();
+
+ CyclicBarrier startTogether = new CyclicBarrier(numCallers);
+ List<Future<ImmutableRoaringBitmap>> searches = new
ArrayList<>(numCallers);
+ for (int i = 0; i < numCallers; i++) {
+ searches.add(callers.submit(() -> {
+ startTogether.await(10, TimeUnit.SECONDS);
+ return index.getDocIds(query, 1, bitmapOf(10));
+ }));
+ }
+ index.awaitWinningRefresher();
+ index.awaitWaitingCallers();
+ index.releaseWinningRefresher();
+ for (Future<ImmutableRoaringBitmap> search : searches) {
+ Assert.assertEquals(search.get(10, TimeUnit.SECONDS).toArray(), new
int[]{10});
+ }
+ Assert.assertEquals(index.getSearcherRefreshCount(), initialRefreshCount
+ 1,
+ "Concurrent readers targeting one writer generation must share
exactly one refresh");
+ // The refresh count alone does not prove coalescing: a caller arriving
after publication would skip refresh
+ // and still leave the count at one. Require actual participation in the
production waiter path.
+ Assert.assertTrue(index.getObservedRefreshWaiters() > 0,
+ "Expected concurrent callers to coalesce onto the in-flight refresh,
but none entered the waiter path");
+ } finally {
+ index.releaseWinningRefresher();
+ callers.shutdownNow();
+ Assert.assertTrue(callers.awaitTermination(10, TimeUnit.SECONDS),
"Concurrent callers did not terminate");
+ }
+ }
+
+ @Test
+ public void testUnfilteredSearchTranslatesSuppliedPinotDocIds() {
+ // commitDocs=4 commits on the 4th add, so the committed-view unfiltered
path sees all rows; with doc ids
+ // offset by 100 the results must be the supplied ids
+ MutableVectorIndex index = create2DIndex(4, 100);
+ try {
+ int[] matches = index.getDocIds(new float[]{1.0F, 0.0F}, 2).toArray();
+ Assert.assertEquals(matches.length, 2);
+ Assert.assertEquals(matches[0], 100);
+ Assert.assertEquals(matches[1], 101);
} finally {
index.close();
}
@@ -100,6 +380,223 @@ public class MutableVectorIndexTest {
}
}
+ /// A consuming segment spends its first seconds with nothing committed: the
constructor commits an empty
+ /// index and the next commit only fires on commitDocs/commitIntervalMs. The
unfiltered path reads that
+ /// committed view, which has no leaves and therefore no doc values, so it
must return no results rather than
+ /// fail while trying to translate them. The filtered path reads a
near-real-time view and does see the rows,
+ /// which is the difference that makes it usable for enforcing the query's
visible-document set.
+ @Test
+ public void testSearchOnIndexWithNothingCommitted() {
+ try (MutableVectorIndex index = createIndexWithoutCommits()) {
+ float[] query = {1.0F, 0.0F, 0.0F, 0.0F, 0.0F};
+ Assert.assertEquals(index.getDocIds(query, 2).toArray(), new int[0],
+ "The committed view has no leaves yet, so an unfiltered search finds
nothing");
+ Assert.assertEquals(index.getDocIds(query, 2, bitmapOf(0, 1)).toArray(),
new int[]{0, 1},
+ "The near-real-time view must see rows that have not been committed
yet");
+ }
+ }
+
+ @Test
+ public void testSearchOnEmptyIndexReturnsEmpty() {
+ try (MutableVectorIndex index =
+ new MutableVectorIndex("mutableVectorIndexTest_" + System.nanoTime(),
COLUMN_NAME, createConfig())) {
+ float[] query = {1.0F, 0.0F, 0.0F, 0.0F, 0.0F};
+ Assert.assertEquals(index.getDocIds(query, 2).toArray(), new int[0]);
+ Assert.assertEquals(index.getDocIds(query, 2, bitmapOf(0, 1)).toArray(),
new int[0]);
+ }
+ }
+
+ /// Lucene short-circuits to an exact scan when the accepted-doc count is at
most topK, so a filtered test
+ /// with |allowed| <= topK never walks the HNSW graph. This one keeps the
allowed set larger than topK so the
+ /// graph traversal itself is exercised against acceptDocs.
+ @Test
+ public void testFilteredApproximateSearchWithMoreAllowedDocsThanTopK() {
+ try (MutableVectorIndex index = createIndex()) {
+ for (int docId = 4; docId < 40; docId++) {
+ addVector(index, new float[]{1.0F, docId, 0.0F, 0.0F, 0.0F}, docId);
+ }
+ MutableRoaringBitmap allowed = new MutableRoaringBitmap();
+ for (int docId = 10; docId < 40; docId++) {
+ allowed.add(docId);
+ }
+
+ int[] matches = index.getDocIds(new float[]{1.0F, 10.0F, 0.0F, 0.0F,
0.0F}, 3, allowed).toArray();
+ Assert.assertEquals(matches.length, 3);
+ for (int match : matches) {
+ Assert.assertTrue(allowed.contains(match), "Filtered search returned a
disallowed doc: " + match);
+ }
+ }
+ }
+
+ private static MutableRoaringBitmap bitmapOf(int... docIds) {
+ MutableRoaringBitmap bitmap = new MutableRoaringBitmap();
+ bitmap.add(docIds);
+ return bitmap;
+ }
+
+ private static boolean advancePhase(Phaser phaser)
+ throws InterruptedException, TimeoutException {
+ int phase = phaser.arrive();
+ return phaser.awaitAdvanceInterruptibly(phase, 10, TimeUnit.SECONDS) >= 0;
+ }
+
+ /// Holds the winning refresher after it publishes
`_searcherRefreshInProgress`, so at least one concurrent search
+ /// is proven to enter the production waiter path before the refresh is
allowed to finish.
+ private static class CoordinatedRefreshMutableVectorIndex extends
MutableVectorIndex {
+ private final CountDownLatch _winningRefresherEntered = new
CountDownLatch(1);
+ private final CountDownLatch _releaseWinningRefresher = new
CountDownLatch(1);
+ private final AtomicInteger _observedRefreshWaiters = new AtomicInteger();
+ private final CountDownLatch _waitingCallers;
+ private volatile boolean _coordinateRefresh;
+
+ CoordinatedRefreshMutableVectorIndex(String segmentName, VectorIndexConfig
config, int expectedWaitingCallers) {
+ super(segmentName, COLUMN_NAME, config);
+ _waitingCallers = new CountDownLatch(expectedWaitingCallers);
+ }
+
+ void coordinateNextRefresh() {
+ _coordinateRefresh = true;
+ }
+
+ void awaitWinningRefresher()
+ throws InterruptedException, TimeoutException {
+ if (!_winningRefresherEntered.await(10, TimeUnit.SECONDS)) {
+ throw new TimeoutException("No filtered-search caller became the
winning refresher");
+ }
+ }
+
+ void awaitWaitingCallers()
+ throws InterruptedException, TimeoutException {
+ if (!_waitingCallers.await(10, TimeUnit.SECONDS)) {
+ throw new TimeoutException("No concurrent caller entered the refresh
waiter path");
+ }
+ }
+
+ void releaseWinningRefresher() {
+ _releaseWinningRefresher.countDown();
+ }
+
+ @Override
+ void beforeSearcherRefresh()
+ throws IOException {
+ if (!_coordinateRefresh) {
+ return;
+ }
+ _winningRefresherEntered.countDown();
+ try {
+ if (!_releaseWinningRefresher.await(10, TimeUnit.SECONDS)) {
+ throw new IOException("Timed out waiting to release the winning
refresher");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException("Interrupted while holding the winning
refresher", e);
+ } finally {
+ _coordinateRefresh = false;
+ }
+ }
+
+ @Override
+ void onSearcherRefreshWait() {
+ _observedRefreshWaiters.incrementAndGet();
+ if (_coordinateRefresh) {
+ _waitingCallers.countDown();
+ }
+ }
+
+ int getObservedRefreshWaiters() {
+ return _observedRefreshWaiters.get();
+ }
+ }
+
+ /// Blocks the first filter membership check until the writer has added its
concurrent rows. This puts the
+ /// synchronization point inside Lucene's actual filtered search, rather
than merely racing two caller threads.
+ private static class ConcurrentWriteCoordinatingBitmap extends
MutableRoaringBitmap {
+ private final CountDownLatch _queryEnteredFilter;
+ private final CountDownLatch _concurrentWritesFinished;
+
+ ConcurrentWriteCoordinatingBitmap(int allowedDocId, CountDownLatch
queryEnteredFilter,
+ CountDownLatch concurrentWritesFinished) {
+ add(allowedDocId);
+ _queryEnteredFilter = queryEnteredFilter;
+ _concurrentWritesFinished = concurrentWritesFinished;
+ }
+
+ @Override
+ public boolean contains(int docId) {
+ _queryEnteredFilter.countDown();
+ try {
+ if (!_concurrentWritesFinished.await(10, TimeUnit.SECONDS)) {
+ throw new AssertionError("Concurrent writer did not finish while the
filtered query was active");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError("Interrupted while coordinating the filtered
query and writer", e);
+ }
+ return super.contains(docId);
+ }
+ }
+
+ /// Config whose commit thresholds are high enough that no commit fires
during the test.
+ private static MutableVectorIndex createIndexWithoutCommits() {
+ Map<String, String> properties = new HashMap<>();
+ properties.put("commitDocs", String.valueOf(Integer.MAX_VALUE));
+ properties.put("commitIntervalMs",
String.valueOf(TimeUnit.DAYS.toMillis(1)));
+ properties.put("vectorIndexType", "HNSW");
+ properties.put("vectorDimension", "5");
+ MutableVectorIndex index = new
MutableVectorIndex("mutableVectorIndexTest_" + System.nanoTime(), COLUMN_NAME,
+ new VectorIndexConfig(false, "HNSW", 5, 1,
VectorIndexConfig.VectorDistanceFunction.EUCLIDEAN, properties));
+ addVector(index, new float[]{1.0F, 0.0F, 0.0F, 0.0F, 0.0F}, 0);
+ addVector(index, new float[]{0.0F, 1.0F, 0.0F, 0.0F, 0.0F}, 1);
+ return index;
+ }
+
+ private static CoordinatedRefreshMutableVectorIndex
createCoordinatedIndexWithoutCommits(int expectedWaiters) {
+ Map<String, String> properties = new HashMap<>();
+ properties.put("commitDocs", String.valueOf(Integer.MAX_VALUE));
+ properties.put("commitIntervalMs",
String.valueOf(TimeUnit.DAYS.toMillis(1)));
+ properties.put("vectorIndexType", "HNSW");
+ properties.put("vectorDimension", "5");
+ VectorIndexConfig config = new VectorIndexConfig(false, "HNSW", 5, 1,
+ VectorIndexConfig.VectorDistanceFunction.EUCLIDEAN, properties);
+ CoordinatedRefreshMutableVectorIndex index = new
CoordinatedRefreshMutableVectorIndex(
+ "mutableVectorIndexCoalescingTest_" + System.nanoTime(), config,
expectedWaiters);
+ addVector(index, new float[]{1.0F, 0.0F, 0.0F, 0.0F, 0.0F}, 0);
+ addVector(index, new float[]{0.0F, 1.0F, 0.0F, 0.0F, 0.0F}, 1);
+ return index;
+ }
+
+ /// MutableIndex#add allows rows in arbitrary doc-id order, so freshness
cannot be tracked by a doc-id
+ /// watermark: a row added later with a lower doc id would be judged already
visible and the refresh skipped,
+ /// silently omitting it. Adds a high doc id, searches (refreshing through
it), then adds a lower one.
+ @Test
+ public void testFilteredSearchSeesOutOfOrderUncommittedDoc() {
+ try (MutableVectorIndex index = createIndexWithoutCommits()) {
+ float[] query = {1.0F, 0.0F, 0.0F, 0.0F, 0.0F};
+ addVector(index, new float[]{1.0F, 0.0F, 0.0F, 0.0F, 0.0F}, 10);
+ Assert.assertEquals(index.getDocIds(query, 1, bitmapOf(10)).toArray(),
new int[]{10},
+ "The searcher must be refreshed through the row just added");
+
+ addVector(index, new float[]{1.0F, 0.0F, 0.0F, 0.0F, 0.0F}, 5);
+ Assert.assertEquals(index.getDocIds(query, 1, bitmapOf(5)).toArray(),
new int[]{5},
+ "A row added out of doc-id order must still be visible to a later
filtered search");
+ }
+ }
+
+ /// Hit translation walks doc values forward, so hits must be visited in
Lucene doc-id order. Querying nearest
+ /// to the last-added row makes score order the reverse of doc-id order,
which exercises that sort.
+ @Test
+ public void testFilteredSearchWithHitsInDescendingScoreOrder() {
+ try (MutableVectorIndex index = createIndexWithoutCommits()) {
+ addVector(index, new float[]{0.0F, 0.0F, 0.0F, 0.0F, 1.0F}, 20);
+ addVector(index, new float[]{0.0F, 0.0F, 0.0F, 1.0F, 0.0F}, 21);
+ addVector(index, new float[]{0.0F, 0.0F, 1.0F, 0.0F, 0.0F}, 22);
+ int[] matches = index.getDocIds(new float[]{0.0F, 0.0F, 1.0F, 0.0F,
0.0F}, 3,
+ bitmapOf(20, 21, 22)).toArray();
+ Assert.assertEquals(matches, new int[]{20, 21, 22},
+ "Every allowed doc must translate to its supplied Pinot doc id
regardless of score order");
+ }
+ }
+
private static MutableVectorIndex createIndex() {
return createIndex("mutableVectorIndexTest_" + System.nanoTime(),
COLUMN_NAME);
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/vector/BaseFilterQueryTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/vector/BaseFilterQueryTest.java
new file mode 100644
index 00000000000..b535bdcdf46
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/vector/BaseFilterQueryTest.java
@@ -0,0 +1,130 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.segment.local.segment.index.readers.vector;
+
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.Weight;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
+
+
+/// Tests the shared Lucene constant-score scaffolding and identity contract
for Pinot doc-id bitmap filters.
+public class BaseFilterQueryTest {
+
+ @Test
+ public void testConstantScoreScaffoldingAndCacheOptOut()
+ throws Exception {
+ TestQuery query = new TestQuery(MutableRoaringBitmap.bitmapOf(2, 4), new
Object(),
+ DocIdSetIterator.range(1, 3));
+ Weight weight = query.createWeight(null, ScoreMode.COMPLETE, 2.5F);
+
+ assertFalse(weight.isCacheable(null));
+ Scorer scorer = weight.scorer(null);
+ assertNotNull(scorer);
+ assertSame(scorer.iterator(), query._iterator);
+ assertEquals(scorer.docID(), -1);
+ assertEquals(scorer.iterator().nextDoc(), 1);
+ assertEquals(scorer.docID(), 1);
+ assertEquals(scorer.score(), 2.5F);
+ assertEquals(scorer.getMaxScore(DocIdSetIterator.NO_MORE_DOCS), 2.5F);
+ }
+
+ @Test
+ public void testLeafWithoutMatchesReturnsNoScorer()
+ throws Exception {
+ TestQuery query = new TestQuery(MutableRoaringBitmap.bitmapOf(1), new
Object(), null);
+ assertNull(query.createWeight(null, ScoreMode.COMPLETE_NO_SCORES,
1.0F).scorer(null));
+ }
+
+ @Test
+ public void testIdentityEqualityIncludesBitmapClassAndDocIdSource() {
+ MutableRoaringBitmap docIds = MutableRoaringBitmap.bitmapOf(1, 3);
+ Object source = new Object();
+ TestQuery query = new TestQuery(docIds, source, DocIdSetIterator.empty());
+ TestQuery equivalent = new TestQuery(docIds, source,
DocIdSetIterator.empty());
+
+ assertEquals(query, query);
+ assertEquals(query, equivalent);
+ assertEquals(query.hashCode(), equivalent.hashCode());
+ assertNotEquals(query, new TestQuery(docIds, new Object(),
DocIdSetIterator.empty()));
+ assertNotEquals(query, new TestQuery(docIds.clone(), source,
DocIdSetIterator.empty()));
+ assertNotEquals(query, new OtherQuery(docIds));
+ assertNotEquals(query, null);
+ assertEquals(query.toString("ignored"), "TestQuery(cardinality=2)");
+ }
+
+ @Test
+ public void testVisitDelegatesToLeafVisitor() {
+ TestQuery query = new TestQuery(MutableRoaringBitmap.bitmapOf(1), new
Object(), DocIdSetIterator.empty());
+ AtomicReference<Query> visited = new AtomicReference<>();
+ query.visit(new QueryVisitor() {
+ @Override
+ public void visitLeaf(Query leafQuery) {
+ visited.set(leafQuery);
+ }
+ });
+ assertSame(visited.get(), query);
+ }
+
+ private static class TestQuery extends BaseFilterQuery {
+ private final Object _source;
+ private final DocIdSetIterator _iterator;
+
+ TestQuery(ImmutableRoaringBitmap docIds, Object source, DocIdSetIterator
iterator) {
+ super(docIds);
+ _source = source;
+ _iterator = iterator;
+ }
+
+ @Override
+ protected DocIdSetIterator createLeafIterator(LeafReaderContext context) {
+ return _iterator;
+ }
+
+ @Override
+ protected boolean equalsDocIdSource(BaseFilterQuery other) {
+ return _source == ((TestQuery) other)._source;
+ }
+ }
+
+ private static class OtherQuery extends BaseFilterQuery {
+ OtherQuery(ImmutableRoaringBitmap docIds) {
+ super(docIds);
+ }
+
+ @Override
+ protected DocIdSetIterator createLeafIterator(LeafReaderContext context) {
+ return DocIdSetIterator.empty();
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
