ESQL: More care when pushing expression to load (elastic#139412)

With this change we only push if the EVAL/WHERE/STATS has exactly one
"leaf" node ancestor *AND* it's valid for pushing.

This prevents pushing into "StubRelations" like we get with this plan:
```
FROM k8s-downsampled
| INLINE STATS tx_max = MAX(network.eth0.tx) BY pod
```

In the future we should be able to push if we can push into *all*
relations rather than checking that there is just one. But baby steps.

Closes elastic#138620

Author: nik9000

muted-tests.yml

@@ -391,12 +391,6 @@ tests:
 - class: org.elasticsearch.xpack.ilm.TimeSeriesLifecycleActionsIT
   method: testWaitForSnapshot
   issue: https://github.com/elastic/elasticsearch/issues/138669
-- class: org.elasticsearch.xpack.esql.optimizer.LocalLogicalPlanOptimizerTests
-  method: testAggregateMetricDoubleInlineStats
-  issue: https://github.com/elastic/elasticsearch/issues/138620
-- class: org.elasticsearch.xpack.esql.optimizer.rules.logical.local.ReplaceDateTruncBucketWithRoundToTests
-  method: testAggregateMetricDoubleInlineStats
-  issue: https://github.com/elastic/elasticsearch/issues/138620
 - class: org.elasticsearch.smoketest.SmokeTestMultiNodeClientYamlTestSuiteIT
   method: test {yaml=index/100_field_name_length_limit/Test field name length limit}
   issue: https://github.com/elastic/elasticsearch/issues/138768

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

@@ -24,11 +24,14 @@
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
+import org.elasticsearch.xpack.esql.plan.logical.Row;
+import org.elasticsearch.xpack.esql.plan.logical.join.StubRelation;
 import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
 import org.elasticsearch.xpack.esql.rule.ParameterizedRule;
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -74,25 +77,98 @@ public class PushExpressionsToFieldLoad extends ParameterizedRule<LogicalPlan, L
 
     @Override
     public LogicalPlan apply(LogicalPlan plan, LocalLogicalOptimizerContext context) {
-        Rule rule = new Rule(context, plan);
+        Rule rule = new Rule(context);
         return plan.transformDown(LogicalPlan.class, rule::doRule);
     }
 
+    /**
+     * Lazily scans the plan for "primaries". A "primary" here is an {@link EsRelation}
+     * we can push a field load into.
+     * <p>
+     *     Every node in the plan in the plan can be traced "down" to a leaf - the source
+     *     of all of its data. This rule can only push expressions into {@link EsRelation},
+     *     but there are lots of other kinds of leaf nodes like {@link StubRelation} and
+     *     {@link Row}. If a node has any of those unsupported ancestors then {@link #primariesFor}
+     *     will return an empty {@link List}. This is the signal the rest of the code uses
+     *     for "can't push".
+     * </p>
+     */
+    private class Primaries {
+        /**
+         * A map from each node to all of its "primaries".  The empty list is special here - it means that the node's
+         * parent doesn't support pushing.
+         * <p>
+         *     Note: The primary itself will be in the map, pointing to itself.
+         * </p>
+         */
+        private Map<LogicalPlan, List<EsRelation>> primaries = new IdentityHashMap<>();
+
+        /**
+         * Find "primaries" for a node. Returning the empty list is special here - it
+         * means that the node's ancestors contain a node to which we cannot push.
+         */
+        List<EsRelation> primariesFor(LogicalPlan plan) {
+            scanSubtree(plan);
+            return primaries.get(plan);
+        }
+
+        /**
+         * Recursively scan the tree under {@code plan}, visiting ancestors
+         * before children, and ignoring any trees we've scanned before.
+         */
+        private void scanSubtree(LogicalPlan plan) {
+            if (primaries.containsKey(plan)) {
+                return;
+            }
+            if (plan.children().isEmpty()) {
+                onLeaf(plan);
+            } else {
+                for (LogicalPlan child : plan.children()) {
+                    scanSubtree(child);
+                }
+                onInner(plan);
+            }
+        }
+
+        private void onLeaf(LogicalPlan plan) {
+            if (plan instanceof EsRelation rel) {
+                if (rel.indexMode() == IndexMode.LOOKUP) {
+                    primaries.put(plan, List.of());
+                } else {
+                    primaries.put(rel, List.of(rel));
+                }
+            } else {
+                primaries.put(plan, List.of());
+            }
+        }
+
+        private void onInner(LogicalPlan plan) {
+            List<EsRelation> result = new ArrayList<>(plan.children().size());
+            for (LogicalPlan child : plan.children()) {
+                List<EsRelation> childPrimaries = primaries.get(child);
+                assert childPrimaries != null : "scanned depth first " + child;
+                if (childPrimaries.isEmpty()) {
+                    log.trace("{} unsupported primaries {}", plan, child);
+                    primaries.put(plan, List.of());
+                    return;
+                }
+                result.addAll(childPrimaries);
+            }
+            log.trace("{} primaries {}", plan, result);
+            primaries.put(plan, result);
+        }
+    }
+
     private class Rule {
         private final Map<Attribute.IdIgnoringWrapper, Attribute> addedAttrs = new HashMap<>();
 
         private final LocalLogicalOptimizerContext context;
-        private final LogicalPlan plan;
+        private final Primaries primaries = new Primaries();
 
-        /**
-         * The primary indices, lazily initialized.
-         */
-        private List<EsRelation> primaries;
         private boolean addedNewAttribute = false;
 
-        private Rule(LocalLogicalOptimizerContext context, LogicalPlan plan) {
+        private Rule(LocalLogicalOptimizerContext context) {
             this.context = context;
-            this.plan = plan;
         }
 
         private LogicalPlan doRule(LogicalPlan plan) {
@@ -115,7 +191,7 @@ private LogicalPlan doRule(LogicalPlan plan) {
         private LogicalPlan transformPotentialInvocation(LogicalPlan plan) {
             LogicalPlan transformedPlan = plan.transformExpressionsOnly(Expression.class, e -> {
                 if (e instanceof BlockLoaderExpression ble) {
-                    return transformExpression(e, ble);
+                    return transformExpression(plan, e, ble);
                 }
                 return e;
             });
@@ -130,12 +206,16 @@ private LogicalPlan transformPotentialInvocation(LogicalPlan plan) {
             return new EsqlProject(Source.EMPTY, transformedPlan, transformedPlan.output());
         }
 
-        private Expression transformExpression(Expression e, BlockLoaderExpression ble) {
+        private Expression transformExpression(LogicalPlan nodeWithExpression, Expression e, BlockLoaderExpression ble) {
             BlockLoaderExpression.PushedBlockLoaderExpression fuse = ble.tryPushToFieldLoading(context.searchStats());
             if (fuse == null) {
                 return e;
             }
-            if (anyPrimaryContains(fuse.field()) == false) {
+            List<EsRelation> planPrimaries = primaries.primariesFor(nodeWithExpression);
+            log.trace("found primaries {} {}", nodeWithExpression, planPrimaries);
+            if (planPrimaries.size() != 1) {
+                // Empty list means that we can't push.
+                // >1 primary is currently unsupported, though we expect to support it later.
                 return e;
             }
             var preference = context.configuration().pragmas().fieldExtractPreference();
@@ -184,26 +264,5 @@ private Expression replaceFieldsForFieldTransformations(Expression e, BlockLoade
             addedAttrs.put(key, newFunctionAttr);
             return newFunctionAttr;
         }
-
-        private List<EsRelation> primaries() {
-            if (primaries == null) {
-                primaries = new ArrayList<>(2);
-                plan.forEachUp(EsRelation.class, r -> {
-                    if (r.indexMode() != IndexMode.LOOKUP) {
-                        primaries.add(r);
-                    }
-                });
-            }
-            return primaries;
-        }
-
-        private boolean anyPrimaryContains(FieldAttribute attr) {
-            for (EsRelation primary : primaries()) {
-                if (primary.outputSet().contains(attr)) {
-                    return true;
-                }
-            }
-            return false;
-        }
     }
 }

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

@@ -44,6 +44,7 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Sum;
 import org.elasticsearch.xpack.esql.expression.function.fulltext.SingleFieldFullTextFunction;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
+import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromAggregateMetricDouble;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.Length;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.StartsWith;
@@ -1438,8 +1439,11 @@ public void testAggregateMetricDoubleTSCommand() {
         assertTrue(esRelation.output().contains(sumfieldAttr));
     }
 
+    /**
+     * Proves that we do <strong>not</strong> push to a stub-relation even if it contains
+     * a field we could otherwise push. Stub relations don't have a "push" concept.
+     */
     public void testAggregateMetricDoubleInlineStats() {
-        // TODO: modify below when we handle fusing and StubRelations properly
         assumeTrue("requires push", EsqlCapabilities.Cap.VECTOR_SIMILARITY_FUNCTIONS_PUSHDOWN.isEnabled());
         String query = """
             FROM k8s-downsampled
@@ -1451,37 +1455,35 @@ public void testAggregateMetricDoubleInlineStats() {
 
         LogicalPlan plan = localPlan(plan(query, tsAnalyzer), new EsqlTestUtils.TestSearchStats());
 
-        // EsqlProject[[@timestamp{f}#15, cluster{f}#16, pod{f}#17, network.eth0.tx{f}#34, tx_max{r}#5]]
+        // EsqlProject[[@timestamp{f}#972, cluster{f}#973, pod{f}#974, network.eth0.tx{f}#991, tx_max{r}#962]]
         var project = as(plan, Project.class);
         assertThat(Expressions.names(project.projections()), contains("@timestamp", "cluster", "pod", "network.eth0.tx", "tx_max"));
-        // TopN[[Order[@timestamp{f}#15,ASC,LAST], Order[cluster{f}#16,ASC,LAST], Order[pod{f}#17,ASC,LAST]],9[INTEGER],false]
+        // TopN[[Order[@timestamp{f}#972,ASC,LAST], Order[cluster{f}#973,ASC,LAST], Order[pod{f}#974,ASC,LAST]],9[INTEGER],false]
         var topN = as(project.child(), TopN.class);
-        // InlineJoin[LEFT,[pod{f}#17],[pod{r}#17]]
+        // InlineJoin[LEFT,[pod{f}#974],[pod{r}#974]]
         var inlineJoin = as(topN.child(), InlineJoin.class);
-        // Aggregate[[pod{f}#17],[MAX($$MAX(network.eth>$MAX$0{r$}#39,true[BOOLEAN],PT0S[TIME_DURATION]) AS tx_max#5, pod{f}#17]]
+        // Aggregate[[pod{f}#974],[MAX($$MAX(network.eth>$MAX$0{r$}#996,true[BOOLEAN],PT0S[TIME_DURATION]) AS tx_max#962, pod{f}#974
         var aggregate = as(inlineJoin.right(), Aggregate.class);
         assertThat(aggregate.groupings(), hasSize(1));
         assertThat(aggregate.aggregates(), hasSize(2));
         as(Alias.unwrap(aggregate.aggregates().get(0)), Max.class);
-        // Eval[[$$network.eth0.tx$AMD_MAX$1489455250{f$}#40 AS $$MAX(network.eth>$MAX$0#39]]
+
+        // Eval[[FROMAGGREGATEMETRICDOUBLE(network.eth0.tx{f}#991,1[INTEGER]) AS $$MAX(network.eth>$MAX$0#996]]
         var eval = as(aggregate.child(), Eval.class);
         assertThat(eval.fields(), hasSize(1));
-
         var alias = as(eval.fields().getFirst(), Alias.class);
-        var fieldAttr = as(alias.child(), FieldAttribute.class);
+        var load = as(alias.child(), FromAggregateMetricDouble.class); // <--- no pushing.
+        var fieldAttr = as(load.field(), FieldAttribute.class);
         assertThat(fieldAttr.fieldName().string(), equalTo("network.eth0.tx"));
-        var field = as(fieldAttr.field(), FunctionEsField.class);
-        var blockLoaderFunctionConfig = as(field.functionConfig(), BlockLoaderFunctionConfig.JustFunction.class);
-        assertThat(blockLoaderFunctionConfig.function(), equalTo(BlockLoaderFunctionConfig.Function.AMD_MAX));
+        as(fieldAttr.field(), EsField.class);
 
-        // TODO: modify this comment when unmuting test
-        // StubRelation[[@timestamp{f}#15, ..., cluster{f}#16, ..., network.eth0.tx{f}#34, ...,
-        // pod{f}#17, ..., $$MAX(network.eth>$MAX$0{r$}#39, $$network.eth0.tx$AMD_MAX$1489455250{f$}#40]]
+        // StubRelation[[@timestamp{f}#972, ... ]]
         var stubRelation = as(eval.child(), StubRelation.class);
-        assertFalse(stubRelation.output().contains(fieldAttr));
-        // EsRelation[k8s-downsampled][@timestamp{f}#15, client.ip{f}#19, cluster{f}#16, e..]
+        assertThat(stubRelation.output(), hasItem(fieldAttr));
+
+        // EsRelation[k8s-downsampled][@timestamp{f}#972, client.ip{f}#976, cluster{f}#973, ..]
         var esRelation = as(inlineJoin.left(), EsRelation.class);
-        assertTrue(esRelation.output().contains(fieldAttr));
+        assertThat(esRelation.output(), hasItem(fieldAttr));
     }
 
     public void testVectorFunctionsWhenFieldMissing() {