ESQL: Fix union types lost attributes in StubRelation for inlinestats

docs/changelog/135547.yaml

@@ -0,0 +1,5 @@
+pr: 135547
+summary: Fix union types lost attributes in `StubRelation` for inlinestats
+area: ES|QL
+type: bug
+issues: []

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

@@ -3498,86 +3498,121 @@ warningRegex:java.lang.IllegalArgumentException: milliSeconds \[-1457696696640\]
 2023-10-23T13:53:55.832Z      |2023-10-23T13:53:55.832987654Z|1698069235832987654|19
 ;
 
-ImplicitCastingMultiTypedDateTruncInlinestats_By-Ignore
-required_capability: inline_stats
-// https://github.com/elastic/elasticsearch/issues/133973
-// optimized incorrectly due to missing references [$$emp_no$converted_to$long{f$}#
+InlineStatsCountUnionType
+required_capability: inline_stats_with_union_types_in_stub_relation
+
+FROM employees, employees_incompatible
+| KEEP emp_no, hire_date 
+| INLINE STATS c = count(emp_no::long)
+| SORT hire_date DESC
+| LIMIT 10
+;
+
+emp_no:unsupported | hire_date:date_nanos     | c:long
+null               | 1999-04-30T00:00:00.000Z | 200
+null               | 1999-04-30T00:00:00.000Z | 200
+null               | 1997-05-19T00:00:00.000Z | 200
+null               | 1997-05-19T00:00:00.000Z | 200
+null               | 1996-11-05T00:00:00.000Z | 200
+null               | 1996-11-05T00:00:00.000Z | 200
+null               | 1995-12-15T00:00:00.000Z | 200
+null               | 1995-12-15T00:00:00.000Z | 200
+null               | 1995-08-22T00:00:00.000Z | 200
+null               | 1995-08-22T00:00:00.000Z | 200
+;
+
+ImplicitCastingMultiTypedDateTruncInlinestats_By
+required_capability: inline_stats_with_union_types_in_stub_relation
 
 FROM employees, employees_incompatible
 | KEEP emp_no, hire_date 
 | INLINE STATS c = count(emp_no::long) BY yr = DATE_TRUNC(1 year, hire_date)
-| SORT yr DESC
-| LIMIT 5
+| SORT yr DESC, hire_date DESC
+| LIMIT 10
 ;
 
-c:long | yr:date_nanos
-2      | 1999-01-01T00:00:00.000Z
-2      | 1997-01-01T00:00:00.000Z
-2      | 1996-01-01T00:00:00.000Z
-10     | 1995-01-01T00:00:00.000Z
-8      | 1994-01-01T00:00:00.000Z
+emp_no:unsupported | hire_date:date_nanos   |c:long | yr:date_nanos           
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
 ;
 
-ImplicitCastingMultiTypedDateTruncInlinestats_ByWithFilter-Ignore
-required_capability: inline_stats
-// https://github.com/elastic/elasticsearch/issues/133973
-// optimized incorrectly due to missing references [$$emp_no$converted_to$long{f$}#
+ImplicitCastingMultiTypedDateTruncInlinestats_ByWithFilter
+required_capability: inline_stats_with_union_types_in_stub_relation
 
 FROM employees, employees_incompatible
 | KEEP emp_no, hire_date 
 | INLINE STATS c = count(emp_no::long) where hire_date > "1996-01-01" BY yr = DATE_TRUNC(1 year, hire_date)
-| SORT yr DESC
-| LIMIT 5
+| SORT yr DESC, hire_date DESC
+| LIMIT 10
 ;
 
-c:long | yr:date_nanos
-2      | 1999-01-01T00:00:00.000Z
-2      | 1997-01-01T00:00:00.000Z
-2      | 1996-01-01T00:00:00.000Z
-0      | 1995-01-01T00:00:00.000Z
-0      | 1994-01-01T00:00:00.000Z
+emp_no:unsupported | hire_date:date_nanos   |c:long | yr:date_nanos           
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|0      |1995-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|0      |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|0      |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|0      |1995-01-01T00:00:00.000Z
 ;
 
-ImplicitCastingMultiTypedDateTruncInlinestats_ByWithEval-Ignore
-required_capability: inline_stats
-// https://github.com/elastic/elasticsearch/issues/133973
-// optimized incorrectly due to missing references [$$emp_no$converted_to$long{f$}#
+ImplicitCastingMultiTypedDateTruncInlinestats_ByWithEval
+required_capability: inline_stats_with_union_types_in_stub_relation
 
 FROM employees, employees_incompatible
 | KEEP emp_no, hire_date 
 | EVAL yr = DATE_TRUNC(1 year, hire_date)
 | INLINE STATS c = count(emp_no::long) BY yr
-| SORT yr DESC
-| LIMIT 5
+| SORT yr DESC, hire_date DESC
+| LIMIT 10
 ;
 
