ESQL: Fix ExchangeSinkExec output type

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

@@ -77,6 +77,11 @@ public ExchangeSinkExec replaceChild(PhysicalPlan newChild) {
         return new ExchangeSinkExec(source(), output, intermediateAgg, newChild);
     }
 
+    /** Like {@link #replaceChild(PhysicalPlan)} but sets the output to the new child's output. */
+    public ExchangeSinkExec replaceChildAndUpdateOutput(PhysicalPlan newChild) {
+        return new ExchangeSinkExec(source(), newChild.output(), intermediateAgg, newChild);
+    }
+
     @Override
     public boolean equals(Object o) {
         if (super.equals(o)) {

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.operator.PlanTimeProfile;
+import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.IndexMode;
@@ -40,6 +41,7 @@
 import org.elasticsearch.xpack.esql.optimizer.LocalLogicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext;
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizer;
+import org.elasticsearch.xpack.esql.optimizer.PhysicalVerifier;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.plan.QueryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
@@ -237,8 +239,6 @@ public static PhysicalPlan localPlan(
         LocalPhysicalPlanOptimizer physicalOptimizer,
         PlanTimeProfile planTimeProfile
     ) {
-        // TODO add a test assertion for the consistency checker (after https://github.com/elastic/elasticsearch/issues/141654, see
-        // https://github.com/elastic/elasticsearch/pull/141082/changes#r2745334028);
         var isCoordPlan = new Holder<>(Boolean.TRUE);
         Set<PhysicalPlan> lookupJoinExecRightChildren = plan.collect(LookupJoinExec.class::isInstance)
             .stream()
@@ -281,7 +281,13 @@ public static PhysicalPlan localPlan(
         });
 
         PhysicalPlan resultPlan = isCoordPlan.get() ? plan : localPhysicalPlan;
-
+        // This check is needed because in test code we sometimes invoke localPlan with a non-ExchangeSinkExec root.
+        if (resultPlan instanceof ExchangeSinkExec sink) {
+            resultPlan = sink.replaceChildAndUpdateOutput(sink.child());
+        }
+        if (Assertions.ENABLED) {
+            PhysicalVerifier.LOCAL_INSTANCE.verify(resultPlan, resultPlan.output());
+        }
         return resultPlan;
     }
 

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

@@ -30,6 +30,7 @@
 import org.elasticsearch.compute.operator.exchange.ExchangeSink;
 import org.elasticsearch.compute.operator.exchange.ExchangeSinkHandler;
 import org.elasticsearch.compute.operator.exchange.ExchangeSourceHandler;
+import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -58,6 +59,7 @@
 import org.elasticsearch.xpack.esql.enrich.LookupFromIndexService;
 import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext;
+import org.elasticsearch.xpack.esql.optimizer.PhysicalVerifier;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSourceExec;
@@ -822,12 +824,32 @@ public 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.ReducedPlan rp when runNodeLevelReduction -> placePlanBetweenExchanges.apply(rp.plan());
+            case PlannerUtils.ReducedPlan rp when runNodeLevelReduction -> {
+                PhysicalPlan reductionSource = new ExchangeSourceExec(
+                    originalPlan.source(),
+                    new ArrayList<>(rp.plan().references()),
+                    originalPlan.isIntermediateAgg()
+                );
+                yield new ReductionPlan(
+                    originalPlan.replaceChild(rp.plan().replaceChildren(List.of(reductionSource))),
+                    originalPlan,
+                    LocalPhysicalOptimization.ENABLED
+                );
+            }
             default -> defaultResult;
         };
         if (planTimeProfile != null) {
             planTimeProfile.addReductionPlanNanos(System.nanoTime() - startTime);
         }
+        reductionPlan = new ReductionPlan(
+            reductionPlan.nodeReducePlan().replaceChildAndUpdateOutput(reductionPlan.nodeReducePlan().child()),
+            reductionPlan.dataNodePlan().replaceChildAndUpdateOutput(reductionPlan.dataNodePlan().child()),
+            reductionPlan.localPhysicalOptimization()
+        );
+        if (Assertions.ENABLED) {
+            PhysicalVerifier.LOCAL_INSTANCE.verify(reductionPlan.nodeReducePlan(), reductionPlan.nodeReducePlan().child().output());
+            PhysicalVerifier.LOCAL_INSTANCE.verify(reductionPlan.dataNodePlan(), reductionPlan.dataNodePlan().child().output());
+        }
         return reductionPlan;
     }
 

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

@@ -128,8 +128,7 @@ public static Optional<ReductionPlan> planReduceDriverTopN(
         }
         var updatedFragment = new Project(Source.EMPTY, withAddedDocToRelation, expectedDataOutput);
         FragmentExec updatedFragmentExec = fragmentExec.withFragment(updatedFragment);
