QL: Add utility for combining QueryBuilders (#99031)

Generify an utility for combining QueryBuilder's while being mindful of
 the number of BoolQueryBuilders created in the process

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/assembler/BoxedQueryRequest.java

@@ -61,7 +61,7 @@ public class BoxedQueryRequest implements QueryRequest {
         timestampField = timestamp;
         keys = keyNames;
         this.optionalKeyNames = optionalKeyNames;
-        RuntimeUtils.addFilter(timestampRange, searchSource);
+        RuntimeUtils.combineFilters(searchSource, timestampRange);
     }
 
     @Override
@@ -181,7 +181,7 @@ public class BoxedQueryRequest implements QueryRequest {
             }
         }
 
-        RuntimeUtils.replaceFilter(keyFilters, newFilters, searchSource);
+        RuntimeUtils.replaceFilter(searchSource, keyFilters, newFilters);
         keyFilters = newFilters;
         return this;
     }

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/assembler/SampleQueryRequest.java

@@ -128,7 +128,7 @@ public class SampleQueryRequest implements QueryRequest {
             }
         }
 
-        RuntimeUtils.replaceFilter(multipleKeyFilters, newFilters, searchSource);
+        RuntimeUtils.replaceFilter(searchSource, multipleKeyFilters, newFilters);
         multipleKeyFilters = newFilters;
     }
 
@@ -158,7 +158,7 @@ public class SampleQueryRequest implements QueryRequest {
         }
 
         SearchSourceBuilder newSource = copySource(searchSource);
-        RuntimeUtils.replaceFilter(singleKeyPairFilters, newFilters, newSource);
+        RuntimeUtils.replaceFilter(newSource, singleKeyPairFilters, newFilters);
         // ask for the minimum needed to get at least N samplese per key
         int minResultsNeeded = maxStages + maxSamplesPerKey - 1;
         newSource.size(minResultsNeeded)

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/search/PITAwareQueryClient.java

@@ -106,7 +106,7 @@ public class PITAwareQueryClient extends BasicQueryClient {
         if (CollectionUtils.isEmpty(indices) == false) {
             request.indices(Strings.EMPTY_ARRAY);
             QueryBuilder indexQuery = indices.length == 1 ? termQuery(GetResult._INDEX, indices[0]) : termsQuery(GetResult._INDEX, indices);
-            RuntimeUtils.addFilter(indexQuery, source);
+            RuntimeUtils.combineFilters(source, indexQuery);
         }
     }
 

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/search/RuntimeUtils.java

@@ -37,6 +37,7 @@ import org.elasticsearch.xpack.ql.expression.gen.pipeline.HitExtractorInput;
 import org.elasticsearch.xpack.ql.expression.gen.pipeline.Pipe;
 import org.elasticsearch.xpack.ql.expression.gen.pipeline.ReferenceInput;
 import org.elasticsearch.xpack.ql.index.IndexResolver;
+import org.elasticsearch.xpack.ql.util.Queries;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -47,6 +48,7 @@ import java.util.Set;
 
 import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
 import static org.elasticsearch.xpack.ql.execution.search.extractor.AbstractFieldHitExtractor.MultiValueSupport.FULL;
