Squashed commit of fang's subquery branch as of Monday 29th September

Author: craigtaverner

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeForkRestTest.java

@@ -62,6 +62,11 @@ protected void shouldSkipTest(String testName) throws IOException {
             testCase.requiredCapabilities.contains(VIEWS_V1.capabilityName())
         );
 
+        assumeFalse(
+            "Tests using subqueries are skipped since we don't support nested subqueries",
+            testCase.requiredCapabilities.contains(SUBQUERY_IN_FROM_COMMAND.capabilityName())
+        );
+
         assumeTrue("Cluster needs to support FORK", hasCapabilities(adminClient(), List.of(FORK_V9.capabilityName())));
     }
 }

x-pack/plugin/esql/qa/testFixtures/src/main/resources/subquery.csv-spec

@@ -108,8 +108,21 @@ timeSeriesCommand
     : TS indexPatternAndMetadataFields
     ;
 
-indexPatternAndMetadataFields:
-    indexPattern (COMMA indexPattern)* metadata?
+indexPatternAndMetadataFields
+    : indexPatternOrSubquery (COMMA indexPatternOrSubquery)* metadata?
+    ;
+
+indexPatternOrSubquery
+    : indexPattern
+    | {this.isDevVersion()}? subquery
+    ;
+
+subquery
+    : LP fromCommand (PIPE subqueryProcessingCommand)* RP
+    ;
+
+subqueryProcessingCommand
+    : processingCommand
     ;
 
 indexPattern

x-pack/plugin/esql/src/main/antlr/EsqlBaseParser.g4

@@ -23,7 +23,11 @@ FROM_ASSIGN : ASSIGN -> type(ASSIGN);
 METADATA : 'metadata';
 
 // we need this for EXPLAIN
-FROM_RP : RP -> type(RP), popMode;
+// change to double popMode to accommodate subquerys in FROM, when see ')' pop out of subquery(default) mode and from mode
+FROM_RP : RP -> type(RP), popMode, popMode;
+
+// accommodate subQuery inside FROM
+FROM_LP : LP -> type(LP), pushMode(DEFAULT_MODE);
 
 // in 8.14 ` were not allowed
 // this has been relaxed in 8.15 since " is used for quoting

x-pack/plugin/esql/src/main/antlr/lexer/From.g4

