From 4ddca132cf7ecdadc6be230919986a0fbdbeb133 Mon Sep 17 00:00:00 2001
From: Craig Taverner <craig@amanzi.com>
Date: Wed, 10 Jul 2024 10:35:54 +0200
Subject: [PATCH] Fix union-types when aggregating on inline conversion
 function (#110652)

A query like:

```
FROM sample_data, sample_data_str
| STATS count=count(*) BY client_ip = TO_IP(client_ip)
| SORT count DESC, client_ip ASC
| KEEP count, client_ip
```

Failed due to unresolved aggregates from the union-type in the grouping key
---
 .../src/main/resources/union_types.csv-spec       | 15 +++++++++++++++
 .../xpack/esql/action/EsqlCapabilities.java       |  7 ++++++-
 .../xpack/esql/analysis/Analyzer.java             | 13 ++++++++-----
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/union_types.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/union_types.csv-spec
index 5783489195458..349f968666132 100644
--- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/union_types.csv-spec
+++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/union_types.csv-spec
@@ -514,6 +514,21 @@ mc:l | count:l
 7    | 2
 ;
 
+multiIndexTsLongStatsStats
+required_capability: union_types
+required_capability: union_types_agg_cast
+
+FROM sample_data, sample_data_ts_long
+| EVAL ts = TO_STRING(@timestamp)
+| STATS count = COUNT(*) BY ts
+| STATS mc = COUNT(count) BY count
+| SORT mc DESC, count ASC
+;
+
+mc:l | count:l
+14   | 1
+;
+
 multiIndexTsLongRenameStats
 required_capability: union_types
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
index fa822b50ffcf5..5002a4a584954 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
@@ -117,7 +117,12 @@ public enum Cap {
          * Fix to GROK validation in case of multiple fields with same name and different types
          * https://github.com/elastic/elasticsearch/issues/110533
          */
-        GROK_VALIDATION;
+        GROK_VALIDATION,
+
+        /**
+         * Fix for union-types when aggregating over an inline conversion with conversion function. Done in #110652.
+         */
+        UNION_TYPES_INLINE_FIX;
 
         private final boolean snapshotOnly;
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java
index fbc98e093c0fb..add1f74cc3f04 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java
@@ -1088,7 +1088,11 @@ protected LogicalPlan doRule(LogicalPlan plan) {
 
             // In ResolveRefs the aggregates are resolved from the groupings, which might have an unresolved MultiTypeEsField.
             // Now that we have resolved those, we need to re-resolve the aggregates.
-            if (plan instanceof EsqlAggregate agg && agg.expressionsResolved() == false) {
+            if (plan instanceof EsqlAggregate agg) {
+                // If the union-types resolution occurred in a child of the aggregate, we need to check the groupings
+                plan = agg.transformExpressionsOnly(FieldAttribute.class, UnresolveUnionTypes::checkUnresolved);
+
+                // Aggregates where the grouping key comes from a union-type field need to be resolved against the grouping key
                 Map<Attribute, Expression> resolved = new HashMap<>();
                 for (Expression e : agg.groupings()) {
                     Attribute attr = Expressions.attribute(e);
@@ -1096,7 +1100,7 @@ protected LogicalPlan doRule(LogicalPlan plan) {
                         resolved.put(attr, e);
                     }
                 }
-                plan = agg.transformExpressionsOnly(UnresolvedAttribute.class, ua -> resolveAttribute(ua, resolved));
+                plan = plan.transformExpressionsOnly(UnresolvedAttribute.class, ua -> resolveAttribute(ua, resolved));
             }
 
             // Otherwise drop the converted attributes after the alias function, as they are only needed for this function, and
@@ -1222,9 +1226,8 @@ protected LogicalPlan rule(LogicalPlan plan) {
             return plan.transformExpressionsOnly(FieldAttribute.class, UnresolveUnionTypes::checkUnresolved);
         }
 
-        private static Attribute checkUnresolved(FieldAttribute fa) {
-            var field = fa.field();
-            if (field instanceof InvalidMappedField imf) {
+        static Attribute checkUnresolved(FieldAttribute fa) {
+            if (fa.field() instanceof InvalidMappedField imf) {
                 String unresolvedMessage = "Cannot use field [" + fa.name() + "] due to ambiguities being " + imf.errorMessage();
                 return new UnresolvedAttribute(fa.source(), fa.name(), fa.qualifier(), fa.id(), unresolvedMessage, null);
             }