xiangfu0 commented on code in PR #18114:
URL: https://github.com/apache/pinot/pull/18114#discussion_r3043562899
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSimilarityFilterOperator.java:
##########
@@ -120,9 +143,25 @@ public VectorSimilarityFilterOperator(VectorIndexReader
vectorIndexReader, Vecto
_distanceFunction =
VectorDistanceUtils.resolveDistanceFunction(vectorIndexConfig);
_requestedExactRerank = searchParams.isExactRerank(_backendType);
_effectiveExactRerank = _requestedExactRerank && forwardIndexReader !=
null;
- _effectiveSearchCount = _effectiveExactRerank
- ? searchParams.getEffectiveMaxCandidates(predicate.getTopK(), numDocs)
- : predicate.getTopK();
+ _hasMetadataFilter = hasMetadataFilter;
+ _hasThresholdPredicate = searchParams.hasDistanceThreshold();
+ _distanceThreshold = searchParams.getDistanceThreshold();
+ // When metadata filter is present, over-fetch ANN candidates to
compensate for filter loss.
+ // Default over-fetch factor is 2x topK for filtered queries without
explicit maxCandidates.
+ // For threshold queries, use a larger candidate pool since we need
distance refinement.
+ int baseSearchCount;
+ if (_hasThresholdPredicate && !searchParams.isMaxCandidatesExplicit()) {
+ // Threshold queries need a larger candidate pool for exact distance
refinement.
+ // Use topK * 10 by default, similar to rerank mode.
+ baseSearchCount = Math.min(predicate.getTopK() * 10, numDocs);
+ } else if (_effectiveExactRerank) {
+ baseSearchCount =
searchParams.getEffectiveMaxCandidates(predicate.getTopK(), numDocs);
+ } else if (hasMetadataFilter && !searchParams.isMaxCandidatesExplicit()) {
+ baseSearchCount = Math.min(predicate.getTopK() *
FILTERED_OVERFETCH_FACTOR, numDocs);
+ } else {
+ baseSearchCount = predicate.getTopK();
+ }
Review Comment:
Fixed. Threshold mode now uses `getEffectiveMaxCandidates(topK, numDocs)`
which honors explicit `vectorMaxCandidates` and defaults to `topK * 10`.
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorQueryExecutionContext.java:
##########
@@ -0,0 +1,234 @@
+/**
+ * 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.core.operator.filter;
+
+import javax.annotation.Nullable;
+import org.apache.pinot.segment.spi.index.creator.VectorBackendCapabilities;
+import org.apache.pinot.segment.spi.index.creator.VectorBackendType;
+import org.apache.pinot.segment.spi.index.creator.VectorExecutionMode;
+
+
+/**
+ * Captures the runtime execution plan for a vector query within a single
segment.
+ *
+ * <p>This context is populated by the query planner (in {@code
FilterPlanNode}) and carried
+ * through to the filter operator so that execution decisions are explicit and
reportable.
+ * It is also surfaced in explain/debug output.</p>
+ *
+ * <p>The context is immutable once created. Use the {@link Builder} to
construct instances.</p>
+ */
+public final class VectorQueryExecutionContext {
+
+ private final VectorExecutionMode _executionMode;
+ private final VectorBackendType _backendType;
+ private final VectorBackendCapabilities _capabilities;
+ private final boolean _exactRerank;
+ private final boolean _hasMetadataFilter;
+ private final boolean _hasThresholdPredicate;
+ private final int _topK;
+ private final int _candidateBudget;
+ private final float _distanceThreshold;
+ @Nullable
+ private final String _fallbackReason;
+
+ private VectorQueryExecutionContext(Builder builder) {
+ _executionMode = builder._executionMode;
+ _backendType = builder._backendType;
+ _capabilities = builder._capabilities;
+ _exactRerank = builder._exactRerank;
+ _hasMetadataFilter = builder._hasMetadataFilter;
+ _hasThresholdPredicate = builder._hasThresholdPredicate;
+ _topK = builder._topK;
+ _candidateBudget = builder._candidateBudget;
+ _distanceThreshold = builder._distanceThreshold;
+ _fallbackReason = builder._fallbackReason;
+ }
+
+ public VectorExecutionMode getExecutionMode() {
+ return _executionMode;
+ }
+
+ public VectorBackendType getBackendType() {
+ return _backendType;
+ }
+
+ public VectorBackendCapabilities getCapabilities() {
+ return _capabilities;
+ }
+
+ public boolean isExactRerank() {
+ return _exactRerank;
+ }
+
+ public boolean hasMetadataFilter() {
+ return _hasMetadataFilter;
+ }
+
+ public boolean hasThresholdPredicate() {
+ return _hasThresholdPredicate;
+ }
+
+ public int getTopK() {
+ return _topK;
+ }
+
+ public int getCandidateBudget() {
+ return _candidateBudget;
+ }
+
+ public float getDistanceThreshold() {
+ return _distanceThreshold;
+ }
+
+ @Nullable
+ public String getFallbackReason() {
+ return _fallbackReason;
+ }
+
+ @Override
+ public String toString() {
+ StringBuilder sb = new StringBuilder("VectorQueryExecutionContext{");
+ sb.append("mode=").append(_executionMode);
+ sb.append(", backend=").append(_backendType);
+ sb.append(", topK=").append(_topK);
+ if (_hasMetadataFilter) {
+ sb.append(", hasFilter=true");
+ }
+ if (_hasThresholdPredicate) {
+ sb.append(", threshold=").append(_distanceThreshold);
+ }
+ if (_exactRerank) {
+ sb.append(", rerank=true, candidateBudget=").append(_candidateBudget);
+ }
+ if (_fallbackReason != null) {
+ sb.append(", fallbackReason=").append(_fallbackReason);
+ }
+ sb.append('}');
+ return sb.toString();
+ }
+
+ /**
+ * Selects the execution mode for a vector query based on backend
capabilities and query shape.
+ *
+ * <p>This is the centralized decision point. The rules are conservative and
explicit:</p>
+ * <ol>
+ * <li>No vector index -> EXACT_SCAN</li>
+ * <li>Threshold query + filter -> ANN_THRESHOLD_THEN_FILTER</li>
+ * <li>Threshold query (no filter) -> ANN_THRESHOLD_SCAN</li>
+ * <li>Filter + rerank -> ANN_THEN_FILTER_THEN_RERANK</li>
+ * <li>Filter (no rerank) -> ANN_THEN_FILTER</li>
+ * <li>Rerank (no filter) -> ANN_TOP_K_WITH_RERANK</li>
+ * <li>Plain top-K -> ANN_TOP_K</li>
+ * </ol>
+ */
+ public static VectorExecutionMode selectExecutionMode(boolean hasVectorIndex,
+ boolean hasMetadataFilter, boolean hasThresholdPredicate, boolean
exactRerank,
+ @Nullable VectorBackendCapabilities capabilities) {
Review Comment:
Fixed. Removed the `capabilities` parameter from `selectExecutionMode()`.
The method now takes 4 booleans. Updated Javadoc to say "based on query shape"
instead of "based on backend capabilities and query shape". `FILTER_THEN_ANN`
remains defined for future backends but is correctly unreachable until a
backend declares `supportsFilterAwareSearch=true`.
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java:
##########
@@ -223,7 +223,13 @@ private BaseFilterOperator
constructPhysicalOperator(FilterContext filter, int n
childFilters = filter.getChildren();
childFilterOperators = new ArrayList<>(childFilters.size());
for (FilterContext childFilter : childFilters) {
- BaseFilterOperator childFilterOperator =
constructPhysicalOperator(childFilter, numDocs);
+ BaseFilterOperator childFilterOperator;
+ if (isVectorSimilarityFilter(childFilter) && childFilters.size() >
1) {
+ // Pass filtered context so vector operator reports correct
execution mode
+ childFilterOperator = constructFilteredVectorOperator(childFilter,
numDocs);
+ } else {
+ childFilterOperator = constructPhysicalOperator(childFilter,
numDocs);
+ }
Review Comment:
Fixed. Added `hasNonVectorSibling(childFilters)` check instead of
`childFilters.size() > 1`. The flag is now only set when the AND has at least
one non-VECTOR_SIMILARITY sibling. Also removed the over-fetch entirely —
`vectorSimilarity(col, q, 10)` now always returns at most 10 docs regardless of
metadata filters.
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VectorTest.java:
##########
@@ -208,6 +211,85 @@ public void testVectorSimilarity(boolean
useMultiStageQueryEngine)
}
}
+ /**
+ * Tests filtered ANN: VECTOR_SIMILARITY combined with a metadata filter.
+ * All returned rows must match the filter predicate and be ordered by
distance.
+ */
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testFilteredAnn(boolean useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ int topK = 5;
+ String queryVector = "ARRAY[1.1" + StringUtils.repeat(", 1.1",
VECTOR_DIM_SIZE - 1) + "]";
+ String targetCategory = "cat_0";
+
+ String filteredQuery = 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_1, queryVector, CATEGORY, getTableName(),
+ VECTOR_1, queryVector, topK * 10, CATEGORY, targetCategory, topK);
+
+ JsonNode result = postQuery(filteredQuery);
+ JsonNode rows = result.get("resultTable").get("rows");
+ assertTrue(rows.size() > 0, "Filtered ANN should return results");
+
+ double prevDist = -1;
+ for (int i = 0; i < rows.size(); i++) {
+ double dist = rows.get(i).get(0).asDouble();
+ String cat = rows.get(i).get(1).asText();
+ assertEquals(cat, targetCategory, "All results must match the filter");
+ assertTrue(dist >= prevDist, "Results must be ordered by distance");
+ prevDist = dist;
+ }
+ }
+
+ /**
+ * Tests that filtered ANN returns fewer or equal results compared to
unfiltered ANN.
+ */
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testFilteredAnnReducesResults(boolean useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ int topK = 20;
+ String queryVector = "ARRAY[1.1" + StringUtils.repeat(", 1.1",
VECTOR_DIM_SIZE - 1) + "]";
+
+ String unfilteredQuery = String.format(
+ "SELECT count(*) FROM %s WHERE vectorSimilarity(%s, %s, %d)",
+ getTableName(), VECTOR_1, queryVector, topK);
+ String filteredQuery = String.format(
+ "SELECT count(*) FROM %s WHERE vectorSimilarity(%s, %s, %d) AND %s =
'cat_0'",
+ getTableName(), VECTOR_1, queryVector, topK, CATEGORY);
+
+ long unfilteredCount =
postQuery(unfilteredQuery).get("resultTable").get("rows").get(0).get(0).asLong();
+ long filteredCount =
postQuery(filteredQuery).get("resultTable").get("rows").get(0).get(0).asLong();
+
+ assertTrue(filteredCount <= unfilteredCount,
+ "Filtered (" + filteredCount + ") should be <= unfiltered (" +
unfilteredCount + ")");
+ assertTrue(filteredCount > 0, "Filtered ANN should return at least 1
result");
+ }
+
+ /**
+ * Tests that EXPLAIN output includes execution mode for vector queries.
+ */
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testExplainShowsExecutionMode(boolean useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ String queryVector = "ARRAY[1.1" + StringUtils.repeat(", 1.1",
VECTOR_DIM_SIZE - 1) + "]";
+
+ String explainQuery = String.format(
+ "set explainAskingServers=true; EXPLAIN PLAN FOR "
+ + "SELECT cosineDistance(%s, %s) AS dist FROM %s "
+ + "WHERE vectorSimilarity(%s, %s, %d) ORDER BY dist ASC LIMIT %d",
+ VECTOR_1, queryVector, getTableName(), VECTOR_1, queryVector, 5, 5);
+
+ JsonNode result = postQuery(explainQuery);
+ String explain = result.get("resultTable").toString();
+ assertTrue(explain.contains("executionMode") ||
explain.contains("VECTOR_SIMILARITY"),
+ "EXPLAIN should contain execution mode info: " + explain);
Review Comment:
Fixed. Now asserts `explain.contains("executionMode")` without the
`VECTOR_SIMILARITY` fallback.
##########
pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/creator/VectorExecutionModeTest.java:
##########
@@ -0,0 +1,95 @@
+/**
+ * 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.spi.index.creator;
+
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/**
+ * Tests for the {@link VectorExecutionMode} enum.
+ */
+public class VectorExecutionModeTest {
+
+ @Test
+ public void testAllModesHaveDescription() {
+ for (VectorExecutionMode mode : VectorExecutionMode.values()) {
+ assertNotNull(mode.getDescription());
+ assertFalse(mode.getDescription().isEmpty(), "Description must not be
empty for " + mode);
+ }
+ }
+
+ @Test
+ public void testIsApproximate() {
+ assertTrue(VectorExecutionMode.ANN_TOP_K.isApproximate());
+ assertTrue(VectorExecutionMode.ANN_TOP_K_WITH_RERANK.isApproximate());
+ assertTrue(VectorExecutionMode.ANN_THEN_FILTER.isApproximate());
+
assertTrue(VectorExecutionMode.ANN_THEN_FILTER_THEN_RERANK.isApproximate());
+ assertTrue(VectorExecutionMode.FILTER_THEN_ANN.isApproximate());
+ assertTrue(VectorExecutionMode.ANN_THRESHOLD_SCAN.isApproximate());
+ assertTrue(VectorExecutionMode.ANN_THRESHOLD_THEN_FILTER.isApproximate());
+ assertFalse(VectorExecutionMode.EXACT_SCAN.isApproximate());
+ }
+
+ @Test
+ public void testHasMetadataFilter() {
+ assertFalse(VectorExecutionMode.ANN_TOP_K.hasMetadataFilter());
+ assertFalse(VectorExecutionMode.ANN_TOP_K_WITH_RERANK.hasMetadataFilter());
+ assertTrue(VectorExecutionMode.ANN_THEN_FILTER.hasMetadataFilter());
+
assertTrue(VectorExecutionMode.ANN_THEN_FILTER_THEN_RERANK.hasMetadataFilter());
+ assertTrue(VectorExecutionMode.FILTER_THEN_ANN.hasMetadataFilter());
+ assertFalse(VectorExecutionMode.ANN_THRESHOLD_SCAN.hasMetadataFilter());
+
assertTrue(VectorExecutionMode.ANN_THRESHOLD_THEN_FILTER.hasMetadataFilter());
+ assertFalse(VectorExecutionMode.EXACT_SCAN.hasMetadataFilter());
+ }
+
+ @Test
+ public void testHasExactRefinement() {
+ assertFalse(VectorExecutionMode.ANN_TOP_K.hasExactRefinement());
+ assertTrue(VectorExecutionMode.ANN_TOP_K_WITH_RERANK.hasExactRefinement());
+ assertFalse(VectorExecutionMode.ANN_THEN_FILTER.hasExactRefinement());
+
assertTrue(VectorExecutionMode.ANN_THEN_FILTER_THEN_RERANK.hasExactRefinement());
+ assertFalse(VectorExecutionMode.FILTER_THEN_ANN.hasExactRefinement());
+ assertTrue(VectorExecutionMode.ANN_THRESHOLD_SCAN.hasExactRefinement());
+
assertTrue(VectorExecutionMode.ANN_THRESHOLD_THEN_FILTER.hasExactRefinement());
+ assertTrue(VectorExecutionMode.EXACT_SCAN.hasExactRefinement());
+ }
+
+ @Test
+ public void testIsThresholdMode() {
+ assertFalse(VectorExecutionMode.ANN_TOP_K.isThresholdMode());
+ assertFalse(VectorExecutionMode.ANN_TOP_K_WITH_RERANK.isThresholdMode());
+ assertFalse(VectorExecutionMode.ANN_THEN_FILTER.isThresholdMode());
+
assertFalse(VectorExecutionMode.ANN_THEN_FILTER_THEN_RERANK.isThresholdMode());
+ assertFalse(VectorExecutionMode.FILTER_THEN_ANN.isThresholdMode());
+ assertTrue(VectorExecutionMode.ANN_THRESHOLD_SCAN.isThresholdMode());
+
assertTrue(VectorExecutionMode.ANN_THRESHOLD_THEN_FILTER.isThresholdMode());
+ assertFalse(VectorExecutionMode.EXACT_SCAN.isThresholdMode());
+ }
+
+ @Test
+ public void testModeCount() {
+ assertEquals(VectorExecutionMode.values().length, 8,
+ "Phase 3 defines exactly 8 execution modes");
Review Comment:
Fixed. Replaced `testModeCount()` with `testExpectedModesExist()` that
asserts each expected mode by name via `valueOf()`. New modes won't break this
test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
