POC-2 - Auto prefiltering for semantic text queries

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

@@ -19,17 +19,21 @@
 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;
 import org.elasticsearch.xcontent.XContentParser;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.elasticsearch.common.lucene.search.Queries.fixNegativeQueryIfNeeded;
 
@@ -299,16 +303,17 @@ public String getWriteableName() {
     @Override
     protected Query doToQuery(SearchExecutionContext context) throws IOException {
         BooleanQuery.Builder booleanQueryBuilder = new BooleanQuery.Builder();
-        addBooleanClauses(context, booleanQueryBuilder, mustClauses, BooleanClause.Occur.MUST);
+        final List<QueryBuilder> prefilters = collectPrefilters();
+        addBooleanClauses(context, booleanQueryBuilder, mustClauses, BooleanClause.Occur.MUST, prefilters);
         try {
             // disable tracking of the @timestamp range for must_not and should clauses
             context.setTrackTimeRangeFilterFrom(false);
-            addBooleanClauses(context, booleanQueryBuilder, mustNotClauses, BooleanClause.Occur.MUST_NOT);
-            addBooleanClauses(context, booleanQueryBuilder, shouldClauses, BooleanClause.Occur.SHOULD);
+            addBooleanClauses(context, booleanQueryBuilder, mustNotClauses, BooleanClause.Occur.MUST_NOT, List.of());
+            addBooleanClauses(context, booleanQueryBuilder, shouldClauses, BooleanClause.Occur.SHOULD, prefilters);
         } finally {
             context.setTrackTimeRangeFilterFrom(true);
         }
-        addBooleanClauses(context, booleanQueryBuilder, filterClauses, BooleanClause.Occur.FILTER);
+        addBooleanClauses(context, booleanQueryBuilder, filterClauses, BooleanClause.Occur.FILTER, List.of());
         BooleanQuery booleanQuery = booleanQueryBuilder.build();
         if (booleanQuery.clauses().isEmpty()) {
             return new MatchAllDocsQuery();
@@ -318,15 +323,25 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
         return adjustPureNegative ? fixNegativeQueryIfNeeded(query) : query;
     }
 
+    private List<QueryBuilder> collectPrefilters() {
+        return Stream.of(mustClauses, mustNotClauses.stream().map(c -> QueryBuilders.boolQuery().mustNot(c)).toList(), filterClauses)
+            .flatMap(Collection::stream)
+            .collect(Collectors.toList());
+    }
+
     private static void addBooleanClauses(
         SearchExecutionContext context,
         BooleanQuery.Builder booleanQueryBuilder,
         List<QueryBuilder> clauses,
-        Occur occurs
+        Occur occurs,
+        List<QueryBuilder> prefilters
     ) throws IOException {
         for (QueryBuilder query : clauses) {
-            Query luceneQuery = query.toQuery(context);
-            booleanQueryBuilder.add(new BooleanClause(luceneQuery, occurs));
+            try (AutoPrefilteringScope autoPrefilteringScope = context.autoPrefilteringScope()) {
+                autoPrefilteringScope.push(prefilters.stream().filter(c -> c != query).toList());
+                Query luceneQuery = query.toQuery(context);
+                booleanQueryBuilder.add(new BooleanClause(luceneQuery, occurs));
+            }
         }
     }
 

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,41 @@
+/*
+ * 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.index.query.QueryBuilder;
+
+import java.util.Deque;
+import java.util.LinkedList;
+import java.util.List;
+
+/**
+ * During query parsing this keeps track of the current prefiltering level.
+ */
+public final class AutoPrefilteringScope implements AutoCloseable {
+
+    private final Deque<List<QueryBuilder>> prefiltersStack = new LinkedList<>();
+
+    public List<QueryBuilder> getPrefilters() {
+        return prefiltersStack.stream().flatMap(List::stream).toList();
+    }
+
+    public void push(List<QueryBuilder> prefilters) {
+        prefiltersStack.push(prefilters);
+    }
+
+    public void pop() {
+        prefiltersStack.pop();
+    }
+
+    @Override
+    public void close() {
+        pop();
+    }
+}

server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java

@@ -56,6 +56,9 @@
  * {@link org.apache.lucene.search.KnnByteVectorQuery}.
  */
 public class KnnVectorQueryBuilder extends AbstractQueryBuilder<KnnVectorQueryBuilder> {
+
+    public static final TransportVersion AUTO_PREFILTERING = TransportVersion.fromName("knn_vector_query_auto_prefiltering");
+
     public static final String NAME = "knn";
     private static final int NUM_CANDS_LIMIT = 10_000;
     private static final float NUM_CANDS_MULTIPLICATIVE_FACTOR = 1.5f;
@@ -121,7 +124,7 @@ public static KnnVectorQueryBuilder fromXContent(XContentParser parser) {
         return PARSER.apply(parser, null);
     }
 
-    private static final TransportVersion VISIT_PERCENTAGE = TransportVersion.fromName("visit_percentage");
+    public static final TransportVersion VISIT_PERCENTAGE = TransportVersion.fromName("visit_percentage");
 
     private final String fieldName;
     private final VectorData queryVector;
@@ -133,6 +136,7 @@ public static KnnVectorQueryBuilder fromXContent(XContentParser parser) {
     private final QueryVectorBuilder queryVectorBuilder;
     private final Supplier<float[]> queryVectorSupplier;
     private final RescoreVectorBuilder rescoreVectorBuilder;
+    private boolean isAutoPrefiltering = false;
 
     public KnnVectorQueryBuilder(
         String fieldName,
@@ -302,6 +306,9 @@ public KnnVectorQueryBuilder(StreamInput in) throws IOException {
         } else {
             this.rescoreVectorBuilder = null;
         }
+        if (in.getTransportVersion().supports(AUTO_PREFILTERING)) {
+            this.isAutoPrefiltering = in.readBoolean();
+        }
 
         this.queryVectorSupplier = null;
     }
@@ -357,6 +364,10 @@ public KnnVectorQueryBuilder addFilterQueries(List<QueryBuilder> filterQueries)
         return this;
     }
 
+    public boolean isAutoPrefiltering() {
+        return isAutoPrefiltering;
+    }
+
     @Override
     protected void doWriteTo(StreamOutput out) throws IOException {
         if (queryVectorSupplier != null) {
@@ -417,6 +428,9 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().supports(TransportVersions.V_8_18_0)) {
             out.writeOptionalWriteable(rescoreVectorBuilder);
         }
+        if (out.getTransportVersion().supports(AUTO_PREFILTERING)) {
+            out.writeBoolean(isAutoPrefiltering);
+        }
     }
 
     @Override
@@ -479,7 +493,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
                 visitPercentage,
                 rescoreVectorBuilder,
                 vectorSimilarity
-            ).boost(boost).queryName(queryName).addFilterQueries(filterQueries);
+            ).boost(boost).queryName(queryName).addFilterQueries(filterQueries).setAutoPrefiltering(isAutoPrefiltering);
         }
         if (queryVectorBuilder != null) {
             SetOnce<float[]> toSet = new SetOnce<>();
@@ -509,7 +523,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
                 visitPercentage,
                 rescoreVectorBuilder,
                 vectorSimilarity
-            ).boost(boost).queryName(queryName).addFilterQueries(filterQueries);
+            ).boost(boost).queryName(queryName).addFilterQueries(filterQueries).setAutoPrefiltering(isAutoPrefiltering);
         }
         boolean changed = false;
         List<QueryBuilder> rewrittenQueries = new ArrayList<>(filterQueries.size());
