Auto prefiltering for queries on dense semantic_text fields (elastic#138989)

`knn` queries allow specifying `filters` that will be applied before
the knn search. This `pre-filtering` allows the `knn` to return `k`
results. If such filters are to applied only after the `knn` executes,
then the `knn` returns the `k` matching results but the filters can
filter out some of them thus potentially returning fewer than `k` results.

`semantic_text` fields can be queried with:

DSL

- `match` queries
- `semantic` queries
- `knn` queries

ES|QL

- `match` queries
- `knn` queries

For DSL, `knn` queries allow users to specify direct prefilters.
However, `match` and `semantic` queries provide no way to do so.
Same goes for ES|QL `match`. Noting that ES|QL `KNN` already implements
auto pre-filtering where conjunctions are pushed down to the `knn` query
as prefilters.

This commit implements semantic_text auto pre-filtering for `semantic_text` queries
in DSL (`match` and `semantic` queries) and ES|QL (`MATCH`).

We achieve this by adding an `AutoPrefilteringScope` object to the
`SearchExecutionContext`. When we convert a `bool` query to a lucene query,
we push its `must`, `filter`, and `must_not` clauses to the `AutoPrefilteringScope`.
At that stage queries have already been rewritten. Semantic queries using
`text_embedding` inference endpoints are rewritten to knn vector queries
that are auto-prefiltering enabled. Then, when an auto-prefiltering
enabled knn vector query is converted to its lucene equivalent, we
fetch the prefilters from the `SearchExecutionContext` and we apply
them to the knn vector query - which supports pre-filtering already.

ES|QL queries that contain `MATCH` automatically benefit from this implementation because
they are rewritten in `bool` queries.

Limitations

DSL

- nested queries are excluded from pre-filtering (elastic#138184)

ES|QL

- filters that are not translatable to lucene queries will be applied as post-filters

Relates elastic#132068

Author: dimitris-athanasiou

docs/changelog/138989.yaml

@@ -0,0 +1,5 @@
+pr: 138989
+summary: Auto prefiltering for queries on dense `semantic_text` fields
+area: Vector Search
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/index/query/BoolQueryBuilder.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.index.query.support.AutoPrefilteringScope;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -318,18 +319,33 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
         return adjustPureNegative ? fixNegativeQueryIfNeeded(query) : query;
     }
 
-    private static void addBooleanClauses(
+    private void addBooleanClauses(
         SearchExecutionContext context,
         BooleanQuery.Builder booleanQueryBuilder,
         List<QueryBuilder> clauses,
         Occur occurs
     ) throws IOException {
         for (QueryBuilder query : clauses) {
-            Query luceneQuery = query.toQuery(context);
-            booleanQueryBuilder.add(new BooleanClause(luceneQuery, occurs));
+            try (AutoPrefilteringScope autoPrefilteringScope = context.autoPrefilteringScope()) {
+                autoPrefilteringScope.push(collectPrefilters(query));
+                Query luceneQuery = query.toQuery(context);
+                booleanQueryBuilder.add(new BooleanClause(luceneQuery, occurs));
+            }
         }
     }
 
+    private List<QueryBuilder> collectPrefilters(QueryBuilder excluded) {
+        List<QueryBuilder> prefilters = new ArrayList<>();
+        mustClauses.stream().filter(q -> q != excluded).forEach(prefilters::add);
+        filterClauses.stream().filter(q -> q != excluded).forEach(prefilters::add);
+        mustNotClauses.stream().filter(q -> q != excluded).map(BoolQueryBuilder::invertQuery).forEach(prefilters::add);
+        return prefilters;
+    }
+
+    private static QueryBuilder invertQuery(QueryBuilder queryBuilder) {
+        return QueryBuilders.boolQuery().mustNot(queryBuilder);
+    }
+
     @Override
     protected int doHashCode() {
         return Objects.hash(adjustPureNegative, minimumShouldMatch, mustClauses, shouldClauses, mustNotClauses, filterClauses);

server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java

@@ -22,15 +22,19 @@
  * Wrapper for instances of {@link QueryBuilder} that have been intercepted using the {@link QueryRewriteInterceptor} to
  * break out of the rewrite phase. These instances are unwrapped on serialization.
  */
-class InterceptedQueryBuilderWrapper implements QueryBuilder {
+public class InterceptedQueryBuilderWrapper implements QueryBuilder {
 
     protected final QueryBuilder queryBuilder;
 
-    InterceptedQueryBuilderWrapper(QueryBuilder queryBuilder) {
+    public InterceptedQueryBuilderWrapper(QueryBuilder queryBuilder) {
         super();
         this.queryBuilder = queryBuilder;
     }
 
+    public QueryBuilder query() {
+        return queryBuilder;
+    }
+
     @Override
     public QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException {
         QueryRewriteInterceptor queryRewriteInterceptor = queryRewriteContext.getQueryRewriteInterceptor();

server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java

@@ -49,6 +49,7 @@
 import org.elasticsearch.index.mapper.ParsedDocument;
 import org.elasticsearch.index.mapper.SourceLoader;
 import org.elasticsearch.index.mapper.SourceToParse;
+import org.elasticsearch.index.query.support.AutoPrefilteringScope;
 import org.elasticsearch.index.query.support.NestedScope;
 import org.elasticsearch.index.similarity.SimilarityService;
 import org.elasticsearch.script.Script;
@@ -103,6 +104,7 @@ public class SearchExecutionContext extends QueryRewriteContext {
 
     private final Map<String, Query> namedQueries = new HashMap<>();
     private NestedScope nestedScope;
+    private AutoPrefilteringScope autoPrefilteringScope;
     private QueryBuilder aliasFilter;
     private boolean rewriteToNamedQueries = false;
 
@@ -291,6 +293,7 @@ private SearchExecutionContext(
         this.bitsetFilterCache = bitsetFilterCache;
         this.indexFieldDataLookup = indexFieldDataLookup;
         this.nestedScope = new NestedScope();
+        this.autoPrefilteringScope = new AutoPrefilteringScope();
         this.searcher = searcher;
         this.requestSize = requestSize;
         this.mapperMetrics = mapperMetrics;
@@ -301,7 +304,7 @@ private void reset() {
         this.lookup = null;
         this.namedQueries.clear();
         this.nestedScope = new NestedScope();
-
+        this.autoPrefilteringScope = new AutoPrefilteringScope();
     }
 
     // Set alias filter, so it can be applied for queries that need it (e.g. knn query)
@@ -556,6 +559,10 @@ public NestedScope nestedScope() {
         return nestedScope;
     }
 
+    public AutoPrefilteringScope autoPrefilteringScope() {
+        return autoPrefilteringScope;
+    }
+
     public IndexVersion indexVersionCreated() {
         return indexSettings.getIndexVersionCreated();
     }

server/src/main/java/org/elasticsearch/index/query/support/AutoPrefilteringScope.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.query.support;
+
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.SearchExecutionContext;
+
+import java.util.Deque;
+import java.util.LinkedList;
+import java.util.List;
+
+/**
+ * Keeps track of queries to be used as prefilters.
+ * During {@link QueryBuilder#toQuery(SearchExecutionContext)}, each query pushes queries to be used
+ * as prefilters to the {@link AutoPrefilteringScope}. Queries that need to apply prefilters can
+ * fetch them by calling {@link #getPrefilters()}.
+ *
+ * The scope is implemented as a stack {@link Deque} of lists of prefilters.
+ * As we move down the query tree, each query may push a list of prefilters.
+ * A query that consumes prefilters fetches a flattened list of all prefilters in scope via {@link #getPrefilters()}.
+ * When the query leaves the scope, {@link #pop()} should be called to remove the latest list of prefilters from the stack.
+ * This way queries in other query tree branches will not fetch irrelevant prefilters.
+ */
+public final class AutoPrefilteringScope implements Releasable {
+
+    private final Deque<List<QueryBuilder>> prefiltersStack = new LinkedList<>();
+
+    /**
+     * Pushes a list of prefilters to the scope.
+     */
+    public void push(List<QueryBuilder> prefilters) {
+        prefiltersStack.push(prefilters);
+    }
+
+    /**
+     * Removes the latest list of prefilters from the scope.
+     */
+    public void pop() {
+        prefiltersStack.pop();
+    }
+
+    /**
+     * Returns all prefilters in scope.
+     */
+    public List<QueryBuilder> getPrefilters() {
+        return prefiltersStack.stream().flatMap(List::stream).toList();
+    }
+
+    @Override
+    public void close() {
+        pop();
+    }
+}

server/src/main/java/org/elasticsearch/index/query/support/AutoPrefilteringUtils.java

@@ -0,0 +1,142 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.query.support;
+
+import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.BoostingQueryBuilder;
+import org.elasticsearch.index.query.ConstantScoreQueryBuilder;
+import org.elasticsearch.index.query.DisMaxQueryBuilder;
+import org.elasticsearch.index.query.InterceptedQueryBuilderWrapper;
+import org.elasticsearch.index.query.NestedQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+
+public final class AutoPrefilteringUtils {
+
+    private AutoPrefilteringUtils() {}
+
+    /**
+     * Prunes any query branch on queries whose type is included in the prunedTypes set.
+     * Note that queries may not preserve their scoring parameters as this assumes the returned
+     * query is going to be used for filtering only.
+     *
+     * @param query the root query to start pruning from
+     * @param prunedTypes the types of queries to prune
+     * @return an {@link Optional} containing the pruned query, or {@link Optional#empty()} if the query was pruned to nothing.
+     */
+    public static Optional<QueryBuilder> pruneQuery(QueryBuilder query, Set<Class<? extends QueryBuilder>> prunedTypes) {
+        if (prunedTypes.contains(query.getClass())) {
+            // Matched a pruned query type, prune away!
+            return Optional.empty();
+        }
+
+        return switch (query) {
+            case BoolQueryBuilder boolQuery -> pruneBoolQuery(boolQuery, prunedTypes);
+            // We only need the positive query here as the negative query is used for scoring only - we are filtering.
+            case BoostingQueryBuilder boostingQuery -> pruneQuery(boostingQuery.positiveQuery(), prunedTypes);
+            case ConstantScoreQueryBuilder constantScoreQuery -> {
+                Optional<QueryBuilder> pruned = pruneQuery(constantScoreQuery.innerQuery(), prunedTypes);
+                yield pruned.map(q -> q == constantScoreQuery.innerQuery() ? constantScoreQuery : new ConstantScoreQueryBuilder(q));
+            }
+            case DisMaxQueryBuilder disMaxQuery -> pruneDisMaxQuery(disMaxQuery, prunedTypes);
+            case FunctionScoreQueryBuilder functionScoreQuery -> pruneFunctionScoreQuery(functionScoreQuery, prunedTypes);
+            case InterceptedQueryBuilderWrapper interceptedQuery -> {
+                Optional<QueryBuilder> pruned = pruneQuery(interceptedQuery.query(), prunedTypes);
+                yield pruned.map(q -> q == interceptedQuery.query() ? interceptedQuery : new InterceptedQueryBuilderWrapper(q));
+            }
+            case NestedQueryBuilder nestedQuery -> {
+                Optional<QueryBuilder> pruned = pruneQuery(nestedQuery.query(), prunedTypes);
+                yield pruned.map(
+                    q -> q == nestedQuery.query() ? nestedQuery : new NestedQueryBuilder(nestedQuery.path(), q, nestedQuery.scoreMode())
+                );
+            }
+            default -> Optional.of(query);
+        };
+    }
+
+    private static Optional<QueryBuilder> pruneBoolQuery(BoolQueryBuilder boolQuery, Set<Class<? extends QueryBuilder>> prunedTypes) {
+        BoolQueryBuilder prunedBool = new BoolQueryBuilder();
+        pruneQueries(boolQuery.must(), prunedTypes).forEach(prunedBool::must);
+        pruneQueries(boolQuery.should(), prunedTypes).forEach(prunedBool::should);
+        pruneQueries(boolQuery.filter(), prunedTypes).forEach(prunedBool::filter);
+        pruneQueries(boolQuery.mustNot(), prunedTypes).forEach(prunedBool::mustNot);
+        adjustMinimumShouldMatchForPrunedShouldClauses(boolQuery, prunedBool);
+
+        if (prunedBool.equals(boolQuery)) {
+            return Optional.of(boolQuery);
+        }
+
+        // No need to preserve scoring parameters for filtering.
+        return prunedBool.hasClauses() ? Optional.of(prunedBool) : Optional.empty();
+    }
+
+    private static void adjustMinimumShouldMatchForPrunedShouldClauses(BoolQueryBuilder originalBool, BoolQueryBuilder prunedBool) {
+        if (originalBool.minimumShouldMatch() == null) {
+            return;
+        }
+        if (prunedBool.should().size() == originalBool.should().size()) {
+            prunedBool.minimumShouldMatch(originalBool.minimumShouldMatch());
+        } else {
+            int originalMsm = Queries.calculateMinShouldMatch(originalBool.should().size(), originalBool.minimumShouldMatch());
+            int numPrunedClauses = originalBool.should().size() - prunedBool.should().size();
+            // We need to adjust the minimum should match to account for the pruned clauses.
+            // We considered the following approaches:
+            // 1. strict approach: set to min(remaining_should_clauses, original_msm)
+            // 2. lenient approach: if msm is set and at least one should clause is pruned, prune all should clauses.
+            // 3. middle ground approach: set to max(0, original_msm - remaining_should_clauses)
+            // Let us imagine a query with 5 should clauses. 2 get pruned. msm is 3. 1 remaining clause matches.
+            // Approach 1 would make the entire bool query to not match as we would retain msm of 3 but only 1 clause would match.
+            // We do not know whether the pruned clauses would match or not. Thus, this approach seems too restrictive.
+            // Approach 2 would mean we prune all should clauses and the query would match,
+            // even if none of the remaining should clauses match.
+            // Approach 3 would mean we adjust the msm to 3 - 2 = 1. This would mean that the query would match if at least one
+            // of the remaining clauses matches.
+            // We opt for the lenient approach. It is as if we assume the pruned clauses matched. Seems to be the best compromise.
+            int prunedMsm = Math.max(0, originalMsm - numPrunedClauses);
+            prunedBool.minimumShouldMatch(prunedMsm);
+        }
+    }
+
+    private static Optional<QueryBuilder> pruneDisMaxQuery(DisMaxQueryBuilder disMaxQuery, Set<Class<? extends QueryBuilder>> prunedTypes) {
+        DisMaxQueryBuilder builder = new DisMaxQueryBuilder();
+        for (QueryBuilder innerQuery : disMaxQuery.innerQueries()) {
+            Optional<QueryBuilder> pruned = pruneQuery(innerQuery, prunedTypes);
+            pruned.ifPresent(builder::add);
+        }
+        if (builder.innerQueries().equals(disMaxQuery.innerQueries())) {
+            return Optional.of(disMaxQuery);
+        }
+        // No need to preserve tiebreaks for filtering.
+        return builder.innerQueries().isEmpty() ? Optional.empty() : Optional.of(builder);
+    }
+
+    private static Optional<QueryBuilder> pruneFunctionScoreQuery(
+        FunctionScoreQueryBuilder functionScoreQuery,
+        Set<Class<? extends QueryBuilder>> prunedTypes
+    ) {
+        // We could remove the function score entirely as it should not be helpful for filtering,
+        // but leaving it in for now to preserve the original query.
+        Optional<QueryBuilder> pruned = pruneQuery(functionScoreQuery.query(), prunedTypes);
+        return pruned.map(
+            q -> q == functionScoreQuery.query()
+                ? functionScoreQuery
+                : new FunctionScoreQueryBuilder(q, functionScoreQuery.filterFunctionBuilders())
+        );
+    }
+
+    private static List<QueryBuilder> pruneQueries(List<QueryBuilder> queries, Set<Class<? extends QueryBuilder>> prunedTypes) {
+        return queries.stream().map(q -> pruneQuery(q, prunedTypes)).filter(Optional::isPresent).map(Optional::get).toList();
+    }
+}