replace query with queryBuilderAndTags in EsQueryExec

Author: fang-xing-esql

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

@@ -81,7 +81,7 @@ protected static List<Batch<PhysicalPlan>> rules(boolean optimizeForEsSource) {
         }
         // execute the SubstituteRoundToWithQueryAndTags rule once after all the other pushdown rules are applied, as this rule generate
         // multiple QueryBuilders according the number of RoundTo points, it should be applied after all the other eligible pushdowns are
-        // done.
+        // done, and it should be executed only once.
         @SuppressWarnings("unchecked")
         var substituteRoundToWithQueryAndTags = new Batch<PhysicalPlan>(
             "Substitute RoundTo with QueryAndTags",

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushFiltersToSource.java

@@ -126,10 +126,10 @@ private static PhysicalPlan rewrite(
                 queryExec.indexMode(),
                 queryExec.indexNameWithModes(),
                 queryExec.output(),
-                query,
                 queryExec.limit(),
                 queryExec.sorts(),
-                queryExec.estimatedRowSize()
+                queryExec.estimatedRowSize(),
+                List.of(new EsQueryExec.QueryBuilderAndTags(query, List.of()))
             );
             // If the eval contains other aliases, not just field attributes, we need to keep them in the plan
             PhysicalPlan plan = evalFields.isEmpty() ? queryExec : new EvalExec(filterExec.source(), queryExec, evalFields);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTags.java

@@ -426,7 +426,6 @@ private static PhysicalPlan planRoundTo(RoundTo roundTo, EvalExec evalExec, EsQu
             queryExec.indexMode(),
             queryExec.indexNameWithModes(),
             newAttributes,
-            queryExec.query(),
             queryExec.limit(),
             queryExec.sorts(),
             queryExec.estimatedRowSize(),

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.plan.physical;
 
-import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -53,7 +52,6 @@ public class EsQueryExec extends LeafExec implements EstimatesRowSize {
     private final IndexMode indexMode;
     private final Map<String, IndexMode> indexNameWithModes;
     private final List<Attribute> attrs;
-    private final QueryBuilder query;
     private final Expression limit;
     private final List<Sort> sorts;
 
@@ -64,8 +62,12 @@ public class EsQueryExec extends LeafExec implements EstimatesRowSize {
     private final Integer estimatedRowSize;
 
     /**
-     * queryBuilderAndTags is build on data node from the {@code RoundTo} function by {@code PushRoundToToSource} rule.
-     * It will be used by {@code sourcePhysicalOperation} to create {@code LuceneSliceQueue.QueryAndTags}
+     * queryBuilderAndTags may contain one or multiple {@code QueryBuilder}s, it is built on data node.
+     * If there is one {@code RoundTo} function in the query plan, {@code ReplaceRoundToWithQueryAndTags} rule will build multiple
+     * {@code QueryBuilder}s in the list, otherwise it is expected to contain only one {@code QueryBuilder} and its tag is
+     * null.
+     * It will be used by {@code EsPhysicalOperationProviders}.{@code sourcePhysicalOperation} to create
+     * {@code LuceneSliceQueue.QueryAndTags}
      */
     private final List<QueryBuilderAndTags> queryBuilderAndTags;
 
@@ -88,7 +90,7 @@ public SortBuilder<?> sortBuilder() {
         }
 
         private static FieldSort readFrom(StreamInput in) throws IOException {
-            return new EsQueryExec.FieldSort(
+            return new FieldSort(
                 FieldAttribute.readFrom(in),
                 in.readEnum(Order.OrderDirection.class),
                 in.readEnum(Order.NullsPosition.class)
@@ -118,7 +120,12 @@ public FieldAttribute field() {
         }
     }
 
-    public record QueryBuilderAndTags(QueryBuilder query, List<Object> tags) {};
+    public record QueryBuilderAndTags(QueryBuilder query, List<Object> tags) {
+        @Override
+        public String toString() {
+            return "QueryBuilderAndTags{" + "queryBuilder=[" + query + "], tags=" + tags.toString() + "}";
+        }
+    };
 
     public EsQueryExec(
         Source source,
@@ -128,21 +135,17 @@ public EsQueryExec(
         List<Attribute> attributes,
         QueryBuilder query
     ) {
-        this(source, indexPattern, indexMode, indexNameWithModes, attributes, query, null, null, null, null);
-    }
-
-    public EsQueryExec(
-        Source source,
-        String indexPattern,
-        IndexMode indexMode,
-        Map<String, IndexMode> indexNameWithModes,
-        List<Attribute> attrs,
-        QueryBuilder query,
-        Expression limit,
-        List<Sort> sorts,
-        Integer estimatedRowSize
-    ) {
-        this(source, indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts, estimatedRowSize, null);
+        this(
+            source,
+            indexPattern,
+            indexMode,
+            indexNameWithModes,
+            attributes,
+            null,
+            null,
+            null,
+            List.of(new QueryBuilderAndTags(query, List.of()))
+        );
     }
 
     public EsQueryExec(
@@ -151,7 +154,6 @@ public EsQueryExec(
         IndexMode indexMode,
         Map<String, IndexMode> indexNameWithModes,
         List<Attribute> attrs,
-        QueryBuilder query,
         Expression limit,
         List<Sort> sorts,
         Integer estimatedRowSize,
@@ -162,10 +164,10 @@ public EsQueryExec(
         this.indexMode = indexMode;
         this.indexNameWithModes = indexNameWithModes;
         this.attrs = attrs;
-        this.query = query;
         this.limit = limit;
         this.sorts = sorts;
         this.estimatedRowSize = estimatedRowSize;
+        // cannot keep the ctor with QueryBuilder as it has the same number of arguments as this ctor, EsqlNodeSubclassTests will fail
         this.queryBuilderAndTags = queryBuilderAndTags;
     }
 
@@ -193,7 +195,17 @@ private static EsQueryExec readFrom(StreamInput in) throws IOException {
         var rowSize = in.readOptionalVInt();
         // Ignore sorts from the old serialization format
         // queryBuilderAndTags is built on data node, serialization is not needed
-        return new EsQueryExec(source, indexPattern, indexMode, indexNameWithModes, attrs, query, limit, NO_SORTS, rowSize, null);
+        return new EsQueryExec(
+            source,
+            indexPattern,
+            indexMode,
+            indexNameWithModes,
+            attrs,
+            limit,
+            NO_SORTS,
+            rowSize,
+            List.of(new QueryBuilderAndTags(query, List.of()))
+        );
     }
 
     private static Sort readSort(StreamInput in) throws IOException {
@@ -215,7 +227,7 @@ public void writeTo(StreamOutput out) throws IOException {
         }
         EsRelation.writeIndexMode(out, indexMode());
         out.writeNamedWriteableCollection(output());
-        out.writeOptionalNamedWriteable(query());
+        out.writeOptionalNamedWriteable(query()); // TODO Do we ever serialize EsQueryExec or query() ?
         out.writeOptionalNamedWriteable(limit());
         out.writeOptionalCollection(NO_SORTS, EsQueryExec::writeSort);
         out.writeOptionalVInt(estimatedRowSize());
@@ -249,7 +261,6 @@ protected NodeInfo<EsQueryExec> info() {
             indexMode,
             indexNameWithModes,
             attrs,
-            query,
             limit,
             sorts,
             estimatedRowSize,
@@ -269,8 +280,13 @@ public Map<String, IndexMode> indexNameWithModes() {
         return indexNameWithModes;
     }
 
+    /**
+     * query is merged into queryBuilderAndTags, keep this method as it is called by too many places.
+     * If this method is called, the caller looks for the original queryBuilder, before {@code ReplaceRoundToWithQueryAndTags} converts it
+     * to multiple queries with tags.
+     */
     public QueryBuilder query() {
-        return query;
+        return queryWithoutTag();
     }
 
     @Override
@@ -312,7 +328,7 @@ public PhysicalPlan estimateRowSize(State state) {
         }
         return Objects.equals(this.estimatedRowSize, size)
             ? this
-            : new EsQueryExec(source(), indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts, size, queryBuilderAndTags);
+            : new EsQueryExec(source(), indexPattern, indexMode, indexNameWithModes, attrs, limit, sorts, size, queryBuilderAndTags);
     }
 
     public EsQueryExec withLimit(Expression limit) {
@@ -324,7 +340,6 @@ public EsQueryExec withLimit(Expression limit) {
                 indexMode,
                 indexNameWithModes,
                 attrs,
-                query,
                 limit,
                 sorts,
                 estimatedRowSize,
@@ -349,43 +364,73 @@ public EsQueryExec withSorts(List<Sort> sorts) {
                 indexMode,
                 indexNameWithModes,
                 attrs,
-                query,
                 limit,
                 sorts,
                 estimatedRowSize,
                 queryBuilderAndTags
             );
     }
 
+    /**
+     * query is merged into queryBuilderAndTags, keep this method as it is called by too many places.
+     * If this method is called, the caller looks for the original queryBuilder, before {@code ReplaceRoundToWithQueryAndTags} converts it
+     * to multiple queries with tags.
+     */
     public EsQueryExec withQuery(QueryBuilder query) {
-        return Objects.equals(this.query, query)
+        QueryBuilder thisQuery = queryWithoutTag();
+        return Objects.equals(thisQuery, query)
             ? this
             : new EsQueryExec(
                 source(),
                 indexPattern,
                 indexMode,
                 indexNameWithModes,
                 attrs,
-                query,
                 limit,
                 sorts,
                 estimatedRowSize,
-                queryBuilderAndTags
+                List.of(new QueryBuilderAndTags(query, List.of()))
             );
     }
 
     public List<QueryBuilderAndTags> queryBuilderAndTags() {
         return queryBuilderAndTags;
     }
 
+    /**
+     * Returns the original queryBuilder before {@code ReplaceRoundToWithQueryAndTags} converts it to multiple queryBuilder with tags.
+     * If we reach here, the caller is looking for the original query before the rule converts it. If there are multiple queries in
+     * queryBuilderAndTags or if the single query in queryBuilderAndTags already has a tag, that means
+     * {@code ReplaceRoundToWithQueryAndTags} already applied to the original query, the original query cannot be retrieved any more,
+     * exception will be thrown.
+     */
+    private QueryBuilder queryWithoutTag() {
+        QueryBuilder queryWithoutTag;
+        if (queryBuilderAndTags == null || queryBuilderAndTags.isEmpty()) {
+            return null;
+        } else if (queryBuilderAndTags.size() == 1) {
+            QueryBuilderAndTags firstQuery = this.queryBuilderAndTags.get(0);
+            if (firstQuery.tags().isEmpty()) {
+                queryWithoutTag = firstQuery.query();
+            } else {
+                throw new UnsupportedOperationException("query is converted to query with tags: " + "[" + firstQuery + "]");
+            }
+        } else {
+            throw new UnsupportedOperationException(
+                "query is converted to multiple queries and tags: " + "[" + this.queryBuilderAndTags + "]"
+            );
+        }
+        return queryWithoutTag;
+    }
+
     public boolean canSubstituteRoundToWithQueryBuilderAndTags() {
         // TimeSeriesSourceOperator and LuceneTopNSourceOperator do not support QueryAndTags
         return indexMode != IndexMode.TIME_SERIES && (sorts == null || sorts.isEmpty());
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts, queryBuilderAndTags);
+        return Objects.hash(indexPattern, indexMode, indexNameWithModes, attrs, limit, sorts, queryBuilderAndTags);
     }
 
     @Override
@@ -403,7 +448,6 @@ public boolean equals(Object obj) {
             && Objects.equals(indexMode, other.indexMode)
             && Objects.equals(indexNameWithModes, other.indexNameWithModes)
             && Objects.equals(attrs, other.attrs)
-            && Objects.equals(query, other.query)
             && Objects.equals(limit, other.limit)
             && Objects.equals(sorts, other.sorts)
             && Objects.equals(estimatedRowSize, other.estimatedRowSize)
@@ -419,9 +463,6 @@ public String nodeString() {
             + "indexMode["
             + indexMode
             + "], "
-            + "query["
-            + (query != null ? Strings.toString(query, false, true) : "")
-            + "]"
             + NodeUtils.limitedToString(attrs)
             + ", limit["
             + (limit != null ? limit.toString() : "")
@@ -430,7 +471,7 @@ public String nodeString() {
             + "] estimatedRowSize["
             + estimatedRowSize
             + "] queryBuilderAndTags ["
-            + (queryBuilderAndTags != null ? queryBuilderAndTags.toString() : "")
+            + (queryBuilderAndTags != null ? queryBuilderAndTags : "")
             + "]";
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java

@@ -289,7 +289,8 @@ public final PhysicalOperation sourcePhysicalOperation(EsQueryExec esQueryExec,
             for (Sort sort : sorts) {
                 sortBuilders.add(sort.sortBuilder());
             }
-            // LuceneTopNSourceOperator does not support QueryAndTags
+            // LuceneTopNSourceOperator does not support QueryAndTags, check if there are multiple queries or if the single query has tags,
+            // UnsupportedOperationException will be thrown
             luceneFactory = new LuceneTopNSourceOperator.Factory(
                 shardContexts,
                 querySupplier(esQueryExec.query()),
@@ -303,9 +304,7 @@ public final PhysicalOperation sourcePhysicalOperation(EsQueryExec esQueryExec,
         } else {
             luceneFactory = new LuceneSourceOperator.Factory(
                 shardContexts,
-                esQueryExec.queryBuilderAndTags() != null && esQueryExec.queryBuilderAndTags().isEmpty() == false
-                    ? querySupplier(esQueryExec.queryBuilderAndTags())
-                    : querySupplier(esQueryExec.query()),
+                querySupplier(esQueryExec.queryBuilderAndTags()),
                 context.queryPragmas().dataPartitioning(physicalSettings.defaultDataPartitioning()),
                 context.queryPragmas().taskConcurrency(),
                 context.pageSize(rowEstimatedSize),

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushTopNToSourceTests.java

@@ -582,7 +582,7 @@ public TestPhysicalPlanBuilder limit(int limit) {
 
         public TopNExec build() {
             List<Attribute> attributes = new ArrayList<>(fields.values());
-            PhysicalPlan child = new EsQueryExec(Source.EMPTY, this.index, indexMode, Map.of(), attributes, null, null, List.of(), 0);
+            PhysicalPlan child = new EsQueryExec(Source.EMPTY, this.index, indexMode, Map.of(), attributes, null, List.of(), 0, null);
             if (aliases.isEmpty() == false) {
                 child = new EvalExec(Source.EMPTY, child, aliases);
             }