Repository: elastic/elasticsearch
Pull request #127225: ESQL: Fix count optimization with pushable union types
Author: login: alex-spies
Labels: name: >bug
name: Team:Analytics
description: Meta label for analytical engine team (ESQL/Aggs/Geo)
name: auto-backport
description: Automatically create backport pull requests when merged
name: :Analytics/ES|QL
description: AKA ESQL
name: v8.18.1
name: v8.19.0
name: v9.0.1
name: v9.1.0
name: v8.17.6
Created: 2025-04-23T10:12:30Z | Merged: 2025-04-28T11:50:08Z

Fix https://github.com/elastic/elasticsearch/issues/127200



Diff:
diff --git a/docs/changelog/127225.yaml b/docs/changelog/127225.yaml
new file mode 100644
index 0000000000000..b161a5c40bfdf
--- /dev/null
+++ b/docs/changelog/127225.yaml
@@ -0,0 +1,6 @@
+pr: 127225
+summary: Fix count optimization with pushable union types
+area: ES|QL
+type: bug
+issues:
+ - 127200
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 8b19bc589fcff..aa08799943ade 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
@@ -393,6 +393,73 @@ count:long | @timestamp:date
          0 | 2023-10-23T13:50:00.000Z
 ;
 
+multiIndexIpStatsNonPushableCount
+// Could be pushed to Lucene if we knew whether ip fields are single valued or not.
+// May still be pushed down on multi-node environments if a node has only the index where client_ip is keyword.
+required_capability: union_types
+required_capability: fix_count_pushdown_for_union_types
+
+FROM sample_data, sample_data_str
+| STATS count=count(client_ip::ip)
+;
+
+count:long
+14
+;
+
+multiIndexIpStatsNonPushableCountEval
+// Could be pushed to Lucene if we knew whether ip fields are single valued or not.
+// May still be pushed down on multi-node environments if a node has only the index where client_ip is keyword.
+required_capability: union_types
+required_capability: fix_count_pushdown_for_union_types
+
+FROM sample_data, sample_data_str
+| EVAL client_ip = client_ip::ip
+| STATS count=count(client_ip)
+;
+
+count:long
+14
+;
+
+multiIndexIpStatsNonPushableCountWithFilter
+// Currently not pushable: has 2 aggs and we don't consider multi-typed fields aggregatable.
+required_capability: union_types
+required_capability: fix_count_pushdown_for_union_types
+
+FROM sample_data, sample_data_ts_long
+| STATS count_matching=count(@timestamp::long) WHERE @timestamp::long >= 1698069301543,
+        total_count=count(@timestamp::long)
+;
+
+count_matching:long | total_count:long
+2                   | 14
+;
+
+multiIndexIpStatsPushableCount
+required_capability: union_types
+required_capability: fix_count_pushdown_for_union_types
+
+FROM sample_data, sample_data_ts_long
+| STATS count=count(@timestamp::long)
+;
+
+count:long
+14
+;
+
+multiIndexIpStatsPushableCountEval
+required_capability: union_types
+required_capability: fix_count_pushdown_for_union_types
+
+FROM sample_data, sample_data_ts_long
+| EVAL @timestamp = @timestamp::long
+| STATS count=count(@timestamp)
+;
+
+count:long
+14
+;
 
 multiIndexIpStringStatsInline2
 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 3551ceeb3c068..281035a51359d 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
