EsQueryExec refactor

Author: fang-xing-esql

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/PlanWritables.java

@@ -36,7 +36,6 @@
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.DissectExec;
 import org.elasticsearch.xpack.esql.plan.physical.EnrichExec;
-import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.EvalExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
@@ -105,7 +104,6 @@ public static List<NamedWriteableRegistry.Entry> physical() {
             CompletionExec.ENTRY,
             DissectExec.ENTRY,
             EnrichExec.ENTRY,
-            EsQueryExec.ENTRY,
             EsSourceExec.ENTRY,
             EvalExec.ENTRY,
             ExchangeExec.ENTRY,

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

@@ -7,8 +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;
 import org.elasticsearch.index.IndexMode;
@@ -28,32 +26,19 @@
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
 import org.elasticsearch.xpack.esql.expression.Order;
-import org.elasticsearch.xpack.esql.index.EsIndex;
-import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
-import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
-import static org.elasticsearch.TransportVersions.ESQL_SKIP_ES_INDEX_SERIALIZATION;
-
 public class EsQueryExec extends LeafExec implements EstimatesRowSize {
-    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
-        PhysicalPlan.class,
-        "EsQueryExec",
-        EsQueryExec::readFrom
-    );
-
     public static final EsField DOC_ID_FIELD = new EsField("_doc", DataType.DOC_DATA_TYPE, Map.of(), false);
-    public static final List<Sort> NO_SORTS = List.of();  // only exists to mimic older serialization, but we no longer serialize sorts
 
     private final String indexPattern;
     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;
 
@@ -128,21 +113,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 +132,6 @@ public EsQueryExec(
         IndexMode indexMode,
         Map<String, IndexMode> indexNameWithModes,
         List<Attribute> attrs,
-        QueryBuilder query,
         Expression limit,
         List<Sort> sorts,
         Integer estimatedRowSize,
@@ -162,69 +142,21 @@ 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;
     }
 
-    /**
-     * The matching constructor is used during physical plan optimization and needs valid sorts. But we no longer serialize sorts.
-     * If this cluster node is talking to an older instance it might receive a plan with sorts, but it will ignore them.
-     */
-    private static EsQueryExec readFrom(StreamInput in) throws IOException {
-        var source = Source.readFrom((PlanStreamInput) in);
-        String indexPattern;
-        Map<String, IndexMode> indexNameWithModes;
-        if (in.getTransportVersion().onOrAfter(ESQL_SKIP_ES_INDEX_SERIALIZATION)) {
-            indexPattern = in.readString();
-            indexNameWithModes = in.readMap(IndexMode::readFrom);
-        } else {
-            var index = EsIndex.readFrom(in);
-            indexPattern = index.name();
-            indexNameWithModes = index.indexNameWithModes();
-        }
-        var indexMode = EsRelation.readIndexMode(in);
-        var attrs = in.readNamedWriteableCollectionAsList(Attribute.class);
-        var query = in.readOptionalNamedWriteable(QueryBuilder.class);
-        var limit = in.readOptionalNamedWriteable(Expression.class);
-        in.readOptionalCollectionAsList(EsQueryExec::readSort);
-        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);
-    }
-
-    private static Sort readSort(StreamInput in) throws IOException {
-        return FieldSort.readFrom(in);
-    }
-
-    private static void writeSort(StreamOutput out, Sort sort) {
-        throw new IllegalStateException("sorts are no longer serialized");
-    }
-
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        Source.EMPTY.writeTo(out);
-        if (out.getTransportVersion().onOrAfter(ESQL_SKIP_ES_INDEX_SERIALIZATION)) {
-            out.writeString(indexPattern);
-            out.writeMap(indexNameWithModes, (o, v) -> IndexMode.writeTo(v, out));
-        } else {
-            new EsIndex(indexPattern, Map.of(), indexNameWithModes).writeTo(out);
-        }
-        EsRelation.writeIndexMode(out, indexMode());
-        out.writeNamedWriteableCollection(output());
-        out.writeOptionalNamedWriteable(query());
-        out.writeOptionalNamedWriteable(limit());
-        out.writeOptionalCollection(NO_SORTS, EsQueryExec::writeSort);
-        out.writeOptionalVInt(estimatedRowSize());
-        // queryBuilderAndTags is built on data node, serialization is not needed
+        throw new UnsupportedOperationException("not serialized");
     }
 
     @Override
     public String getWriteableName() {
-        return ENTRY.name;
+        throw new UnsupportedOperationException("not serialized");
     }
 
     public static boolean isSourceAttribute(Attribute attr) {
@@ -249,7 +181,6 @@ protected NodeInfo<EsQueryExec> info() {
             indexMode,
             indexNameWithModes,
             attrs,
-            query,
             limit,
             sorts,
             estimatedRowSize,
@@ -269,8 +200,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 +248,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 +260,6 @@ public EsQueryExec withLimit(Expression limit) {
                 indexMode,
                 indexNameWithModes,
                 attrs,
-                query,
                 limit,
                 sorts,
                 estimatedRowSize,
@@ -349,28 +284,32 @@ 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()))
             );
     }
 
@@ -383,9 +322,35 @@ public boolean canSubstituteRoundToWithQueryBuilderAndTags() {
         return indexMode != IndexMode.TIME_SERIES && (sorts == null || sorts.isEmpty());
     }
 
+    /**
+     * 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;
+    }
+
     @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 +368,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 +383,6 @@ public String nodeString() {
             + "indexMode["
             + indexMode
             + "], "
-            + "query["
-            + (query != null ? Strings.toString(query, false, true) : "")
-            + "]"
             + NodeUtils.limitedToString(attrs)
             + ", limit["
             + (limit != null ? limit.toString() : "")

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, if there are multiple queries or if the single query has tags,
+            // UnsupportedOperationException will be thrown by esQueryExec.query()
             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,17 @@ 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,
+                List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of()))
+            );
             if (aliases.isEmpty() == false) {
                 child = new EvalExec(Source.EMPTY, child, aliases);
             }