fix predicate push down

Author: fang-xing-esql

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

@@ -92,6 +92,49 @@ null           | null           | null              | 1        | 172.21.2.162
 null           | null           | null              | 4        | 172.21.3.15
 ;
 
+subqueryInFromWithStatsInSubqueryConjunctiveFilterInMainQuery
+required_capability: fork_v9
+required_capability: subquery_in_from_command
+
+FROM employees, (FROM sample_data
+                                | STATS cnt = count(*) by client_ip )
+                              , (FROM sample_data_str
+                                | STATS cnt = count(*) by client_ip )
+           metadata _index
+| EVAL client_ip = client_ip::ip
+| WHERE client_ip == "172.21.3.15" AND cnt >0
+| SORT emp_no, client_ip
+| KEEP _index,  emp_no, languages, cnt, client_ip
+;
+
+_index:keyword | emp_no:integer | languages:integer | cnt:long | client_ip:ip
+null           | null           | null              | 4        | 172.21.3.15
+null           | null           | null              | 4        | 172.21.3.15
+;
+
+subqueryInFromWithStatsInSubqueryDisjunctiveFilterInMainQuery
+required_capability: fork_v9
+required_capability: subquery_in_from_command
+
+FROM employees, (FROM sample_data
+                                | STATS cnt = count(*) by client_ip )
+                              , (FROM sample_data_str
+                                | STATS cnt = count(*) by client_ip )
+           metadata _index
+| EVAL client_ip = client_ip::ip
+| WHERE ( emp_no >= 10091 AND emp_no < 10094)  OR client_ip == "172.21.3.15"
+| SORT emp_no, client_ip
+| KEEP _index,  emp_no, languages, cnt, client_ip
+;
+
+_index:keyword | emp_no:integer | languages:integer | cnt:long | client_ip:ip
+employees      | 10091          | 3                 | null     | null
+employees      | 10092          | 1                 | null     | null
+employees      | 10093          | 3                 | null     | null
+null           | null           | null              | 4        | 172.21.3.15
+null           | null           | null              | 4        | 172.21.3.15
+;
+
 subqueryInFromWithStatsInMainQuery
 required_capability: fork_v9
 required_capability: subquery_in_from_command

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineFilters.java

@@ -329,24 +329,31 @@ private static LogicalPlan maybePushDownPastUnionAll(Filter filter, UnionAll uni
         if (pushable.isEmpty()) {
             return filter; // nothing to push down
         }
-        // Preserve the filter on top of UnionAll if not all pushable predicates can be pushed down into UnionAll children.
-        // This happens when the pushable predicate contains ReferenceAttribute that cannot be mapped to children's output correctly.
-        boolean preserveOriginalFilterOnTopOfUnionAll = false;
-        // 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
+        // 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 or a limit added by fork and
+        // then the real child, if there is unknown pattern, keep the filter and UnionAll plan unchanged
         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;
-                } else {
-                    preserveOriginalFilterOnTopOfUnionAll = true;
-                }
+            LogicalPlan newChild = switch (child) {
+                case Project project -> maybePushDownFilterPastProjectForUnionAllChild(pushable, project);
+                case Limit limit -> maybePushDownFilterPastLimitForUnionAllChild(pushable, limit);
+                default -> null; // TODO add a general push down for unexpected pattern
+            };
+
+            if (newChild == null) {
+                // Unexpected pattern, keep plan unchanged without pushing down filters
+                return filter;
+            }
+
+            if (newChild != child) {
+                changed = true;
                 newChildren.add(newChild);
-            } else { // unexpected pattern, just add the child as is
-                newChildren.add(child);
+            } else {
+                // Theoretically, all the pushable predicates should be pushed down into each child,
+                // in case one child is not changed, preserve the filter on top of UnionAll to make sure
+                // correct results are returned and avoid infinite loop of the rule.
+                return filter;
             }
         }
 