-c:long | yr:date_nanos
-2      | 1999-01-01T00:00:00.000Z
-2      | 1997-01-01T00:00:00.000Z
-2      | 1996-01-01T00:00:00.000Z
-10     | 1995-01-01T00:00:00.000Z
-8      | 1994-01-01T00:00:00.000Z
+emp_no:unsupported | hire_date:date_nanos   |c:long | yr:date_nanos           
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
 ;
 
-ImplicitCastingMultiTypedDateTruncInlinestats_ByWithEvalWithFilter-Ignore
-required_capability: inline_stats
-// https://github.com/elastic/elasticsearch/issues/133973
-// optimized incorrectly due to missing references [$$emp_no$converted_to$long{f$}#
+ImplicitCastingMultiTypedDateTruncInlinestats_ByWithEvalWithFilter
+required_capability: inline_stats_with_union_types_in_stub_relation
 
 FROM employees, employees_incompatible
 | KEEP emp_no, hire_date 
 | EVAL yr = DATE_TRUNC(1 year, hire_date)
 | INLINE STATS c = count(emp_no::long) where hire_date > "1991-01-01" BY yr
-| SORT yr DESC
-| LIMIT 5
+| SORT yr DESC, hire_date DESC
+| LIMIT 10
 ;
 
-c:long | yr:date_nanos
-2      | 1999-01-01T00:00:00.000Z
-2      | 1997-01-01T00:00:00.000Z
-2      | 1996-01-01T00:00:00.000Z
-10     | 1995-01-01T00:00:00.000Z
-8      | 1994-01-01T00:00:00.000Z
+emp_no:unsupported | hire_date:date_nanos   |c:long | yr:date_nanos           
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1999-04-30T00:00:00.000Z|2      |1999-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1997-05-19T00:00:00.000Z|2      |1997-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1996-11-05T00:00:00.000Z|2      |1996-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-12-15T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
+null               |1995-08-22T00:00:00.000Z|10     |1995-01-01T00:00:00.000Z
 ;
 
 ImplicitCastingMultiTypedBucketDateNanosByYear

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

@@ -1543,6 +1543,8 @@ public enum Cap {
         /** INLINE STATS supports remote indices */
         INLINE_STATS_SUPPORTS_REMOTE(INLINESTATS_V11.enabled),
 
+        INLINE_STATS_WITH_UNION_TYPES_IN_STUB_RELATION(INLINE_STATS.enabled),
+
         /**
          * Support TS command in non-snapshot builds
          */

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

@@ -21,6 +21,9 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin.replaceStub;
+import static org.elasticsearch.xpack.esql.plan.logical.join.StubRelation.computeOutput;
+
 /**
  * Replace any evaluation from the inlined aggregation side (right side) to the left side (source) to perform the matching.
  * In INLINE STATS m = MIN(x) BY a + b the right side contains STATS m = MIN(X) BY a + b.
@@ -83,8 +86,7 @@ protected LogicalPlan rule(InlineJoin plan) {
         if (groupingAlias.size() > 0) {
             left = new Eval(plan.source(), plan.left(), groupingAlias);
         }
-
         // replace the old stub with the new out to capture the new output
-        return plan.replaceChildren(left, InlineJoin.replaceStub(new StubRelation(right.source(), left.output()), right));
+        return plan.replaceChildren(left, replaceStub(new StubRelation(right.source(), computeOutput(right, left)), right));
     }
 }

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

@@ -55,7 +55,7 @@ public class InlineJoin extends Join {
      * Replaces the source of the target plan with a stub preserving the output of the source plan.
      */
     public static LogicalPlan stubSource(UnaryPlan sourcePlan, LogicalPlan target) {
-        return sourcePlan.replaceChild(new StubRelation(sourcePlan.source(), target.output()));
+        return sourcePlan.replaceChild(new StubRelation(sourcePlan.source(), StubRelation.computeOutput(sourcePlan, target)));
     }
 
     /**

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

@@ -18,8 +18,11 @@
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 
 import static java.util.Collections.emptyList;
 
@@ -36,15 +39,23 @@ public class StubRelation extends LeafPlan {
         StubRelation::new
     );
 
-    public static final StubRelation EMPTY = new StubRelation(null, null);
-
     private final List<Attribute> output;
 
     public StubRelation(Source source, List<Attribute> output) {
         super(source);
         this.output = output;
     }
 
+    /*
+     * The output of a StubRelation must also include any synthetic attributes referenced by the source plan (union types is a great
+     * example of those attributes that has some special treatment throughout the planning phases, especially in the EsRelation).
+     */
+    public static List<Attribute> computeOutput(LogicalPlan source, LogicalPlan target) {
+        Set<Attribute> stubRelationOutput = new LinkedHashSet<>(target.output());
+        stubRelationOutput.addAll(source.references().stream().filter(Attribute::synthetic).toList());
+        return new ArrayList<>(stubRelationOutput);
+    }
+
     public StubRelation(StreamInput in) throws IOException {
         this(Source.readFrom((PlanStreamInput) in), emptyList());
     }