@@ -384,6 +384,12 @@ public enum Cap {
          */
         UNION_TYPES_AGG_CAST,
 
+        /**
+         * When pushing down {@code STATS count(field::type)} for a union type field, we wrongly used a synthetic attribute name in the
+         * query instead of the actual field name. This led to 0 counts instead of the correct result.
+         */
+        FIX_COUNT_PUSHDOWN_FOR_UNION_TYPES,
+
         /**
          * Fix to GROK validation in case of multiple fields with same name and different types
          * https://github.com/elastic/elasticsearch/issues/110533
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java
index 0d92c5830ff8b..5a99781b2afde 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java
@@ -66,12 +66,12 @@ public LogicalPlan apply(LogicalPlan plan, LocalLogicalOptimizerContext localLog
 
     private LogicalPlan missingToNull(LogicalPlan plan, Predicate<FieldAttribute> shouldBeRetained) {
         if (plan instanceof EsRelation relation) {
-            // Remove missing fields from the EsRelation because this is not where we will obtain them from; replace them by an Eval right
-            // after, instead. This allows us to safely re-use the attribute ids of the corresponding FieldAttributes.
+            // For any missing field, place an Eval right after the EsRelation to assign null values to that attribute (using the same name
+            // id!), thus avoiding that InsertFieldExtrations inserts a field extraction later.
             // This means that an EsRelation[field1, field2, field3] where field1 and field 3 are missing will be replaced by
             // Project[field1, field2, field3] <- keeps the ordering intact
             // \_Eval[field1 = null, field3 = null]
-            // \_EsRelation[field2]
+            // \_EsRelation[field1, field2, field3]
             List<Attribute> relationOutput = relation.output();
             Map<DataType, Alias> nullLiterals = Maps.newLinkedHashMapWithExpectedSize(DataType.types().size());
             List<NamedExpression> newProjections = new ArrayList<>(relationOutput.size());
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
index 3dab1e8b31eac..d7abc58bc7300 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
@@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
                                     if (context.searchStats().isSingleValue(fName)) {
-                                        fieldName = fa.name();
+                                        fieldName = fName;
                                         query = QueryBuilders.existsQuery(fieldName);
                                     }
                                 }
                             }
                             if (fieldName != null) {
                                 if (count.hasFilter()) {
+                                    // Note: currently, it seems like we never actually perform stats pushdown if we reach here.
+                                    // That's because stats pushdown only works for 1 agg function (without BY); but in that case, filters
+                                    // are extracted into a separate filter node upstream from the aggregation (and hopefully pushed into
+                                    // the EsQueryExec separately).
                                     if (canPushToSource(count.filter()) == false) {
                                         return null; // can't push down
                                     }


Reviews:
user: login: idegtiarenko
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T06:25:16Z
user: login: craigtaverner
body: Interesting! I did not know the bug was about Lucene pushdown of STATS.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T10:32:38Z
user: login: alex-spies
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T10:36:34Z
user: login: alex-spies
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T10:38:33Z
user: login: alex-spies
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T10:46:17Z
user: login: bpintea
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T11:53:30Z
user: login: alex-spies
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T08:50:59Z

Review comments:
diff_hunk: @@ -1038,7 +1038,13 @@ public enum Cap {
         /**
          * Support for the SAMPLE command
          */
-        SAMPLE(Build.current().isSnapshot());
+        SAMPLE(Build.current().isSnapshot()),
+
+        /**
+         * When pushing down {@code STATS count(field::type)} for a union type field, we wrongly used a synthetic attribute name in the
+         * query instead of the actual field name. This led to 0 counts instead of the correct result.
+         */
+        FIX_COUNT_PUSHDOWN_FOR_UNION_TYPES;
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
user: login: craigtaverner
body: We should always add new capabilities near similar capabilities, or at least avoid adding them at the end, because the end maximizes merge conflicts.
created_at: 2025-04-28T10:22:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -66,12 +66,12 @@ public LogicalPlan apply(LogicalPlan plan, LocalLogicalOptimizerContext localLog
 
     private LogicalPlan missingToNull(LogicalPlan plan, Predicate<FieldAttribute> shouldBeRetained) {
         if (plan instanceof EsRelation relation) {
-            // Remove missing fields from the EsRelation because this is not where we will obtain them from; replace them by an Eval right
-            // after, instead. This allows us to safely re-use the attribute ids of the corresponding FieldAttributes.
+            // For any missing field, place an Eval right after the EsRelation to assign null values to that attribute (using the same name
+            // id!), thus avoiding that InsertFieldExtrations inserts a field extraction later.
             // This means that an EsRelation[field1, field2, field3] where field1 and field 3 are missing will be replaced by
             // Project[field1, field2, field3] <- keeps the ordering intact
             // \_Eval[field1 = null, field3 = null]
-            // \_EsRelation[field2]
+            // \_EsRelation[field1, field2, field3]
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java
user: login: craigtaverner
body: If the field is missing, would it be in the EsRelation at all?
created_at: 2025-04-28T10:26:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
user: login: craigtaverner
body: Wow! Is this the fix? I imagine this could have impacts in many places, so could this fix other bugs we've not noticed?
created_at: 2025-04-28T10:27:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
                                     if (context.searchStats().isSingleValue(fName)) {
-                                        fieldName = fa.name();
+                                        fieldName = fName;
                                         query = QueryBuilders.existsQuery(fieldName);
                                     }
                                 }
                             }
                             if (fieldName != null) {
                                 if (count.hasFilter()) {
+                                    // Note: currently, it seems like we never actually perform stats pushdown if we reach here.
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
user: login: craigtaverner
body: It seems the question of what exactly gets pushed down and why is subtle. I wonder if we want some documentation on this?
created_at: 2025-04-28T10:31:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -66,12 +66,12 @@ public LogicalPlan apply(LogicalPlan plan, LocalLogicalOptimizerContext localLog
 
     private LogicalPlan missingToNull(LogicalPlan plan, Predicate<FieldAttribute> shouldBeRetained) {
         if (plan instanceof EsRelation relation) {
-            // Remove missing fields from the EsRelation because this is not where we will obtain them from; replace them by an Eval right
-            // after, instead. This allows us to safely re-use the attribute ids of the corresponding FieldAttributes.
+            // For any missing field, place an Eval right after the EsRelation to assign null values to that attribute (using the same name
+            // id!), thus avoiding that InsertFieldExtrations inserts a field extraction later.
             // This means that an EsRelation[field1, field2, field3] where field1 and field 3 are missing will be replaced by
             // Project[field1, field2, field3] <- keeps the ordering intact
             // \_Eval[field1 = null, field3 = null]
-            // \_EsRelation[field2]
+            // \_EsRelation[field1, field2, field3]
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java
user: login: alex-spies
body: Yes! Because we're in the local optimizer, the field _does_ exist overall, and thus is put into the EsRelation after the field caps call on the coordinator. But on the local node it's missing! (Or in the search stats that this optimization run uses, to be more precise) This optimizer rule applies exclusively to such fields.
created_at: 2025-04-28T10:36:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
user: login: alex-spies
body: Yep, this is the fix. Simple oversight to use the correct field name - in the past, the field _attribute_ name and the field name were the same, but union types had to break with this pattern.

I do not see other situations that this may fix, too, because the specific stats-pushdown optimization only triggers in a narrow slice of queries, anyway.
created_at: 2025-04-28T10:38:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
                                     if (context.searchStats().isSingleValue(fName)) {
-                                        fieldName = fa.name();
+                                        fieldName = fName;
                                         query = QueryBuilders.existsQuery(fieldName);
                                     }
                                 }
                             }
                             if (fieldName != null) {
                                 if (count.hasFilter()) {
+                                    // Note: currently, it seems like we never actually perform stats pushdown if we reach here.
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
user: login: alex-spies
body: When we touch this rule again, I think we should add unit tests that demonstrate exactly what is pushed down and how.

I also found that union type filters don't seem to be pushed to Lucene, yet - maybe we could improve the documentation as part of that work?

If you agree, I can open up an issue for this.
created_at: 2025-04-28T10:46:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -66,12 +66,12 @@ public LogicalPlan apply(LogicalPlan plan, LocalLogicalOptimizerContext localLog
 
     private LogicalPlan missingToNull(LogicalPlan plan, Predicate<FieldAttribute> shouldBeRetained) {
         if (plan instanceof EsRelation relation) {
-            // Remove missing fields from the EsRelation because this is not where we will obtain them from; replace them by an Eval right
-            // after, instead. This allows us to safely re-use the attribute ids of the corresponding FieldAttributes.
+            // For any missing field, place an Eval right after the EsRelation to assign null values to that attribute (using the same name
+            // id!), thus avoiding that InsertFieldExtrations inserts a field extraction later.
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceMissingFieldWithNull.java
user: login: bpintea
body: 👍 
created_at: 2025-04-28T10:55:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
                                     if (context.searchStats().isSingleValue(fName)) {
-                                        fieldName = fa.name();
+                                        fieldName = fName;
                                         query = QueryBuilders.existsQuery(fieldName);
                                     }
                                 }
                             }
                             if (fieldName != null) {
                                 if (count.hasFilter()) {
+                                    // Note: currently, it seems like we never actually perform stats pushdown if we reach here.
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
user: login: bpintea
body: > // Note: currently, it seems like we never actually perform stats pushdown if we reach here.

We evolved to this. Before the aggs filter were extracted into an upstream filter and pushed down, this code was "alive". 
created_at: 2025-04-28T11:14:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
user: login: bpintea
body: What I'm wondering is if we could make this mistake again and if we could prevent it.
Also, should we insert filters here using `#fieldName()`?
https://github.com/elastic/elasticsearch/blob/9e0a5af898b0d78c48fb22e758ae59fe5a1c8266/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/InferNonNullAggConstraint.java#L58 
created_at: 2025-04-28T11:53:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,15 +94,19 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
                                     if (context.searchStats().isSingleValue(fName)) {
-                                        fieldName = fa.name();
+                                        fieldName = fName;
                                         query = QueryBuilders.existsQuery(fieldName);
                                     }
                                 }
                             }
                             if (fieldName != null) {
                                 if (count.hasFilter()) {
+                                    // Note: currently, it seems like we never actually perform stats pushdown if we reach here.
path: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java
user: login: alex-spies
body: Yep yep, and if we manage to evolve the optimization to multiple stats, this code could come alive again.

There's an argument to be made that it should be deleted as long as it's dead - but properly double checking that this code path _really is never ever used atm_ was beyond the scope that I could allocate to this bug fix.

So leaving a comment was the next best thing for the time being, I think  :)
created_at: 2025-04-29T08:50:59Z
author_association: CONTRIBUTOR

PR comments:
user: login: elasticsearchmachine
created_at: 2025-04-23T10:13:27Z
body: Hi @alex-spies, I've created a changelog YAML for you.
author_association: COLLABORATOR
user: login: elasticsearchmachine
created_at: 2025-04-25T16:49:29Z
body: Pinging @elastic/es-analytical-engine (Team:Analytics)
author_association: COLLABORATOR
user: login: elasticsearchmachine
created_at: 2025-04-28T11:52:13Z
body: ## 💚 Backport successful
| Status | Branch | Result |
|:------:|:------:|:------:|
| ✅ |  [8.18](https://github.com/elastic/elasticsearch/pull/127459)  | [<img src="https://img.shields.io/github/pulls/detail/state/elastic/elasticsearch/127459">](https://github.com/elastic/elasticsearch/pull/127459) |
| ✅ |  [8.19](https://github.com/elastic/elasticsearch/pull/127460)  | [<img src="https://img.shields.io/github/pulls/detail/state/elastic/elasticsearch/127460">](https://github.com/elastic/elasticsearch/pull/127460) |
| ✅ |  [9.0](https://github.com/elastic/elasticsearch/pull/127461)  | [<img src="https://img.shields.io/github/pulls/detail/state/elastic/elasticsearch/127461">](https://github.com/elastic/elasticsearch/pull/127461) |
| ✅ |  [8.17](https://github.com/elastic/elasticsearch/pull/127462)  | [<img src="https://img.shields.io/github/pulls/detail/state/elastic/elasticsearch/127462">](https://github.com/elastic/elasticsearch/pull/127462) |
author_association: COLLABORATOR
user: login: bpintea
created_at: 2025-04-28T13:03:51Z
body: Oh, too late.

> Also, should we insert filters here using #fieldName()?

I will follow up on this.
author_association: CONTRIBUTOR
user: login: alex-spies
created_at: 2025-04-29T08:58:51Z
body: @bpintea , your comment https://github.com/elastic/elasticsearch/pull/127225#discussion_r2063503255 is an excellent find. Dang, we missed another usage of `FieldName#name` where we really meant `#fieldName`.

Thank you for following up on this! Maybe let's also add a javadoc comment to `NamedExpression#name`, or more specifically `FieldAttribute#name` to warn against using the attribute name (which can be synthetic).

To remove the sharp edge, we could consider refactoring methods using the _actual field name_ to use an `EsField` object (each `FieldAttribute` has one!) rather than a `String` to ensure, at compile time, that we never hand them a bogus name accidentally. More specifically, `SearchStats`' methods should maybe not use `String` everywhere, but maybe `EsField` or a dedicated `String` wrapper record that we can call `FielName` or so, just to leverage the compiler.
author_association: CONTRIBUTOR
user: login: alex-spies
created_at: 2025-04-29T17:06:59Z
body: > @bpintea , your comment [#127225 (comment)](https://github.com/elastic/elasticsearch/pull/127225#discussion_r2063503255) is an excellent find. Dang, we missed another usage of `FieldName#name` where we really meant `#fieldName`.
> 
> Thank you for following up on this! Maybe let's also add a javadoc comment to `NamedExpression#name`, or more specifically `FieldAttribute#name` to warn against using the attribute name (which can be synthetic).
> 
> To remove the sharp edge, we could consider refactoring methods using the _actual field name_ to use an `EsField` object (each `FieldAttribute` has one!) rather than a `String` to ensure, at compile time, that we never hand them a bogus name accidentally. More specifically, `SearchStats`' methods should maybe not use `String` everywhere, but maybe `EsField` or a dedicated `String` wrapper record that we can call `FielName` or so, just to leverage the compiler.

I think it's worth tracking this in an issue. Opened https://github.com/elastic/elasticsearch/issues/127521
author_association: CONTRIBUTOR