Trying to add a PrefilteredQueryBuilder as a wrapper - does not work

carlosdelest · carlosdelest · commit fab91a4d2b26 · 2025-09-26T08:53:44.000+02:00
diff --git a/server/src/main/resources/transport/upper_bounds/8.18.csv b/server/src/main/resources/transport/upper_bounds/8.18.csv
@@ -1 +1 @@
-initial_elasticsearch_8_18_8,8840010
+transform_check_for_dangling_tasks,8840011
diff --git a/server/src/main/resources/transport/upper_bounds/8.19.csv b/server/src/main/resources/transport/upper_bounds/8.19.csv
@@ -1 +1 @@
-initial_elasticsearch_8_19_5,8841069
+transform_check_for_dangling_tasks,8841070
diff --git a/server/src/main/resources/transport/upper_bounds/9.0.csv b/server/src/main/resources/transport/upper_bounds/9.0.csv
@@ -1 +1 @@
-initial_elasticsearch_9_0_8,9000017
+transform_check_for_dangling_tasks,9000018
diff --git a/server/src/main/resources/transport/upper_bounds/9.1.csv b/server/src/main/resources/transport/upper_bounds/9.1.csv
@@ -1 +1 @@
-initial_elasticsearch_9_1_5,9112008
+transform_check_for_dangling_tasks,9112009
diff --git a/server/src/main/resources/transport/upper_bounds/9.2.csv b/server/src/main/resources/transport/upper_bounds/9.2.csv
@@ -1 +1 @@
-semantic_query_filter_queries,9169000
+prefiltered_query_builder,9171000
diff --git a/x-pack/plugin/esql/build.gradle b/x-pack/plugin/esql/build.gradle
@@ -53,6 +53,7 @@ dependencies {
     exclude(group: "org.elasticsearch.plugin", module: "esql")
   }
   testImplementation project(':test:framework')
+  testImplementation testArtifact(project(':server'))
   testImplementation(testArtifact(project(xpackModule('core'))))
   testImplementation project(path: xpackModule('enrich'))
   testImplementation project(path: xpackModule('spatial'))
diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/SemanticMatchTestCase.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/SemanticMatchTestCase.java
@@ -92,6 +92,41 @@ public void testCopyToWithSemanticText() throws IOException {
         );
     }
 