@@ -534,7 +548,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
                 visitPercentage,
                 rescoreVectorBuilder,
                 vectorSimilarity
-            ).boost(boost).queryName(queryName).addFilterQueries(rewrittenQueries);
+            ).boost(boost).queryName(queryName).addFilterQueries(rewrittenQueries).setAutoPrefiltering(isAutoPrefiltering);
         }
         if (ctx.convertToInnerHitsRewriteContext() != null) {
             QueryBuilder exactKnnQuery = new ExactKnnQueryBuilder(queryVector, fieldName, vectorSimilarity);
@@ -579,8 +593,9 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
         }
         DenseVectorFieldType vectorFieldType = (DenseVectorFieldType) fieldType;
 
-        List<Query> filtersInitial = new ArrayList<>(filterQueries.size());
-        for (QueryBuilder query : this.filterQueries) {
+        List<QueryBuilder> allApplicableFilters = getAllApplicableFilters(context);
+        List<Query> filtersInitial = new ArrayList<>(allApplicableFilters.size());
+        for (QueryBuilder query : allApplicableFilters) {
             filtersInitial.add(query.toQuery(context));
         }
         if (context.getAliasFilter() != null) {
@@ -650,6 +665,14 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
         );
     }
 
+    private List<QueryBuilder> getAllApplicableFilters(SearchExecutionContext context) {
+        List<QueryBuilder> applicableFilters = new ArrayList<>(filterQueries);
+        if (isAutoPrefiltering) {
+            applicableFilters.addAll(context.autoPrefilteringScope().getPrefilters());
+        }
+        return applicableFilters;
+    }
+
     private static Query buildFilterQuery(List<Query> filters) {
         BooleanQuery.Builder builder = new BooleanQuery.Builder();
         for (Query f : filters) {
@@ -671,7 +694,8 @@ protected int doHashCode() {
             filterQueries,
             vectorSimilarity,
             queryVectorBuilder,
-            rescoreVectorBuilder
+            rescoreVectorBuilder,
+            isAutoPrefiltering
         );
     }
 
@@ -685,11 +709,17 @@ protected boolean doEquals(KnnVectorQueryBuilder other) {
             && Objects.equals(filterQueries, other.filterQueries)
             && Objects.equals(vectorSimilarity, other.vectorSimilarity)
             && Objects.equals(queryVectorBuilder, other.queryVectorBuilder)
-            && Objects.equals(rescoreVectorBuilder, other.rescoreVectorBuilder);
+            && Objects.equals(rescoreVectorBuilder, other.rescoreVectorBuilder)
+            && isAutoPrefiltering == other.isAutoPrefiltering;
     }
 
     @Override
     public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.V_8_0_0;
     }
+
+    public KnnVectorQueryBuilder setAutoPrefiltering(boolean isAutoPrefiltering) {
+        this.isAutoPrefiltering = isAutoPrefiltering;
+        return this;
+    }
 }

server/src/main/java/org/elasticsearch/search/vectors/RescoreVectorBuilder.java

@@ -38,7 +38,7 @@ public class RescoreVectorBuilder implements Writeable, ToXContentObject {
         PARSER.declareFloat(ConstructingObjectParser.constructorArg(), OVERSAMPLE_FIELD);
     }
 
-    private static final TransportVersion RESCORE_VECTOR_ALLOW_ZERO = TransportVersion.fromName("rescore_vector_allow_zero");
+    public static final TransportVersion RESCORE_VECTOR_ALLOW_ZERO = TransportVersion.fromName("rescore_vector_allow_zero");
 
     // Oversample is required as of now as it is the only field in the rescore vector
     private final float oversample;

server/src/main/resources/transport/definitions/referable/knn_vector_query_auto_prefiltering.csv

@@ -0,0 +1 @@
+9216000

server/src/main/resources/transport/upper_bounds/9.3.csv

@@ -1 +1 @@
-inference_api_eis_authorization_persistent_task,9215000
+knn_vector_query_auto_prefiltering,9216000