-        // TODO This ignores the possible change in output, see #141654
-        ExchangeSinkExec updatedDataPlan = originalPlan.replaceChild(updatedFragmentExec);
+        ExchangeSinkExec updatedDataPlan = originalPlan.replaceChildAndUpdateOutput(updatedFragmentExec);
 
         // Replace the TopN child with the data driver as the source.
         PhysicalPlan reductionPlan = toPhysical(fragmentExec.fragment(), context).transformDown(TopNExec.class, t -> {
@@ -138,9 +137,8 @@ public static Optional<ReductionPlan> planReduceDriverTopN(
             boolean fragmentIsSorted = updatedFragment.child() instanceof TopN;
             return fragmentIsSorted ? t.replaceChild(exchangeExec).withSortedInput() : t.replaceChild(exchangeExec);
         });
-        ExchangeSinkExec reductionPlanWithSize = originalPlan.replaceChild(
-            EstimatesRowSize.estimateRowSize(updatedFragmentExec.estimatedRowSize(), reductionPlan)
-        );
+        PhysicalPlan sizedReductionPlan = EstimatesRowSize.estimateRowSize(updatedFragmentExec.estimatedRowSize(), reductionPlan);
+        ExchangeSinkExec reductionPlanWithSize = originalPlan.replaceChildAndUpdateOutput(sizedReductionPlan);
 
         // 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!

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testBasicTopNLateMaterialization/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}],false]
 \_ProjectExec[[_doc{f}, hire_date{f}]]
   \_FieldExtractExec[hire_date{f}]<[],[]>
     \_EsQueryExec[employees], indexMode[standard], [_doc{f}], limit[20], sort[[FieldSort[field=hire_date{f}, direction=ASC, nulls=LAST]]] estimatedRowSize[24] queryBuilderAndTags [[QueryBuilderAndTags[query=null, tags=[]]]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testBasicTopNLateMaterialization/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, hire_date{f}]]
 \_TopN[[Order[hire_date{f},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testExpressionSortTopNKeepAfterSort/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[hire_date{f}, $$order_by{r}],false]
+ExchangeSinkExec[[_doc{f}, $$order_by{r}],false]
 \_ProjectExec[[_doc{f}, $$order_by{r}]]
   \_TopNExec[[Order[$$order_by{r},ASC,LAST]],20[INTEGER],36]
     \_EvalExec[[SIN(height{f}) * 2[INTEGER] AS $$order_by]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testExpressionSortTopNKeepAfterSort/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[hire_date{f}, $$order_by{r}],false]
+ExchangeSinkExec[[_doc{f}, $$order_by{r}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, $$order_by{r}]]
 \_TopN[[Order[$$order_by{r},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testExpressionSortTopNKeepBeforeSort/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[height{f}, hire_date{f}, $$order_by{r}],false]
+ExchangeSinkExec[[_doc{f}, height{f}, $$order_by{r}],false]
 \_ProjectExec[[_doc{f}, height{f}, $$order_by{r}]]
   \_TopNExec[[Order[$$order_by{r},ASC,LAST]],20[INTEGER],36]
     \_EvalExec[[SIN(height{f}) * 2[INTEGER] AS $$order_by]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testExpressionSortTopNKeepBeforeSort/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[height{f}, hire_date{f}, $$order_by{r}],false]
+ExchangeSinkExec[[_doc{f}, height{f}, $$order_by{r}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, height{f}, $$order_by{r}]]
 \_TopN[[Order[$$order_by{r},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testLookupJoinOnDataNode/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[avg_worked_seconds{f}, birth_date{f}, emp_no{f}, first_name{f}, gender{f}, height{f}, height.float{f}, height.half_float{f}, height.scaled_float{f}, hire_date{f}, is_rehired{f}, job_positions{f}, languages{f}, languages.byte{f}, languages.long{f}, languages.short{f}, last_name{f}, salary{f}, salary_change{f}, salary_change.int{f}, salary_change.keyword{f}, salary_change.long{f}, still_hired{f}, language_code{r}, language_name{f}],false]
+ExchangeSinkExec[[_doc{f}, emp_no{f}, languages{f}, language_code{r}, language_name{f}],false]
 \_ProjectExec[[_doc{f}, emp_no{f}, languages{f}, language_code{r}, language_name{f}]]
   \_TopNExec[[Order[emp_no{f},ASC,LAST]],20[INTEGER],82]
     \_FieldExtractExec[emp_no{f}]<[],[]>

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testLookupJoinOnDataNode/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[avg_worked_seconds{f}, birth_date{f}, emp_no{f}, first_name{f}, gender{f}, height{f}, height.float{f}, height.half_float{f}, height.scaled_float{f}, hire_date{f}, is_rehired{f}, job_positions{f}, languages{f}, languages.byte{f}, languages.long{f}, languages.short{f}, last_name{f}, salary{f}, salary_change{f}, salary_change.int{f}, salary_change.keyword{f}, salary_change.long{f}, still_hired{f}, language_code{r}, language_name{f}],false]
+ExchangeSinkExec[[_doc{f}, emp_no{f}, languages{f}, language_code{r}, language_name{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, emp_no{f}, languages{f}, language_code{r}, language_name{f}]]
 \_TopN[[Order[emp_no{f},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testMultipleFieldSortTopN/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, height{f}, hire_date{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}, height{f}],false]
 \_ProjectExec[[_doc{f}, hire_date{f}, height{f}]]
   \_FieldExtractExec[hire_date{f}, height{f}]<[],[]>
     \_EsQueryExec[employees], indexMode[standard], [_doc{f}], limit[20], sort[[FieldSort[field=hire_date{f}, direction=ASC, nulls=LAST], FieldSort[field=height{f}, direction=ASC, nulls=LAST]]] estimatedRowSize[32] queryBuilderAndTags [[QueryBuilderAndTags[query=null, tags=[]]]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testMultipleFieldSortTopN/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, height{f}, hire_date{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}, height{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, hire_date{f}, height{f}]]
 \_TopN[[Order[hire_date{f},ASC,LAST], Order[height{f},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testMultipleTopN/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}],false]
 \_ProjectExec[[_doc{f}, hire_date{f}]]
   \_FieldExtractExec[hire_date{f}]<[],[]>
     \_EsQueryExec[employees], indexMode[standard], [_doc{f}], limit[20], sort[[FieldSort[field=hire_date{f}, direction=ASC, nulls=LAST]]] estimatedRowSize[24] queryBuilderAndTags [[QueryBuilderAndTags[query=null, tags=[]]]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testMultipleTopN/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, hire_date{f}]]
 \_TopN[[Order[hire_date{f},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testPushedDownTopN/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, height{f}],false]
+ExchangeSinkExec[[_doc{f}, height{f}],false]
 \_ProjectExec[[_doc{f}, height{f}]]
   \_FieldExtractExec[height{f}]<[],[]>
     \_EsQueryExec[employees], indexMode[standard], [_doc{f}], limit[20], sort[[FieldSort[field=height{f}, direction=ASC, nulls=LAST]]] estimatedRowSize[24] queryBuilderAndTags [[QueryBuilderAndTags[query=null, tags=[]]]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testPushedDownTopN/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, height{f}],false]
+ExchangeSinkExec[[_doc{f}, height{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, height{f}]]
 \_TopN[[Order[height{f},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testSomeFieldsNeededBeforeLateMaterialization/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, salary{f}, hire_date{f}],false]
 \_ProjectExec[[_doc{f}, salary{f}, hire_date{f}]]
   \_FieldExtractExec[salary{f}, hire_date{f}]<[],[]>
     \_EsQueryExec[employees], indexMode[standard], [_doc{f}], limit[20], sort[[FieldSort[field=hire_date{f}, direction=ASC, nulls=LAST]]] estimatedRowSize[28] queryBuilderAndTags [[QueryBuilderAndTags[query={

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testSomeFieldsNeededBeforeLateMaterialization/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, salary{f}, hire_date{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, salary{f}, hire_date{f}]]
 \_TopN[[Order[hire_date{f},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testTopNThenStats/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}],false]
 \_ProjectExec[[_doc{f}, hire_date{f}]]
   \_FieldExtractExec[hire_date{f}]<[],[]>
     \_EsQueryExec[employees], indexMode[standard], [_doc{f}], limit[20], sort[[FieldSort[field=hire_date{f}, direction=ASC, nulls=LAST]]] estimatedRowSize[24] queryBuilderAndTags [[QueryBuilderAndTags[query=null, tags=[]]]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testTopNThenStats/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, hire_date{f}]]
 \_TopN[[Order[hire_date{f},ASC,LAST]],20[INTEGER],false]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testTopNWithMissingSortField/local_reduce_physical_optimization_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{r}],false]
 \_ProjectExec[[_doc{f}, birth_date{r} AS hire_date]]
   \_EvalExec[[null[DATETIME] AS birth_date]]
     \_EsQueryExec[employees], indexMode[standard], [_doc{f}], limit[20], sort[] estimatedRowSize[12] queryBuilderAndTags [[QueryBuilderAndTags[query=null, tags=[]]]]

x-pack/plugin/esql/src/test/resources/org/elasticsearch/xpack/esql/plugin/golden_tests/LateMaterializationPlannerGoldenTests/testTopNWithMissingSortField/local_reduce_planned_data_driver.expected

@@ -1,4 +1,4 @@
-ExchangeSinkExec[[emp_no{f}, hire_date{f}, salary{f}],false]
+ExchangeSinkExec[[_doc{f}, hire_date{f}],false]
 \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[<>
 Project[[_doc{f}, hire_date{f}]]
 \_TopN[[Order[hire_date{f},ASC,LAST]],20[INTEGER],false]