Fix FORK with union-types (elastic#134033)

---------

Co-authored-by: ncordon <[email protected]>
Co-authored-by: Nacho Cordón <[email protected]>

Author: craigtaverner

docs/changelog/134033.yaml

@@ -0,0 +1,6 @@
+pr: 134033
+summary: Fix FORK with union-types
+area: ES|QL
+type: bug
+issues:
+ - 133973

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeForkRestTest.java

@@ -55,11 +55,6 @@ protected void shouldSkipTest(String testName) throws IOException {
             testCase.requiredCapabilities.contains(UNMAPPED_FIELDS.capabilityName())
         );
 
-        assumeFalse(
-            "Tests using implicit_casting_date_and_date_nanos are not supported for now",
-            testCase.requiredCapabilities.contains(IMPLICIT_CASTING_DATE_AND_DATE_NANOS.capabilityName())
-        );
-
         assumeTrue("Cluster needs to support FORK", hasCapabilities(adminClient(), List.of(FORK_V9.capabilityName())));
     }
 }

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

@@ -1397,6 +1397,23 @@ event_duration:long | _index:keyword       | ts:date                   | ts_str:
 8268153             | sample_data_ts_nanos | 2023-10-23T13:52:55.015Z  | 2023-10-23T13:52:55.015123456Z  | 1698069175015123456 |  172.21.3.15   |  172.21.3.15
 ;
 
+multiIndexInlinestatsOfMultiTypedField
+required_capability: inlinestats_v11
+required_capability: fork_union_types
+// https://github.com/elastic/elasticsearch/issues/133973
+FROM apps, apps_short
+| INLINESTATS s = sum(id::integer)
+| SORT name, version
+| LIMIT 5
+;
+
+id:unsupported | name:keyword | version:version | s:long       
+null           | aaaaa        | 1               | 210
+null           | aaaaa        | 1               | 210
+null           | aaaaa        | 1.2.3.4         | 210
+null           | aaaaa        | 1.2.3.4         | 210
+null           | bbbbb        | 2.1             | 210
+;
 
 multiIndexIndirectUseOfUnionTypesInSort
 required_capability: union_types
@@ -1652,6 +1669,35 @@ id:integer | name:keyword | count:long
 14         | mmmmm        | 2
 ;
 
+shortIntegerWideningFork
+required_capability: union_types
+required_capability: fork_union_types
+required_capability: casting_operator
+required_capability: union_types_numeric_widening
+required_capability: fork_v9
+
+FROM apps, apps_short
+| EVAL x = id::integer
+| FORK (WHERE true) (WHERE true)
+| DROP _fork
+| SORT x ASC, name ASC, version ASC
+| EVAL x = x + 1
+| LIMIT 10
+;
+
+id:unsupported | name:keyword | version:version | x:integer
+null           | aaaaa        | 1               | 2
+null           | aaaaa        | 1               | 2
+null           | aaaaa        | 1               | 2
+null           | aaaaa        | 1               | 2
+null           | bbbbb        | 2.1             | 3
+null           | bbbbb        | 2.1             | 3
+null           | bbbbb        | 2.1             | 3
+null           | bbbbb        | 2.1             | 3
+null           | ccccc        | 2.3.4           | 4
+null           | ccccc        | 2.3.4           | 4
+;
+
 ImplicitCastingMultiTypedFieldsKeepSort
 required_capability: date_nanos_type
 required_capability: implicit_casting_date_and_date_nanos

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

@@ -1100,6 +1100,11 @@ public enum Cap {
          */
         FORK_V9,
 
+        /**
+         * Support for union types in FORK
+         */
+        FORK_UNION_TYPES,
+
         /**
          * Support for the {@code leading_zeros} named parameter.
          */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -809,7 +809,9 @@ private LogicalPlan resolveFork(Fork fork, AnalyzerContext context) {
             // We don't want to keep the same attributes that are outputted by the FORK branches.
             // Keeping the same attributes can have unintended side effects when applying optimizations like constant folding.
             for (Attribute attr : outputUnion) {
-                newOutput.add(new ReferenceAttribute(attr.source(), null, attr.name(), attr.dataType()));
+                newOutput.add(
+                    new ReferenceAttribute(attr.source(), null, attr.name(), attr.dataType(), Nullability.FALSE, null, attr.synthetic())
+                );
             }
 
             return changed ? new Fork(fork.source(), newSubPlans, newOutput) : fork;
@@ -1946,10 +1948,9 @@ static Attribute checkUnresolved(FieldAttribute fa) {
         private static LogicalPlan planWithoutSyntheticAttributes(LogicalPlan plan) {
             List<Attribute> output = plan.output();
             List<Attribute> newOutput = new ArrayList<>(output.size());
-
             for (Attribute attr : output) {
                 // Do not let the synthetic union type field attributes end up in the final output.
-                if (attr.synthetic() && attr instanceof FieldAttribute) {
+                if (attr.synthetic() && attr != NO_FIELDS.getFirst()) {
                     continue;
                 }
                 newOutput.add(attr);

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

@@ -29,6 +29,8 @@
 import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
+import static org.elasticsearch.xpack.esql.analysis.Analyzer.NO_FIELDS;
+
 /**
  * A Fork is a n-ary {@code Plan} where each child is a sub plan, e.g.
  * {@code FORK [WHERE content:"fox" ] [WHERE content:"dog"] }
@@ -121,7 +123,7 @@ public static List<Attribute> outputUnion(List<LogicalPlan> subplans) {
                     continue;
                 }
 
-                if (names.contains(attr.name()) == false && attr.name().equals(Analyzer.NO_FIELDS_NAME) == false) {
+                if (names.contains(attr.name()) == false && attr != NO_FIELDS.getFirst()) {
                     names.add(attr.name());
                     output.add(attr);
                 }