DRAFT: LOAD/NULLIFY golden tests

x-pack/plugin/esql/build.gradle

@@ -641,6 +641,7 @@ tasks.named("test").configure {
   if (System.getProperty("golden.overwrite") != null || project.hasProperty("golden.overwrite")) {
     systemProperty "golden.overwrite", "true"
   }
+
   systemProperty "policy.directory", file("${projectDir}").absolutePath
   systemProperty "java.security.policy", file("${projectDir}/test.policy").absolutePath
 }

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

@@ -371,6 +371,9 @@ protected void doTest() throws Throwable {
     }
 
     protected final void doTest(String query) throws Throwable {
+        if (query.toUpperCase().startsWith("SET") == false) {
+            query = "SET unmapped_fields=\"load\"; " + query;
+        }
         RequestObjectBuilder builder = new RequestObjectBuilder(randomFrom(XContentType.values()));
 
         if (query.toUpperCase(Locale.ROOT).contains("LOOKUP_\uD83D\uDC14")) {

x-pack/plugin/esql/qa/testFixtures/src/main/resources/data/all-types.csv

@@ -1 +1,6 @@
-alias_integer,boolean,byte,constant_keyword-foo,date,date_nanos,double,float,half_float,scaled_float,integer,ip,keyword,long,unsigned_long,short,text,version,wildcard,semantic_text,dense_vector
+boolean:boolean,byte:byte,constant_keyword-foo:keyword,date:date,date_nanos:date_nanos,double:double,float:float,half_float:half_float,scaled_float:scaled_float,integer:integer,ip:ip,keyword:keyword,long:long,unsigned_long:unsigned_long,short:short,text:text,version:version,wildcard:keyword
+true,1,foo,2024-01-01T00:00:00.000Z,2024-01-01T00:00:00.000000001Z,1.1,1.1,1.1,1.1,1,127.0.0.1,key1,100,100,10,text1,1.0.0,wild1
+false,2,foo,2024-01-02T00:00:00.000Z,2024-01-02T00:00:00.000000002Z,2.2,2.2,2.2,2.2,2,127.0.0.2,key2,200,200,20,text2,2.0.0,wild2
+true,3,foo,2024-01-03T00:00:00.000Z,2024-01-03T00:00:00.000000003Z,3.3,3.3,3.3,3.3,3,127.0.0.3,key3,300,300,30,text3,3.0.0,wild3
+false,4,foo,2024-01-04T00:00:00.000Z,2024-01-04T00:00:00.000000004Z,4.4,4.4,4.4,4.4,4,127.0.0.4,key4,400,400,40,text4,4.0.0,wild4
+true,5,foo,2024-01-05T00:00:00.000Z,2024-01-05T00:00:00.000000005Z,5.5,5.5,5.5,5.5,5,127.0.0.5,key5,500,500,50,text5,5.0.0,wild5

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

@@ -729,3 +729,22 @@ count(*):long | message::INT:integer
 13            | null
 1             | 42
 ;
+
+fieldAliasAndNonExistent
+required_capability: unmapped_fields
+required_capability: optional_fields
+required_capability: field_alias_support
+
+SET unmapped_fields="load"\;
+FROM all_types
+| KEEP integer, alias_integer, does_not_exist
+| SORT integer
+;
+
+integer:integer | alias_integer:integer | does_not_exist:keyword
+1               | 1                     | null
+2               | 2                     | null
+3               | 3                     | null
+4               | 4                     | null
+5               | 5                     | null
+;

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

@@ -12,6 +12,22 @@ null
 null
 ;
 
+partialMappingUnmappedFromSource
+required_capability: optional_fields_nullify_tech_preview
+
+SET unmapped_fields="nullify"\;
+FROM partial_mapping_sample_data
+| KEEP @timestamp, unmapped_message
+| SORT @timestamp DESC
+| LIMIT 3
+;
+
+@timestamp:date          | unmapped_message:null
+2024-10-23T13:55:01.543Z | null
+2024-10-23T13:53:55.832Z | null
+2024-10-23T13:52:55.015Z | null
+;
+
 keepStar
 required_capability: optional_fields_nullify_tech_preview
 
@@ -482,3 +498,93 @@ ROW x = 1
 x:integer      |does_not_exist:null |y:keyword | language_name:keyword
 1              |null                |null      |null
 ;
+
+partiallyUnmappedKeywordMultiIndex
+required_capability: optional_fields_nullify_tech_preview
+required_capability: index_metadata_field
+
+SET unmapped_fields="nullify"\;
+FROM sample_data, no_mapping_sample_data METADATA _index
+| KEEP _index, @timestamp, message
+| SORT _index, @timestamp DESC
+;
+
+_index:keyword         | @timestamp:date           | message:keyword
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+sample_data            | 2023-10-23T13:55:01.543Z  | Connected to 10.1.0.1
+sample_data            | 2023-10-23T13:53:55.832Z  | Connection error
+sample_data            | 2023-10-23T13:52:55.015Z  | Connection error
+sample_data            | 2023-10-23T13:51:54.732Z  | Connection error
+sample_data            | 2023-10-23T13:33:34.937Z  | Disconnected
+sample_data            | 2023-10-23T12:27:28.948Z  | Connected to 10.1.0.2
+sample_data            | 2023-10-23T12:15:03.360Z  | Connected to 10.1.0.3
+;
+
+partiallyUnmappedNonKeywordMultiIndexCast
+required_capability: optional_fields_nullify_tech_preview
+required_capability: index_metadata_field
+
+SET unmapped_fields="nullify"\;
+FROM sample_data, no_mapping_sample_data METADATA _index
+| EVAL duration = event_duration::DOUBLE
+| KEEP _index, @timestamp, duration
+| SORT _index, @timestamp DESC
+;
+
+_index:keyword         | @timestamp:date           | duration:double
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+sample_data            | 2023-10-23T13:55:01.543Z  | 1756467.0
+sample_data            | 2023-10-23T13:53:55.832Z  | 5033755.0
+sample_data            | 2023-10-23T13:52:55.015Z  | 8268153.0
+sample_data            | 2023-10-23T13:51:54.732Z  | 725448.0
+sample_data            | 2023-10-23T13:33:34.937Z  | 1232382.0
+sample_data            | 2023-10-23T12:27:28.948Z  | 2764889.0
+sample_data            | 2023-10-23T12:15:03.360Z  | 3450233.0
+;
+
+partiallyUnmappedMixedTypesMultiIndexCast
+required_capability: optional_fields_nullify_tech_preview
+required_capability: index_metadata_field
+
+SET unmapped_fields="nullify"\;
+FROM sample_data, logs, no_mapping_sample_data METADATA _index
+| EVAL msg = message::KEYWORD
+| KEEP _index, @timestamp, msg
+| SORT _index, @timestamp DESC
+;
+
+_index:keyword         | @timestamp:date           | msg:keyword
+logs                   | 2023-10-23T13:57:01.544Z  | Running cats (cycle 3)
+logs                   | 2023-10-23T13:56:01.544Z  | Running cats (cycle 2)
+logs                   | 2023-10-23T13:56:01.543Z  | No response
+logs                   | 2023-10-23T13:55:01.546Z  | More java stuff
+logs                   | 2023-10-23T13:55:01.545Z  | Doing java stuff for 192.168.86.038
+logs                   | 2023-10-23T13:55:01.544Z  | Running cats (cycle 1)
+logs                   | 2023-10-23T13:55:01.543Z  | Pinging 192.168.86.046
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+no_mapping_sample_data | null                      | null
+sample_data            | 2023-10-23T13:55:01.543Z  | Connected to 10.1.0.1
+sample_data            | 2023-10-23T13:53:55.832Z  | Connection error
+sample_data            | 2023-10-23T13:52:55.015Z  | Connection error
+sample_data            | 2023-10-23T13:51:54.732Z  | Connection error
+sample_data            | 2023-10-23T13:33:34.937Z  | Disconnected
+sample_data            | 2023-10-23T12:27:28.948Z  | Connected to 10.1.0.2
+sample_data            | 2023-10-23T12:15:03.360Z  | Connected to 10.1.0.3
+;

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

@@ -845,6 +845,11 @@ public enum Cap {
          */
         SOURCE_FIELD_MAPPING,
 
+        /**
+         * Support for field aliases in mappings.
+         */
+        FIELD_ALIAS_SUPPORT,
+
         /**
          * Allow filter per individual aggregation.
          */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/core/expression/Alias.java

@@ -143,7 +143,7 @@ public String toString() {
 
     @Override
     public String nodeString(NodeStringFormat format) {
-        return child.nodeString() + " AS " + name() + "#" + id();
+        return child.nodeString(format) + " AS " + name() + "#" + id();
     }
 
     /**

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/core/expression/Attribute.java

@@ -221,7 +221,7 @@ public String toString() {
 
     @Override
     public String nodeString(NodeStringFormat format) {
-        return toString();
+        return toString(format);
     }
 
     protected abstract String label();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/core/expression/FieldAttribute.java

@@ -271,4 +271,14 @@ public boolean isMetric() {
     public EsField field() {
         return field;
     }
+
+    @Override
+    public String nodeString(NodeStringFormat format) {
+        return switch (format) {
+            case FULL -> Strings.format(
+                "%s{%s(%s)%s}#%s".formatted(qualifiedName(), label(), field.getWriteableName(), synthetic() ? "$" : "", id())
+            );
+            case LIMITED -> super.nodeString(format);
+        };
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/FragmentExec.java

@@ -137,7 +137,7 @@ public String nodeString(NodeStringFormat format) {
         sb.append(estimatedRowSize);
         sb.append(", reducer=[");
         sb.append("], fragment=[<>\n");
-        sb.append(fragment.toString());
+        sb.append(fragment.toString(format));
         sb.append("<>]]");
         return sb.toString();
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java

@@ -262,6 +262,8 @@ public static PhysicalPlan localPlan(
             return EstimatesRowSize.estimateRowSize(f.estimatedRowSize(), localOptimized);
         });
 
+        // TODO add a test assertion for the consistency checker (see
+        // https://github.com/elastic/elasticsearch/pull/141082/changes#r2745334028);
         PhysicalPlan resultPlan = isCoordPlan.get() ? plan : localPhysicalPlan;
 
         return resultPlan;

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

@@ -303,6 +303,7 @@ void runComputeOnRemoteCluster(
                         () -> exchangeSink.createExchangeSink(() -> {})
                     ),
                     coordinatorPlan,
+                    NodeReduceLocalPhysicalOptimization.ENABLED,
                     configuration.profile() ? new PlanTimeProfile() : null,
                     computeListener.acquireCompute()
                 );

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

@@ -206,7 +206,7 @@ public void execute(
         ActionListener<Result> listener
     ) {
         assert ThreadPool.assertCurrentThreadPool(
-            EsqlPlugin.ESQL_WORKER_THREAD_POOL_NAME,
+            ESQL_WORKER_THREAD_POOL_NAME,
             ThreadPool.Names.SYSTEM_READ,
             ThreadPool.Names.SEARCH,
             ThreadPool.Names.SEARCH_COORDINATION
@@ -284,7 +284,14 @@ public void execute(
                     })
                 )
             ) {
-                runCompute(rootTask, computeContext, mainPlan, planTimeProfile, localListener.acquireCompute());
+                runCompute(
+                    rootTask,
+                    computeContext,
+                    mainPlan,
+                    NodeReduceLocalPhysicalOptimization.ENABLED,
+                    planTimeProfile,
+                    localListener.acquireCompute()
+                );
 
                 for (int i = 0; i < subplans.size(); i++) {
                     var subplan = subplans.get(i);
@@ -387,7 +394,14 @@ public void executePlan(
                     })
                 )
             ) {
-                runCompute(rootTask, computeContext, coordinatorPlan, planTimeProfile, computeListener.acquireCompute());
+                runCompute(
+                    rootTask,
+                    computeContext,
+                    coordinatorPlan,
+                    NodeReduceLocalPhysicalOptimization.ENABLED,
+                    planTimeProfile,
+                    computeListener.acquireCompute()
+                );
                 return;
             }
         } else {
@@ -468,6 +482,7 @@ public void executePlan(
                             exchangeSinkSupplier
                         ),
                         coordinatorPlan,
+                        NodeReduceLocalPhysicalOptimization.ENABLED,
                         planTimeProfile,
                         localListener.acquireCompute()
                     );
@@ -642,6 +657,7 @@ void runCompute(
         CancellableTask task,
         ComputeContext context,
         PhysicalPlan plan,
+        NodeReduceLocalPhysicalOptimization nodeReduceLocalPhysicalOptimization,
         PlanTimeProfile planTimeProfile,
         ActionListener<DriverCompletionInfo> listener
     ) {
@@ -674,15 +690,18 @@ void runCompute(
 
             List<SearchExecutionContext> localContexts = new ArrayList<>();
             context.searchExecutionContexts().iterable().forEach(localContexts::add);
-            var localPlan = PlannerUtils.localPlan(
-                plannerSettings,
-                context.flags(),
-                localContexts,
-                context.configuration(),
-                context.foldCtx(),
-                plan,
-                planTimeProfile
-            );
+            var localPlan = switch (nodeReduceLocalPhysicalOptimization) {
+                case ENABLED -> PlannerUtils.localPlan(
+                    plannerSettings,
+                    context.flags(),
+                    localContexts,
+                    context.configuration(),
+                    context.foldCtx(),
+                    plan,
+                    planTimeProfile
+                );
+                case DISABLED -> plan;
+            };
             if (LOGGER.isDebugEnabled()) {
                 LOGGER.debug("Local plan for {}:\n{}", context.description(), localPlan);
             }
@@ -758,7 +777,8 @@ ActionListener<Void> addCompletionInfo(
         });
     }
 
-    static ReductionPlan reductionPlan(
+    // public for testing
+    public static ReductionPlan reductionPlan(
         PlannerSettings plannerSettings,
         EsqlFlags flags,
         Configuration configuration,
@@ -770,14 +790,19 @@ static ReductionPlan reductionPlan(
     ) {
         long startTime = planTimeProfile == null ? 0 : System.nanoTime();
         PhysicalPlan source = new ExchangeSourceExec(originalPlan.source(), originalPlan.output(), originalPlan.isIntermediateAgg());
-        ReductionPlan defaultResult = new ReductionPlan(originalPlan.replaceChild(source), originalPlan);
+        ReductionPlan defaultResult = new ReductionPlan(
+            originalPlan.replaceChild(source),
+            originalPlan,
+            NodeReduceLocalPhysicalOptimization.ENABLED
+        );
         if (reduceNodeLateMaterialization == false && runNodeLevelReduction == false) {
             return defaultResult;
         }
 
         Function<PhysicalPlan, ReductionPlan> placePlanBetweenExchanges = p -> new ReductionPlan(
             originalPlan.replaceChild(p.replaceChildren(List.of(source))),
-            originalPlan
+            originalPlan,
+            NodeReduceLocalPhysicalOptimization.ENABLED
         );
         // The default plan is just the exchange source piped directly into the exchange sink.
         ReductionPlan reductionPlan = switch (PlannerUtils.reductionPlan(originalPlan)) {
@@ -791,7 +816,10 @@ static ReductionPlan reductionPlan(
                 )
                     // Fallback to the behavior listed below, i.e., a regular top n reduction without loading new fields.
                     .orElseGet(() -> runNodeLevelReduction ? placePlanBetweenExchanges.apply(topN.plan()) : defaultResult);
-            case PlannerUtils.TopNReduction topN when runNodeLevelReduction -> placePlanBetweenExchanges.apply(topN.plan());
+            case PlannerUtils.TopNReduction topN when runNodeLevelReduction ->
+                // The TopN reduction plan should not be further optimized locally on the node reduce driver, since we took great pains to
+                // preplan in advance, including all the necessary field extractions!
+                placePlanBetweenExchanges.apply(topN.plan()).withoutNodeReduceLocalPhysicalOptimization();
             case PlannerUtils.ReducedPlan rp when runNodeLevelReduction -> placePlanBetweenExchanges.apply(rp.plan());
             default -> defaultResult;
         };