+    public void testMatchWithAndPrefilter() throws IOException {
+        assumeTrue("semantic text capability not available", EsqlCapabilities.Cap.SEMANTIC_TEXT_FIELD_CAPS.isEnabled());
+
+        StringBuilder bulkRequestBody = new StringBuilder();
+        int specificValueCount = 0;
+        for (int i = 0; i < 100; i++) {
+            bulkRequestBody.append("{\"index\":{\"_index\":\"test-semantic-prefilter\"}}\n");
+            String prefilterValue = "other_value";
+            if (i % 10 == 0) {
+                prefilterValue = "specific_value";
+                specificValueCount++;
+            }
+            bulkRequestBody.append("{\"semantic_text_field\":\"inference test\", \"prefilter_field\":\"")
+                .append(prefilterValue)
+                .append("\"}\n");
+        }
+
+        Request bulkRequest = new Request("POST", "/_bulk");
+        bulkRequest.addParameter("refresh", "true");
+        bulkRequest.setJsonEntity(bulkRequestBody.toString());
+        assertEquals(200, adminClient().performRequest(bulkRequest).getStatusLine().getStatusCode());
+
+        String query = """
+            FROM test-semantic-prefilter
+            | WHERE MATCH(semantic_text_field, "something") AND prefilter_field == "specific_value"
+            | LIMIT 10
+            """;
+        Map<String, Object> result = runEsqlQuery(query);
+
+        @SuppressWarnings("unchecked")
+        List<List<Object>> values = (List<List<Object>>) result.get("values");
+        assertNotNull(values);
+        assertEquals(specificValueCount, values.size());
+    }
+
     @Before
     public void setUpIndices() throws IOException {
         var settings = Settings.builder().build();
@@ -139,6 +174,19 @@ public void setUpIndices() throws IOException {
                 }
             """;
         createIndex(adminClient(), "test-semantic4", settings, mapping4);
+
+        String mappingFilter = """
+                "properties": {
+                  "semantic_text_field": {
+                   "type": "semantic_text",
+                   "inference_id": "test_dense_inference"
+                  },
+                  "prefilter_field": {
+                    "type": "keyword"
+                  }
+                }
+            """;
+        createIndex(adminClient(), "test-semantic-prefilter", settings, mappingFilter);
     }
 
     @Before
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java
@@ -164,7 +164,7 @@ public Translatable translatable(LucenePushdownPredicates pushdownPredicates) {
     }
 
     @Override
-    public Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
+    public final Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         return queryBuilder != null ? new TranslationAwareExpressionQuery(source(), queryBuilder) : translate(pushdownPredicates, handler);
     }
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/Match.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/Match.java
@@ -416,11 +416,6 @@ public Expression withPrefilters(List<Expression> filterExpressions) {
         return new Match(source(), field(), query(), options(), queryBuilder(), filterExpressions);
     }
 
-    @Override
-    protected Expression replaceFilteredQueryBuilder(QueryBuilder queryBuilder) {
-        return new Match(source(), field, query(), options(), queryBuilder, prefilterExpressions());
-    }
-
     @Override
     public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
         return (plan, failures) -> {
@@ -455,14 +450,20 @@ public Object queryAsObject() {
     }
 
     @Override
-    protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
+    protected Query translateInnerQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         var fieldAttribute = fieldAsFieldAttribute();
         Check.notNull(fieldAttribute, "Match must have a field attribute as the first argument");
         String fieldName = getNameFromFieldAttribute(fieldAttribute);
         // Make query lenient so mixed field types can be queried when a field type is incompatible with the value provided
+
         return new MatchQuery(source(), fieldName, queryAsObject(), matchQueryOptions());
     }
 
+    @Override
+    public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
+        return new Match(source(), field(), query(), options(), queryBuilder, prefilterExpressions());
+    }
+
     private FieldAttribute fieldAsFieldAttribute() {
         return fieldAsFieldAttribute(field);
     }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/MatchOperator.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/MatchOperator.java
@@ -92,7 +92,7 @@ public Expression replaceChildren(List<Expression> newChildren) {
     }
 
     @Override
-    public Expression replaceFilteredQueryBuilder(QueryBuilder queryBuilder) {
+    public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
         return new MatchOperator(source(), field, query(), queryBuilder, prefilterExpressions());
     }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/PrefilteredFullTextFunction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/PrefilteredFullTextFunction.java
@@ -11,14 +11,14 @@
 import org.elasticsearch.search.vectors.FilteredQueryBuilder;
 import org.elasticsearch.xpack.esql.capabilities.TranslationAware;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
+import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
+import org.elasticsearch.xpack.esql.querydsl.query.PrefilteredQuery;
 
-import java.util.ArrayList;
 import java.util.List;
-
-import static org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates.DEFAULT;
-import static org.elasticsearch.xpack.esql.planner.TranslatorHandler.TRANSLATOR_HANDLER;
+import java.util.stream.Collectors;
 
 public abstract class PrefilteredFullTextFunction extends FullTextFunction {
 
@@ -57,29 +57,28 @@ public List<Expression> prefilterExpressions() {
         return filterExpressions;
     }
 
-    public final Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
-        if (queryBuilder instanceof FilteredQueryBuilder<?> filteredQueryBuilder) {
-            queryBuilder = filteredQueryBuilder.addFilterQueries(prefilterQueryBuilders());
-        }
-
-        return replaceFilteredQueryBuilder(queryBuilder);
+    @Override
+    public final Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
+        PrefilteredQuery prefilteredQuery = new PrefilteredQuery(
+            source(),
+            translateInnerQuery(pushdownPredicates, handler),
+            prefiltersAsQueries(pushdownPredicates, handler)
+        );
+
+        return prefilteredQuery;
     }
 
-    private List<QueryBuilder> prefilterQueryBuilders() {
-        List<QueryBuilder> filterQueries = new ArrayList<>();
-        for (Expression filterExpression : prefilterExpressions()) {
-            if (filterExpression instanceof TranslationAware translationAware) {
-                // We can only translate filter expressions that are translatable. In case any is not translatable,
-                // Knn won't be pushed down so it's safe not to translate all filters and check them when creating an evaluator
-                // for the non-pushed down query
-                if (translationAware.translatable(DEFAULT) == Translatable.YES) {
-                    filterQueries.add(TRANSLATOR_HANDLER.asQuery(DEFAULT, filterExpression).toQueryBuilder());
-                }
-            }
-        }
-
-        return filterQueries;
+    protected abstract Query translateInnerQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler);
+
+    private List<Query> prefiltersAsQueries(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
+        return prefilterExpressions()
+            .stream()
+            // We can only translate filter expressions that are translatable. In case any is not translatable,
+            // Knn won't be pushed down so it's safe not to translate all filters and check them when creating an evaluator
+            // for the non-pushed down query
+            .filter(f -> f instanceof TranslationAware)
+            .map(TranslationAware.class::cast)
+            .map(ta -> ta.asQuery(pushdownPredicates, handler))
+            .collect(Collectors.toList());
     }
-
-    protected abstract Expression replaceFilteredQueryBuilder(QueryBuilder queryBuilder);
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/PrefilteredQueryBuilder.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/PrefilteredQueryBuilder.java
@@ -0,0 +1,149 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.expression.function.fulltext;
+
+import org.apache.lucene.search.Query;
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.index.query.AbstractQueryBuilder;
+import org.elasticsearch.index.query.NestedQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryRewriteContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.search.vectors.KnnVectorQueryBuilder;
+import org.elasticsearch.xcontent.ConstructingObjectParser;
+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.List;
+import java.util.Objects;
+
+import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
+
+public class PrefilteredQueryBuilder extends AbstractQueryBuilder<PrefilteredQueryBuilder> {
+
+    private static final TransportVersion PREFILTERED_QUERY_BUILDER = TransportVersion.fromName("prefiltered_query_builder");
+
+    public static final String NAME = "prefilter";
+
+    private static final ParseField QUERY_FIELD = new ParseField("query");
+    private static final ParseField FILTER_FIELD = new ParseField("filter");
+
+    @SuppressWarnings("unchecked")
+    private static final ConstructingObjectParser<PrefilteredQueryBuilder, Void> PARSER = new ConstructingObjectParser<>(
+        NAME,
+        false,
+        args -> new PrefilteredQueryBuilder((QueryBuilder) args[0], (List<QueryBuilder>) args[1])
+    );
+
+    static {
+        PARSER.declareObject(constructorArg(), (p, c) -> parseInnerQueryBuilder(p), QUERY_FIELD);
+        PARSER.declareFieldArray(
+            constructorArg(),
+            (p, c) -> AbstractQueryBuilder.parseTopLevelQuery(p),
+            FILTER_FIELD,
+            ObjectParser.ValueType.OBJECT_ARRAY
+        );
+        declareStandardFields(PARSER);
+    }
+
+    private final QueryBuilder innerQuery;
+    private final List<QueryBuilder> filters;
+
+    public PrefilteredQueryBuilder(QueryBuilder innerQuery, List<QueryBuilder> filters) {
+        assert innerQuery != null : "inner query cannot be null";
+        assert filters != null && filters.isEmpty() == false : "filters cannot be null and can't be empty";
+
+        this.innerQuery = innerQuery;
+        this.filters = filters;
+    }
+
+    public PrefilteredQueryBuilder(StreamInput in) throws IOException {
+        super(in);
+        this.innerQuery = in.readNamedWriteable(QueryBuilder.class);
+        this.filters = in.readNamedWriteableCollectionAsList(QueryBuilder.class);
+    }
+
+    QueryBuilder getInnerQuery() {
+        return innerQuery;
+    }
+
+    List<QueryBuilder> getFilters() {
+        return filters;
+    }
+
+    @Override
+    protected void doWriteTo(StreamOutput out) throws IOException {
+        out.writeNamedWriteable(innerQuery);
+        out.writeNamedWriteableCollection(filters);
+    }
+
+    public static PrefilteredQueryBuilder fromXContent(XContentParser parser) throws IOException {
+        return PARSER.apply(parser, null);
+    }
+
+    @Override
+    protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject(NAME);
+        builder.field(QUERY_FIELD.getPreferredName(), innerQuery);
+        builder.startArray(FILTER_FIELD.getPreferredName());
+        for (QueryBuilder filterQuery : filters) {
+            filterQuery.toXContent(builder, params);
+        }
+        builder.endArray();
+        boostAndQueryNameToXContent(builder);
+        builder.endObject();
+    }
+
+    @Override
+    protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
+        QueryBuilder rewrittenInnerQuery = innerQuery.rewrite(queryRewriteContext);
+        if (rewrittenInnerQuery != innerQuery) {
+            if (rewrittenInnerQuery instanceof NestedQueryBuilder nestedQueryBuilder) {
+                QueryBuilder nestedInnerQuery = nestedQueryBuilder.query();
+                if (nestedInnerQuery instanceof KnnVectorQueryBuilder knnVectorQueryBuilder) {
+                    knnVectorQueryBuilder.addFilterQueries(filters);
+                }
+            }
+            return new PrefilteredQueryBuilder(rewrittenInnerQuery, filters);
+        }
+
+        return this;
+    }
+
+    @Override
+    protected Query doToQuery(SearchExecutionContext context) throws IOException {
+        return innerQuery.toQuery(context);
+//        throw new IllegalStateException(NAME + " should have been rewritten to another query type");
+    }
+
+    @Override
+    protected boolean doEquals(PrefilteredQueryBuilder other) {
+        return Objects.equals(innerQuery, other.innerQuery)
+            && Objects.equals(filters, other.filters);
+    }
+
+    @Override
+    protected int doHashCode() {
+        return Objects.hash(innerQuery, filters);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return PREFILTERED_QUERY_BUILDER;
+    }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Knn.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Knn.java
@@ -194,11 +194,6 @@ public Expression options() {
         return options;
     }
 
-    @Override
-    public DataType dataType() {
-        return DataType.BOOLEAN;
-    }
-
     @Override
     protected TypeResolution resolveParams() {
         return resolveField().and(resolveQuery()).and(Options.resolve(options(), source(), THIRD, ALLOWED_OPTIONS));
@@ -243,12 +238,12 @@ public List<Number> queryAsObject() {
     }
 
     @Override
-    protected Expression replaceFilteredQueryBuilder(QueryBuilder queryBuilder) {
+    public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
         return new Knn(source(), field(), query(), options(), k(), queryBuilder, prefilterExpressions());
     }
 
     @Override
-    protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
+    protected Query translateInnerQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         assert k() != null : "Knn function must have a k value set before translation";
         var fieldAttribute = Match.fieldAsFieldAttribute(field());
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/querydsl/query/PrefilteredQuery.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/querydsl/query/PrefilteredQuery.java
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/fulltext/PrefilteredQueryBuilderTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/fulltext/PrefilteredQueryBuilderTests.java

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-initial_elasticsearch_8_18_8,8840010`
	`1`	`+transform_check_for_dangling_tasks,8840011`
Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-initial_elasticsearch_8_19_5,8841069`
	`1`	`+transform_check_for_dangling_tasks,8841070`
Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-initial_elasticsearch_9_0_8,9000017`
	`1`	`+transform_check_for_dangling_tasks,9000018`
Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-initial_elasticsearch_9_1_5,9112008`
	`1`	`+transform_check_for_dangling_tasks,9112009`
Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-semantic_query_filter_queries,9169000`
	`1`	`+prefiltered_query_builder,9171000`
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@ dependencies {`
`53`	`53`	`exclude(group: "org.elasticsearch.plugin", module: "esql")`
`54`	`54`	`}`
`55`	`55`	`testImplementation project(':test:framework')`
	`56`	`+ testImplementation testArtifact(project(':server'))`
`56`	`57`	`testImplementation(testArtifact(project(xpackModule('core'))))`
`57`	`58`	`testImplementation project(path: xpackModule('enrich'))`
`58`	`59`	`testImplementation project(path: xpackModule('spatial'))`
Original file line number	Diff line number	Diff line change
`@@ -164,7 +164,7 @@ public Translatable translatable(LucenePushdownPredicates pushdownPredicates) {`
`164`	`164`	`}`
`165`	`165`
`166`	`166`	`@Override`
`167`		`- public Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {`
	`167`	`+ public final Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {`
`168`	`168`	`return queryBuilder != null ? new TranslationAwareExpressionQuery(source(), queryBuilder) : translate(pushdownPredicates, handler);`
`169`	`169`	`}`
`170`	`170`
Original file line number	Diff line number	Diff line change
`@@ -92,7 +92,7 @@ public Expression replaceChildren(List<Expression> newChildren) {`
`92`	`92`	`}`
`93`	`93`
`94`	`94`	`@Override`
`95`		`- public Expression replaceFilteredQueryBuilder(QueryBuilder queryBuilder) {`
	`95`	`+ public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {`
`96`	`96`	`return new MatchOperator(source(), field, query(), queryBuilder, prefilterExpressions());`
`97`	`97`	`}`
`98`	`98`	`}`