@@ -1125,6 +1125,11 @@ public enum Cap {
          */
         VIEWS_V1(Build.current().isSnapshot()),
 
+        /**
+         * Support non-correlated subqueries in the FROM clause.
+         */
+        SUBQUERY_IN_FROM_COMMAND(Build.current().isSnapshot()),
+
         /**
          * Support for the {@code leading_zeros} named parameter.
          */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -151,6 +151,10 @@ public boolean includeCCSMetadata() {
         return includeCCSMetadata;
     }
 
+    public Predicate<String> skipOnFailurePredicate() {
+        return skipOnFailurePredicate;
+    }
+
     /**
      * Call when ES|QL "planning" phase is complete and query execution (in ComputeService) is about to start.
      * Note this is currently only built for a single phase planning/execution model. When INLINE STATS

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlExecutionInfo.java

@@ -20,7 +20,8 @@ public record AnalyzerContext(
     IndexResolution indexResolution,
     Map<String, IndexResolution> lookupResolution,
     EnrichResolution enrichResolution,
-    InferenceResolution inferenceResolution
+    InferenceResolution inferenceResolution,
+    Map<String, IndexResolution> subqueryResolution
 ) {
     // Currently for tests only, since most do not test lookups
     // TODO: make this even simpler, remove the enrichResolution for tests that do not require it (most tests)
@@ -31,6 +32,17 @@ public AnalyzerContext(
         EnrichResolution enrichResolution,
         InferenceResolution inferenceResolution
     ) {
-        this(configuration, functionRegistry, indexResolution, Map.of(), enrichResolution, inferenceResolution);
+        this(configuration, functionRegistry, indexResolution, Map.of(), enrichResolution, inferenceResolution, Map.of());
+    }
+
+    public AnalyzerContext(
+        Configuration configuration,
+        EsqlFunctionRegistry functionRegistry,
+        IndexResolution indexResolution,
+        Map<String, IndexResolution> lookupResolution,
+        EnrichResolution enrichResolution,
+        InferenceResolution inferenceResolution
+    ) {
+        this(configuration, functionRegistry, indexResolution, lookupResolution, enrichResolution, inferenceResolution, Map.of());
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.analysis;
 
 import org.elasticsearch.index.IndexMode;
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.core.util.Holder;
 import org.elasticsearch.xpack.esql.expression.function.UnresolvedFunction;
 import org.elasticsearch.xpack.esql.plan.IndexPattern;
@@ -16,7 +17,9 @@
 import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 /**
  * This class is part of the planner.  Acts somewhat like a linker, to find the indices and enrich policies referenced by the query.
@@ -29,9 +32,10 @@ public record PreAnalysis(
         List<Enrich> enriches,
         List<IndexPattern> lookupIndices,
         boolean supportsAggregateMetricDouble,
-        boolean supportsDenseVector
+        boolean supportsDenseVector,
+        Set<IndexPattern> subqueryIndices
     ) {
-        public static final PreAnalysis EMPTY = new PreAnalysis(null, null, List.of(), List.of(), false, false);
+        public static final PreAnalysis EMPTY = new PreAnalysis(null, null, List.of(), List.of(), false, false, Set.of());
     }
 
     public PreAnalysis preAnalyze(LogicalPlan plan) {
@@ -46,12 +50,18 @@ protected PreAnalysis doPreAnalyze(LogicalPlan plan) {
         Holder<IndexMode> indexMode = new Holder<>();
         Holder<IndexPattern> index = new Holder<>();
         List<IndexPattern> lookupIndices = new ArrayList<>();
+        Set<IndexPattern> subqueryIndices = new HashSet<>();
         plan.forEachUp(UnresolvedRelation.class, p -> {
             if (p.indexMode() == IndexMode.LOOKUP) {
                 lookupIndices.add(p.indexPattern());
             } else if (indexMode.get() == null || indexMode.get() == p.indexMode()) {
                 indexMode.set(p.indexMode());
-                index.set(p.indexPattern());
+                // the index pattern from main query is always the first to be seen
+                index.setIfAbsent(p.indexPattern());
+                // collect subquery index patterns
+                if (EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled()) {
+                    collectSubqueryIndexPattern(p, subqueryIndices, index.get());
+                }
             } else {
                 throw new IllegalStateException("index mode is already set");
             }
@@ -96,7 +106,32 @@ protected PreAnalysis doPreAnalyze(LogicalPlan plan) {
             unresolvedEnriches,
             lookupIndices,
             indexMode.get() == IndexMode.TIME_SERIES || supportsAggregateMetricDouble.get(),
-            supportsDenseVector.get()
+            supportsDenseVector.get(),
+            subqueryIndices
         );
     }
+
+    private void collectSubqueryIndexPattern(
+        UnresolvedRelation relation,
+        Set<IndexPattern> subqueryIndices,
+        IndexPattern mainIndexPattern
+    ) {
+        if (relation.preAnalyzed()) {
+            return;
+        }
+
+        IndexPattern pattern = relation.indexPattern();
+        boolean isLookup = relation.indexMode() == IndexMode.LOOKUP;
+        boolean isMainIndexPattern = pattern == mainIndexPattern;
+        /*if the subquery's index pattern is the same as the main query, it won't be added
+        * to the subquery indices set, if Analyzer doesn't find the subquery' indexResolution,
+        * it falls back to the main query's indexResolution
+         */
+        if (isLookup || isMainIndexPattern) {
+            return;
+        }
+        subqueryIndices.add(pattern);
+        System.out.println("collected subquery index pattern: " + pattern);
+        System.out.println("subquery indices now: " + subqueryIndices);
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/AnalyzerContext.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
+import org.elasticsearch.core.Tuple;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
@@ -15,18 +16,22 @@
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.util.CollectionUtils;
 import org.elasticsearch.xpack.esql.expression.predicate.Predicates;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
+import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.RegexExtract;
+import org.elasticsearch.xpack.esql.plan.logical.Subquery;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
+import org.elasticsearch.xpack.esql.plan.logical.UnionAll;
 import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
 import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
 import org.elasticsearch.xpack.esql.plan.logical.join.Join;
@@ -109,6 +114,12 @@ protected LogicalPlan rule(Filter filter, LogicalOptimizerContext ctx) {
             // See also https://github.com/elastic/elasticsearch/issues/127497
             // Push down past INLINE STATS if the condition is on the groupings
             return pushDownPastJoin(filter, join, ctx.foldCtx());
+        } else if (child instanceof UnionAll unionAll) {
+            // Push down filters that can be evaluated using only the output of the UnionAll
+            plan = maybePushDownPastUnionAll(filter, unionAll);
+        } else if (child instanceof Subquery subquery) {
+            // subquery is a placeholder, push down the filter to the child of the subquery
+            plan = subquery.replaceChild(new Filter(filter.source(), subquery.child(), filter.condition()));
         }
         // cannot push past a Limit, this could change the tailing result set returned
         return plan;
@@ -283,4 +294,160 @@ private static LogicalPlan maybePushDownPastUnary(
         }
         return plan;
     }