+import static org.elasticsearch.xpack.ql.util.Queries.Clause.FILTER;
 
 public final class RuntimeUtils {
 
@@ -184,62 +186,29 @@ public final class RuntimeUtils {
         return Arrays.asList(response.getHits().getHits());
     }
 
-    // optimized method that adds filter to existing bool queries without additional wrapping
-    // additionally checks whether the given query exists for safe decoration
-    public static SearchSourceBuilder addFilter(QueryBuilder filter, SearchSourceBuilder source) {
-        BoolQueryBuilder bool = null;
-        QueryBuilder query = source.query();
-
-        if (query instanceof BoolQueryBuilder boolQueryBuilder) {
-            bool = boolQueryBuilder;
-            if (filter != null && bool.filter().contains(filter) == false) {
-                bool.filter(filter);
-            }
-        } else {
-            bool = boolQuery();
-            if (query != null) {
-                bool.filter(query);
-            }
-            if (filter != null) {
-                bool.filter(filter);
-            }
-
-            source.query(bool);
-        }
-        return source;
+    /**
+     * optimized method that adds filter to existing bool queries without additional wrapping
+     * additionally checks whether the given query exists for safe decoration
+     */
+    public static SearchSourceBuilder combineFilters(SearchSourceBuilder source, QueryBuilder filter) {
+        var query = Queries.combine(FILTER, Arrays.asList(source.query(), filter));
+        query = query == null ? boolQuery() : query;
+        return source.query(query);
     }
 
     public static SearchSourceBuilder replaceFilter(
+        SearchSourceBuilder source,
         List<QueryBuilder> oldFilters,
-        List<QueryBuilder> newFilters,
-        SearchSourceBuilder source
+        List<QueryBuilder> newFilters
     ) {
-        BoolQueryBuilder bool = null;
-        QueryBuilder query = source.query();
-
-        if (query instanceof BoolQueryBuilder boolQueryBuilder) {
-            bool = boolQueryBuilder;
-            if (oldFilters != null) {
-                bool.filter().removeAll(oldFilters);
-            }
-
-            if (newFilters != null) {
-                bool.filter().addAll(newFilters);
-            }
-        }
-        // no bool query means no old filters
-        else {
-            bool = boolQuery();
-            if (query != null) {
-                bool.filter(query);
-            }
-            if (newFilters != null) {
-                bool.filter().addAll(newFilters);
-            }
-
-            source.query(bool);
-        }
-        return source;
+        var query = source.query();
+        query = removeFilters(query, oldFilters);
+        query = Queries.combine(
+            FILTER,
+            org.elasticsearch.xpack.ql.util.CollectionUtils.combine(Collections.singletonList(query), newFilters)
+        );
+        query = query == null ? boolQuery() : query;
+        return source.query(query);
     }
 
     public static SearchSourceBuilder wrapAsFilter(SearchSourceBuilder source) {
@@ -252,4 +221,13 @@ public final class RuntimeUtils {
         source.query(bool);
         return source;
     }
+
+    public static QueryBuilder removeFilters(QueryBuilder query, List<QueryBuilder> filters) {
+        if (query instanceof BoolQueryBuilder boolQueryBuilder) {
+            if (org.elasticsearch.xpack.ql.util.CollectionUtils.isEmpty(filters) == false) {
+                boolQueryBuilder.filter().removeAll(filters);
+            }
+        }
+        return query;
+    }
 }

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/sequence/TumblingWindow.java

@@ -48,7 +48,7 @@ import java.util.Set;
 import static java.util.stream.Collectors.toList;
 import static org.elasticsearch.action.ActionListener.runAfter;
 import static org.elasticsearch.xpack.eql.execution.ExecutionUtils.copySource;
-import static org.elasticsearch.xpack.eql.execution.search.RuntimeUtils.addFilter;
+import static org.elasticsearch.xpack.eql.execution.search.RuntimeUtils.combineFilters;
 import static org.elasticsearch.xpack.eql.execution.search.RuntimeUtils.searchHits;
 import static org.elasticsearch.xpack.eql.util.SearchHitUtils.qualifiedIndex;
 
@@ -314,7 +314,7 @@ public class TumblingWindow implements Executable {
                         builder.sort(r.timestampField(), SortOrder.ASC);
                     }
                     addKeyFilter(i, sequence, builder);
-                    RuntimeUtils.addFilter(range, builder);
+                    RuntimeUtils.combineFilters(builder, range);
                     result.add(RuntimeUtils.prepareRequest(builder.size(1).trackTotalHits(false), false, Strings.EMPTY_ARRAY));
                 } else {
                     leading = false;
@@ -331,7 +331,7 @@ public class TumblingWindow implements Executable {
         }
         for (int i = 0; i < keys.size(); i++) {
             Attribute k = keys.get(i);
-            addFilter(new TermQueryBuilder(k.qualifiedName(), sequence.key().asList().get(i)), builder);
+            combineFilters(builder, new TermQueryBuilder(k.qualifiedName(), sequence.key().asList().get(i)));
         }
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java

@@ -46,6 +46,8 @@ import org.elasticsearch.xpack.ql.planner.QlTranslatorHandler;
 import org.elasticsearch.xpack.ql.querydsl.query.Query;
 import org.elasticsearch.xpack.ql.rule.ParameterizedRuleExecutor;
 import org.elasticsearch.xpack.ql.rule.Rule;
+import org.elasticsearch.xpack.ql.util.Queries;
+import org.elasticsearch.xpack.ql.util.Queries.Clause;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -56,7 +58,6 @@ import java.util.Set;
 import java.util.function.Supplier;
 
 import static java.util.Arrays.asList;
-import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
 import static org.elasticsearch.xpack.ql.expression.predicate.Predicates.splitAnd;
 import static org.elasticsearch.xpack.ql.optimizer.OptimizerRules.TransformDirection.UP;
 
@@ -191,11 +192,7 @@ public class LocalPhysicalPlanOptimizer extends ParameterizedRuleExecutor<Physic
                 }
                 if (pushable.size() > 0) { // update the executable with pushable conditions
                     QueryBuilder planQuery = TRANSLATOR_HANDLER.asQuery(Predicates.combineAnd(pushable)).asBuilder();
-                    QueryBuilder query = planQuery;
-                    QueryBuilder filterQuery = queryExec.query();
-                    if (filterQuery != null) {
-                        query = boolQuery().filter(filterQuery).filter(planQuery);
-                    }
+                    var query = Queries.combine(Clause.FILTER, asList(queryExec.query(), planQuery));
                     queryExec = new EsQueryExec(
                         queryExec.source(),
                         queryExec.index(),

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/stats/SearchStats.java

@@ -135,7 +135,7 @@ public class SearchStats {
     }
 
     //
-    // @see org.elasticsearch.search.query.TopDocsCollectorManagerFactory#shortcutTotalHitCount(IndexReader, Query)
+    // @see org.elasticsearch.search.query.QueryPhaseCollectorManager#shortcutTotalHitCount(IndexReader, Query)
     //
     private static int countEntries(IndexReader indexReader, String field) {
         int count = 0;

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/Queries.java

@@ -0,0 +1,77 @@
+/*
+ * 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.ql.util;
+
+import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+
+import java.util.List;
+import java.util.function.Function;
+
+import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
+
+/**
+ * Utilities for Elasticsearch queries.
+ */
+public class Queries {
+
+    public enum Clause {
+        FILTER(BoolQueryBuilder::filter),
+        MUST(BoolQueryBuilder::must),
+        MUST_NOT(BoolQueryBuilder::mustNot),
+        SHOULD(BoolQueryBuilder::should);
+
+        final Function<BoolQueryBuilder, List<QueryBuilder>> innerQueries;
+
+        Clause(Function<BoolQueryBuilder, List<QueryBuilder>> innerQueries) {
+            this.innerQueries = innerQueries;
+        }
+    }
+
+    /**
+     * Combines the given queries while attempting to NOT create a new bool query and avoid
+     * unnecessary nested queries.
+     * The method tries to detect if the first query is a bool query - if that is the case it will
+     * reuse that for adding the rest of the clauses.
+     */
+    public static QueryBuilder combine(Clause clause, List<QueryBuilder> queries) {
+        QueryBuilder firstQuery = null;
+        BoolQueryBuilder bool = null;
+
+        for (QueryBuilder query : queries) {
+            if (query == null) {
+                continue;
+            }
+            if (firstQuery == null) {
+                firstQuery = query;
+                if (firstQuery instanceof BoolQueryBuilder bqb) {
+                    bool = bqb;
+                }
+            }
+            // at least two entries, start copying
+            else {
+                // lazy init the root bool
+                if (bool == null) {
+                    bool = combine(clause, boolQuery(), firstQuery);
+                }
+                // keep adding queries to it
+                bool = combine(clause, bool, query);
+            }
+        }
+
+        return bool == null ? firstQuery : bool;
+    }
+
+    private static BoolQueryBuilder combine(Clause clause, BoolQueryBuilder bool, QueryBuilder query) {
+        var list = clause.innerQueries.apply(bool);
+        if (list.contains(query) == false) {
+            list.add(query);
+        }
+        return bool;
+    }
+}

x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/util/QueriesTests.java

@@ -0,0 +1,129 @@
+/*
+ * 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.ql.util;
+
+import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.test.ESTestCase;
+
+import static java.util.Arrays.asList;
+import static org.hamcrest.Matchers.everyItem;
+import static org.hamcrest.Matchers.in;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.sameInstance;
+
+public class QueriesTests extends ESTestCase {
+
+    private static QueryBuilder randomNonBoolQuery() {
+        return randomFrom(
+            random(),
+            QueryBuilders::matchAllQuery,
+            QueryBuilders::idsQuery,
+            () -> QueryBuilders.rangeQuery(randomRealisticUnicodeOfLength(5)),
+            () -> QueryBuilders.termQuery(randomAlphaOfLength(5), randomAlphaOfLength(5)),
+            () -> QueryBuilders.existsQuery(randomAlphaOfLength(5)),
+            () -> QueryBuilders.geoBoundingBoxQuery(randomAlphaOfLength(5))
+        );
+    }
+
+    private static BoolQueryBuilder randomBoolQuery() {
+        var bool = QueryBuilders.boolQuery();
+        if (randomBoolean()) {
+            bool.filter(randomNonBoolQuery());
+        }
+        if (randomBoolean()) {
+            bool.must(randomNonBoolQuery());
+        }
+        if (randomBoolean()) {
+            bool.mustNot(randomNonBoolQuery());
+        }
+        if (randomBoolean()) {
+            bool.should(randomNonBoolQuery());
+        }
+        return bool;
+    }
+
+    public void testCombineNotCreatingBool() {
+        var clause = randomFrom(Queries.Clause.values());
+        var nonBool = randomNonBoolQuery();
+        assertThat(nonBool, sameInstance(Queries.combine(clause, asList(null, null, nonBool, null))));
+    }
+
+    public void testCombineNonBoolQueries() {
+        var queries = randomArray(2, 10, QueryBuilder[]::new, QueriesTests::randomNonBoolQuery);
+
+        var clause = randomFrom(Queries.Clause.values());
+        var list = asList(queries);
+        var combination = Queries.combine(clause, list);
+
+        assertThat(combination, instanceOf(BoolQueryBuilder.class));
+        var bool = (BoolQueryBuilder) combination;
+        var clauseList = clause.innerQueries.apply(bool);
+        assertThat(list, everyItem(in(clauseList)));
+    }
+
+    public void testCombineBoolQueries() {
+        var queries = randomArray(2, 10, QueryBuilder[]::new, () -> {
+            var bool = QueryBuilders.boolQuery();
+            if (randomBoolean()) {
+                bool.filter(randomNonBoolQuery());
+            }
+            if (randomBoolean()) {
+                bool.must(randomNonBoolQuery());
+            }
+            if (randomBoolean()) {
+                bool.mustNot(randomNonBoolQuery());
+            }
+            if (randomBoolean()) {
+                bool.should(randomNonBoolQuery());
+            }
+            return bool;
+        });
+
+        var clause = randomFrom(Queries.Clause.values());
+        var list = asList(queries);
+        var combination = Queries.combine(clause, list);
+
+        assertThat(combination, instanceOf(BoolQueryBuilder.class));
+        var bool = (BoolQueryBuilder) combination;
+
+        var clauseList = clause.innerQueries.apply(bool);
+
+        for (QueryBuilder query : queries) {
+            if (query != bool) {
+                assertThat(query, in(clauseList));
+            }
+        }
+    }
+
+    public void testCombineMixedBoolAndNonBoolQueries() {
+        var queries = randomArray(2, 10, QueryBuilder[]::new, () -> {
+            if (randomBoolean()) {
+                return QueriesTests.randomBoolQuery();
+            } else {
+                return QueriesTests.randomNonBoolQuery();
+            }
+        });
+
+        var clause = randomFrom(Queries.Clause.values());
+        var list = asList(queries);
+        var combination = Queries.combine(clause, list);
+
+        assertThat(combination, instanceOf(BoolQueryBuilder.class));
+        var bool = (BoolQueryBuilder) combination;
+
+        var clauseList = clause.innerQueries.apply(bool);
+
+        for (QueryBuilder query : queries) {
+            if (query != bool) {
+                assertThat(query, in(clauseList));
+            }
+        }
+    }
+}