@@ -355,34 +362,16 @@ private static LogicalPlan maybePushDownPastUnionAll(Filter filter, UnionAll uni
         }
 
         LogicalPlan newUnionAll = unionAll.replaceChildren(newChildren);
-        if (preserveOriginalFilterOnTopOfUnionAll) {
-            // Preserve the filter on top of UnionAll as some pushable predicates cannot be pushed down
-            // to make sure correct results are returned
-            return filter.replaceChild(newUnionAll);
-        }
         if (nonPushable.isEmpty()) {
             return newUnionAll;
         } else {
             return filter.with(newUnionAll, Predicates.combineAnd(nonPushable));
         }
     }
 
-    private static LogicalPlan maybePushDownFilterPastEvalAndLimitForUnionAllChild(List<Expression> pushable, Project project) {
-        List<Expression> resolvedPushable = new ArrayList<>();
-        // Make sure the pushable predicates can find their corresponding attributes in the child project
-        for (Expression exp : pushable) {
-            Expression replaced = resolveUnionAllOutputByName(exp, project.projections());
-            if (replaced == null || replaced == exp) {
-                // cannot find the attribute in the child project, cannot push down this filter
-                return project;
-            } else {
-                resolvedPushable.add(replaced);
-            }
-        }
-        if (resolvedPushable.size() != pushable.size()) {
-            // Some pushable predicates cannot be resolved to the child project, cannot push down.
-            // This should not happen, however we need to be cautious here, if the predicate is removed from the main query,
-            // and it is not pushed down into the UnionAll child, the result will be incorrect.
+    private static LogicalPlan maybePushDownFilterPastProjectForUnionAllChild(List<Expression> pushable, Project project) {
+        List<Expression> resolvedPushable = resolvePushableAgainstOutput(pushable, project.projections());
+        if (resolvedPushable == null) {
             return project;
         }
         LogicalPlan child = project.child();
@@ -395,6 +384,35 @@ private static LogicalPlan maybePushDownFilterPastEvalAndLimitForUnionAllChild(L
         return project;
     }
 
+    private static LogicalPlan maybePushDownFilterPastLimitForUnionAllChild(List<Expression> pushable, Limit limit) {
+        List<Expression> resolvedPushable = resolvePushableAgainstOutput(pushable, limit.output());
+        if (resolvedPushable == null) {
+            return limit;
+        }
+        return pushDownFilterPastLimitForUnionAllChild(resolvedPushable, limit);
+    }
+
+    /**
+     * Attempts to resolve all pushable expressions against the given output attributes.
+     * Returns a fully resolved list if successful, or null if any expression cannot be resolved.
+     */
+    private static List<Expression> resolvePushableAgainstOutput(List<Expression> pushable, List<? extends NamedExpression> output) {
+        List<Expression> resolved = new ArrayList<>();
+        for (Expression exp : pushable) {
+            Expression replaced = resolveUnionAllOutputByName(exp, output);
+            // Make sure the pushable predicates can find their corresponding attributes in the output
+            if (replaced == null || replaced == exp) {
+                // cannot find the attribute in the child project, cannot push down this filter
+                return null;
+            }
+            resolved.add(replaced);
+        }
+        // If some pushable predicates cannot be resolved against the output, cannot push filter down.
+        // This should not happen, however we need to be cautious here, if the predicate is removed from
+        // the main query, and it is not pushed down into the UnionAll child, the result will be incorrect.
+        return resolved.size() == pushable.size() ? resolved : null;
+    }
+
     private static LogicalPlan pushDownFilterPastEvalForUnionAllChild(List<Expression> pushable, Project project, Eval eval) {
         // if the pushable references any attribute created by the eval, we cannot push down
         AttributeMap<Expression> evalAliases = buildEvaAliases(eval);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/UnionAll.java

@@ -106,6 +106,10 @@ public BiConsumer<LogicalPlan, Failures> postOptimizationPlanVerification() {
         return UnionAll::checkNestedUnionAlls;
     }
 
+    /**
+     * Defer the check for nested UnionAlls until after logical planner as some of the nested subqueries can be flattened
+     * by logical planner in the future.
+     */
     private static void checkNestedUnionAlls(LogicalPlan logicalPlan, Failures failures) {
         if (logicalPlan instanceof UnionAll unionAll) {
             unionAll.forEachDown(UnionAll.class, otherUnionAll -> {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/AbstractLogicalPlanOptimizerTests.java

@@ -203,7 +203,7 @@ public static void init() {
                 EsqlTestUtils.TEST_CFG,
                 new EsqlFunctionRegistry(),
                 getIndexResult,
-                emptyMap(),
+                defaultLookupResolution(),
                 enrichResolution,
                 emptyInferenceResolution(),
                 Map.of("test1", subqueryIndex1, "languages", subqueryIndex2)

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineFiltersTests.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
+import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
@@ -1231,18 +1232,38 @@ public void testPushDownFilterPastUnionAllAndCombineWithFilterInSubquery() {
     public void testPushDownFilterOnReferenceAttributesPastUnionAllDebug() {
         assumeTrue("Requires subquery in FROM command support", EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled());
         var plan = planSubquery("""
-            FROM test, (FROM test1 | where salary < 100000 | EVAL x = 1, y = emp_no, z = emp_no + 1)
+            FROM test
+                       , (FROM test1
+                          | where salary < 100000
+                          | EVAL x = 1, y = emp_no, z = emp_no + 1)
+                       , (FROM languages
+                          | STATS cnt = COUNT(*) by language_code
+                          | RENAME language_code AS z, cnt AS y
+                          | EVAL x = 1)
+                       , (FROM test1
+                          | RENAME languages AS language_code
+                          | LOOKUP JOIN languages_lookup ON language_code
+                          | RENAME emp_no AS x, salary AS y, language_code AS z)
             | WHERE x is not null and y is not null and z > 0
             """);
 
         Limit limit = as(plan, Limit.class);
         UnionAll unionAll = as(limit.child(), UnionAll.class);
-        assertEquals(2, unionAll.children().size());
+        assertEquals(4, unionAll.children().size());
 
         LocalRelation child1 = as(unionAll.children().get(0), LocalRelation.class);
 
         EsqlProject child2 = as(unionAll.children().get(1), EsqlProject.class);
-        Limit childLimit = as(child2.child(), Limit.class);
+        Filter filter = as(child2.child(), Filter.class);
+        IsNotNull isNotNull = as(filter.condition(), IsNotNull.class);
+        ReferenceAttribute y = as(isNotNull.field(), ReferenceAttribute.class);
+        assertEquals("y", y.name());
+        Eval eval = as(filter.child(), Eval.class);
+        List<Alias> aliases = eval.fields();
+        assertEquals(2, aliases.size());
+        assertEquals("language_name", aliases.get(0).name());
+        assertEquals("y", aliases.get(1).name());
+        Limit childLimit = as(eval.child(), Limit.class);
         Subquery subquery = as(childLimit.child(), Subquery.class);
         Project project = as(subquery.child(), Project.class);
         Filter childFilter = as(project.child(), Filter.class);
@@ -1251,8 +1272,8 @@ public void testPushDownFilterOnReferenceAttributesPastUnionAllDebug() {
         assertEquals("z", z.name());
         Literal right = as(greaterThan.right(), Literal.class);
         assertEquals(0, right.value());
-        Eval eval = as(childFilter.child(), Eval.class);
-        List<Alias> aliases = eval.fields();
+        eval = as(childFilter.child(), Eval.class);
+        aliases = eval.fields();
         assertEquals(2, aliases.size());
         Alias aliasX = aliases.get(0);
         assertEquals("x", aliasX.name());
@@ -1261,17 +1282,67 @@ public void testPushDownFilterOnReferenceAttributesPastUnionAllDebug() {
         Alias aliasZ = aliases.get(1);
         assertEquals("z", aliasZ.name());
         childFilter = as(eval.child(), Filter.class);
-        And and = as(childFilter.condition(), And.class);
-        IsNotNull isNotNull = as(and.right(), IsNotNull.class);
-        FieldAttribute emp_no = as(isNotNull.field(), FieldAttribute.class);
-        assertEquals("emp_no", emp_no.name());
-        LessThan lessThan = as(and.left(), LessThan.class);
+        LessThan lessThan = as(childFilter.condition(), LessThan.class);
         FieldAttribute salaryField = as(lessThan.left(), FieldAttribute.class);
         assertEquals("salary", salaryField.name());
         Literal literal = as(lessThan.right(), Literal.class);
         assertEquals(100000, literal.value());
         EsRelation relation = as(childFilter.child(), EsRelation.class);
         assertEquals("test1", relation.indexPattern());
+
+        EsqlProject child3 = as(unionAll.children().get(2), EsqlProject.class);
+        eval = as(child3.child(), Eval.class);
+        limit = as(eval.child(), Limit.class);
+        subquery = as(limit.child(), Subquery.class);
+        eval = as(subquery.child(), Eval.class);
+        filter = as(eval.child(), Filter.class);
+        And and = as(filter.condition(), And.class);
+        isNotNull = as(and.left(), IsNotNull.class);
+        y = as(isNotNull.field(), ReferenceAttribute.class);
+        assertEquals("y", y.name());
+        greaterThan = as(and.right(), GreaterThan.class);
+        z = as(greaterThan.left(), ReferenceAttribute.class);
+        assertEquals("z", z.name());
+        right = as(greaterThan.right(), Literal.class);
+        assertEquals(0, right.value());
+        Aggregate aggregate = as(filter.child(), Aggregate.class);
+        List<Expression> groupings = aggregate.groupings();
+        assertEquals(1, groupings.size());
+        FieldAttribute language_code = as(groupings.get(0), FieldAttribute.class);
+        assertEquals("language_code", language_code.name());
+        List<? extends NamedExpression> aggregates = aggregate.aggregates();
+        assertEquals(2, aggregates.size());
+        assertEquals("y", aggregates.get(0).name());
+        assertEquals("z", aggregates.get(1).name());
+        relation = as(aggregate.child(), EsRelation.class);
+        assertEquals("languages", relation.indexPattern());
+
+        EsqlProject child4 = as(unionAll.children().get(3), EsqlProject.class);
+        filter = as(child4.child(), Filter.class);
+        isNotNull = as(filter.condition(), IsNotNull.class);
+        ReferenceAttribute x = as(isNotNull.field(), ReferenceAttribute.class);
+        assertEquals("y", x.name());
+        eval = as(filter.child(), Eval.class);
+        aliases = eval.fields();
+        assertEquals(4, aliases.size());
+        limit = as(eval.child(), Limit.class);
+        subquery = as(limit.child(), Subquery.class);
+        project = as(subquery.child(), Project.class);
+        Join lookupJoin = as(project.child(), Join.class);
+        Filter leftFilter = as(lookupJoin.left(), Filter.class);
+        and = as(leftFilter.condition(), And.class);
+        isNotNull = as(and.left(), IsNotNull.class);
+        FieldAttribute emp_no = as(isNotNull.field(), FieldAttribute.class);
+        assertEquals("emp_no", emp_no.name());
+        greaterThan = as(and.right(), GreaterThan.class);
+        language_code = as(greaterThan.left(), FieldAttribute.class);
+        assertEquals("languages", language_code.name());
+        right = as(greaterThan.right(), Literal.class);
+        assertEquals(0, right.value());
+        relation = as(leftFilter.child(), EsRelation.class);
+        assertEquals("test1", relation.indexPattern());
+        relation = as(lookupJoin.right(), EsRelation.class);
+        assertEquals("languages_lookup", relation.indexPattern());
     }
 
     public void testPushDownFilterOnReferenceAttributesAndFieldAttributesPastUnionAllDebug() {