+
+    /* Push down filters that can be evaluated by the UnionAll child/leg to each child/leg,
+     * so that the filters can be pushed down further to the data source if possible.
+     * Filters that cannot be pushed down remain above the UnionAll.
+     *
+     * The children of a UnionAll/Fork plan has a similar pattern, as Fork adds EsqlProject,
+     * an optional Eval and Limit on top of its actual children.
+     * UnionAll
+     *   EsqlProject
+     *     Eval (optional)
+     *       Limit
+     *         EsRelation
+     *   EsqlProject
+     *     Eval (optional)
+     *       Limit
+     *         Subquery
+     *
+     * Push down the filter below limit when possible
+     */
+    private static LogicalPlan maybePushDownPastUnionAll(Filter filter, UnionAll unionAll) {
+        List<Expression> pushable = new ArrayList<>();
+        List<Expression> nonPushable = new ArrayList<>();
+        for (Expression exp : Predicates.splitAnd(filter.condition())) {
+            if (exp.references().subsetOf(unionAll.outputSet())) {
+                pushable.add(exp);
+            } else {
+                nonPushable.add(exp);
+            }
+        }
+        if (pushable.isEmpty()) {
+            return filter; // nothing to push down
+        }
+
+        // Push the filter down to each child of the UnionAll, the child of a UnionAll is always a project
+        // followed by an optional eval and then limit added by fork and then the real child
+        List<LogicalPlan> newChildren = new ArrayList<>();
+        boolean changed = false;
+        for (LogicalPlan child : unionAll.children()) {
+            if (child instanceof Project project) {
+                LogicalPlan newChild = maybePushDownFilterPastEvalAndLimitForUnionAllChild(pushable, project);
+                if (newChild != child) {
+                    changed = true;
+                }
+                newChildren.add(newChild);
+            } else { // unexpected pattern, just add the child as is
+                newChildren.add(child);
+            }
+        }
+
+        if (changed == false) { // nothing changed, return the original plan
+            return filter;
+        }
+
+        LogicalPlan newUnionAll = unionAll.replaceChildren(newChildren);
+        if (nonPushable.isEmpty()) {
+            return newUnionAll;
+        } else {
+            return filter.with(newUnionAll, Predicates.combineAnd(nonPushable));
+        }
+    }
+
+    private static LogicalPlan maybePushDownFilterPastEvalAndLimitForUnionAllChild(List<Expression> pushable, Project project) {
+        LogicalPlan child = project.child();
+        if (child instanceof Eval eval) {
+            return pushDownFilterPastEvalForUnionAllChild(pushable, project, eval);
+        } else if (child instanceof Limit limit) {
+            LogicalPlan newLimit = pushDownFilterPastLimitForUnionAllChild(pushable, limit);
+            return project.replaceChild(newLimit);
+        }
+        return project;
+    }
+
+    private static LogicalPlan pushDownFilterPastEvalForUnionAllChild(List<Expression> pushable, Project project, Eval eval) {
+        AttributeMap<Expression> evalAliases = buildEvaAliases(eval);
+
+        Tuple<List<Expression>, List<Expression>> pushablesAndNonPushables = splitPushableAndNonPushablePredicates(
+            pushable,
+            project.projections(),
+            exp -> exp.references().stream().anyMatch(evalAliases::containsKey)
+        );
+
+        LogicalPlan evalChild = eval.child();
+
+        if (evalChild instanceof Limit limit) {
+            LogicalPlan newLimit = pushDownFilterPastLimitForUnionAllChild(pushablesAndNonPushables.v1(), limit);
+            LogicalPlan newEval = eval.replaceChild(newLimit);
+            if (pushablesAndNonPushables.v2().isEmpty()) {
+                return project.replaceChild(newEval);
+            } else {
+                Filter newFilter = new Filter(project.source(), newEval, Predicates.combineAnd(pushablesAndNonPushables.v2()));
+                return project.replaceChild(newFilter);
+            }
+        }
+        return project;
+    }
+
+    private static LogicalPlan pushDownFilterPastLimitForUnionAllChild(List<Expression> pushable, Limit limit) {
+        // check whether the pushable expression's attribute needs to be replaced
+        Tuple<List<Expression>, List<Expression>> pushablesAndNonPushables = splitPushableAndNonPushablePredicates(
+            pushable,
+            limit.output(),
+            exp -> exp.references().subsetOf(limit.outputSet()) == false
+        );
+
+        if (pushablesAndNonPushables.v1().isEmpty() || pushablesAndNonPushables.v2().isEmpty() == false) {
+            return limit;
+        }
+        Expression combined = Predicates.combineAnd(pushablesAndNonPushables.v1());
+        Filter pushed = new Filter(limit.source(), limit.child(), combined);
+        return limit.replaceChild(pushed);
+    }
+
+    private static AttributeMap<Expression> buildEvaAliases(Eval eval) {
+        AttributeMap.Builder<Expression> builder = AttributeMap.builder();
+        for (Alias alias : eval.fields()) {
+            builder.put(alias.toAttribute(), alias.child());
+        }
+        return builder.build();
+    }
+
+    private static Tuple<List<Expression>, List<Expression>> splitPushableAndNonPushablePredicates(
+        List<Expression> predicates,
+        List<? extends NamedExpression> attributes,
+        Predicate<Expression> nonPushableCheck
+    ) {
+        List<Expression> pushable = new ArrayList<>();
+        List<Expression> nonPushable = new ArrayList<>();
+        for (Expression exp : predicates) {
+            Expression replaced = replaceAttributesByName(exp, attributes);
+            if (replaced == null || nonPushableCheck.test(replaced)) {
+                nonPushable.add(replaced);
+            } else {
+                pushable.add(replaced);
+            }
+        }
+        return Tuple.tuple(pushable, nonPushable);
+    }
+
+    private static Expression replaceAttributesByName(Expression expr, List<? extends NamedExpression> namedExpressions) {
+        // Collect all referenced attributes
+        for (Attribute attr : expr.references()) {
+            boolean found = namedExpressions.stream().anyMatch(ne -> ne.name().equals(attr.name()));
+            if (found == false) {
+                return null;
+            }
+        }
+        // Replace attributes by name
+        return expr.transformUp(Attribute.class, attr -> {
+            for (NamedExpression ne : namedExpressions) {
+                if (ne.name().equals(attr.name())) {
+                    return ne.toAttribute();
+                }
+            }
+            return attr;
+        });
+    }
 }