Add reduction plan

Author: carlosdelest

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/PlanTimeProfile.java

@@ -22,31 +22,36 @@
  *
  */
 public final class PlanTimeProfile implements Writeable, ToXContentObject {
+    private long reductionPlanNanos;
     private long logicalOptimizationNanos;
     private long physicalOptimizationNanos;
 
     /**
      * @param logicalOptimizationNanos  Time spent on local logical plan optimization (in nanoseconds)
      * @param physicalOptimizationNanos Time spent on local physical plan optimization (in nanoseconds)
+     * @param reductionPlanNanos Time spent on reduction plan for node_reduce phase (in nanoseconds)
      */
-    public PlanTimeProfile(long logicalOptimizationNanos, long physicalOptimizationNanos) {
+    public PlanTimeProfile(long logicalOptimizationNanos, long physicalOptimizationNanos, long reductionPlanNanos) {
         this.logicalOptimizationNanos = logicalOptimizationNanos;
         this.physicalOptimizationNanos = physicalOptimizationNanos;
+        this.reductionPlanNanos = reductionPlanNanos;
     }
 
     public PlanTimeProfile() {
         this.logicalOptimizationNanos = 0L;
         this.physicalOptimizationNanos = 0L;
+        this.reductionPlanNanos = 0L;
     }
 
     public PlanTimeProfile(StreamInput in) throws IOException {
-        this(in.readVLong(), in.readVLong());
+        this(in.readVLong(), in.readVLong(), in.readVLong());
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeVLong(logicalOptimizationNanos);
         out.writeVLong(physicalOptimizationNanos);
+        out.writeVLong(reductionPlanNanos);
     }
 
     @Override
@@ -57,6 +62,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (physicalOptimizationNanos > 0) {
             builder.field("physical_optimization_nanos", physicalOptimizationNanos);
         }
+        if (reductionPlanNanos > 0) {
+            builder.field("reduction_nanos", physicalOptimizationNanos);
+        }
         return builder;
     }
 
@@ -68,18 +76,23 @@ public void addPhysicalOptimizationPlanTime(long physicalOptimizationPlanTime) {
         this.physicalOptimizationNanos = this.physicalOptimizationNanos + physicalOptimizationPlanTime;
     }
 
+    public void addReductionPlanNanos(long reductionPlanNanos) {
+        this.reductionPlanNanos = this.reductionPlanNanos + reductionPlanNanos;
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (obj == this) return true;
         if (obj == null || obj.getClass() != this.getClass()) return false;
         var that = (PlanTimeProfile) obj;
         return this.logicalOptimizationNanos == that.logicalOptimizationNanos
-            && this.physicalOptimizationNanos == that.physicalOptimizationNanos;
+            && this.physicalOptimizationNanos == that.physicalOptimizationNanos
+            && this.reductionPlanNanos == that.reductionPlanNanos;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(logicalOptimizationNanos, physicalOptimizationNanos);
+        return Objects.hash(logicalOptimizationNanos, physicalOptimizationNanos, reductionPlanNanos);
     }
 
     @Override
@@ -90,6 +103,9 @@ public String toString() {
             + ", "
             + "physicalOptimizationNanos="
             + physicalOptimizationNanos
+            + ", "
+            + "reductionPlanNanos="
+            + reductionPlanNanos
             + ']';
     }
 

x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/RestEsqlIT.java

@@ -372,12 +372,20 @@ public void testProfile() throws IOException {
             assertThat(plan.get("cluster_name"), equalTo("test-cluster"));
             assertThat(plan.get("node_name"), equalTo("test-cluster-0"));
             assertThat(plan.get("plan"), notNullValue());
-            assertThat((int) plan.get("plan_nanos"), greaterThanOrEqualTo(0));
             String description = (String) plan.get("description");
             assertTrue("Unexpected plan description " + description, Set.of("final", "node_reduce", "data").contains(description));
-            if ("final".equals(description) || "data".equals(description)) {
-                assertThat((int) plan.get("logical_optimization_nanos"), greaterThanOrEqualTo(0));
-                assertThat((int) plan.get("physical_optimization_nanos"), greaterThanOrEqualTo(0));
+            switch(description) {
+                case "final", "data" -> {
+                    assertThat((int) plan.get("logical_optimization_nanos"), greaterThanOrEqualTo(0));
+                    assertThat((int) plan.get("physical_optimization_nanos"), greaterThanOrEqualTo(0));
+                    assertFalse(plan.containsKey("reduction_nanos"));
+                }
+                case "node_reduce" -> {
+                    assertThat((int) plan.get("reduction_nanos"), greaterThanOrEqualTo(0));
+                    assertFalse(plan.containsKey("logical_optimization_nanos"));
+                    assertFalse(plan.containsKey("physical_optimization_nanos"));
+                }
+                default -> {}
             }
         }
     }

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

@@ -268,7 +268,8 @@ void runComputeOnRemoteCluster(
             configuration.newFoldContext(),
             plan,
             true,
-            false
+            false,
+            configuration.profile() ? new PlanTimeProfile() : null
         );
         PhysicalPlan coordinatorPlan = reductionPlan.nodeReducePlan();
         final AtomicReference<ComputeResponse> finalResponse = new AtomicReference<>();

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

@@ -742,8 +742,10 @@ static ReductionPlan reductionPlan(
         FoldContext foldCtx,
         ExchangeSinkExec originalPlan,
         boolean runNodeLevelReduction,
-        boolean reduceNodeLateMaterialization
+        boolean reduceNodeLateMaterialization,
+        PlanTimeProfile planTimeProfile
     ) {
+        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);
         if (reduceNodeLateMaterialization == false && runNodeLevelReduction == false) {
@@ -755,7 +757,7 @@ static ReductionPlan reductionPlan(
             originalPlan
         );
         // The default plan is just the exchange source piped directly into the exchange sink.
-        return switch (PlannerUtils.reductionPlan(originalPlan)) {
+        ReductionPlan reductionPlan = switch (PlannerUtils.reductionPlan(originalPlan)) {
             case PlannerUtils.TopNReduction topN when reduceNodeLateMaterialization ->
                 // In the case of TopN, the source output type is replaced since we're pulling the FieldExtractExec to the reduction node,
                 // so essential we are splitting the TopNExec into two parts, similar to other aggregations, but unlike other aggregations,
@@ -770,6 +772,10 @@ static ReductionPlan reductionPlan(
             case PlannerUtils.ReducedPlan rp when runNodeLevelReduction -> placePlanBetweenExchanges.apply(rp.plan());
             default -> defaultResult;
         };
+        if (planTimeProfile != null) {
+            planTimeProfile.addReductionPlanNanos(System.nanoTime() - startTime);
+        }
+        return reductionPlan;
     }
 
     String newChildSession(String session) {

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

@@ -569,6 +569,9 @@ public void messageReceived(DataNodeRequest request, TransportChannel channel, T
         // We can avoid synchronization (for the most part) since the array elements are never modified, and the array is only added to,
         // with its size being known before we start the computation.
         PlanTimeProfile planTimeProfile = null;
+        if (configuration.profile()) {
+            planTimeProfile = new PlanTimeProfile();
+        }
         if (request.plan() instanceof ExchangeSinkExec plan) {
             reductionPlan = ComputeService.reductionPlan(
                 computeService.plannerSettings(),
@@ -577,11 +580,9 @@ public void messageReceived(DataNodeRequest request, TransportChannel channel, T
                 configuration.newFoldContext(),
                 plan,
                 request.runNodeLevelReduction(),
-                request.reductionLateMaterialization()
+                request.reductionLateMaterialization(),
+                planTimeProfile
             );
-            if (configuration.profile()) {
-                planTimeProfile = new PlanTimeProfile();
-            }
         } else {
             listener.onFailure(new IllegalStateException("expected exchange sink for a remote compute; got " + request.plan()));
             return;

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponseProfileTests.java

@@ -74,7 +74,7 @@ private List<PlanProfile> randomPlanProfiles() {
     }
 
     private PlanTimeProfile randomPlanTimeProfile() {
-        return randomBoolean() ? null : new PlanTimeProfile(randomNonNegativeLong(), randomNonNegativeLong());
+        return randomBoolean() ? null : new PlanTimeProfile(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong());
     }
 
     private OperatorStatus randomOperatorStatus() {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/ComputeListenerTests.java

@@ -95,7 +95,7 @@ private DriverCompletionInfo randomCompletionInfo() {
     }
 
     private PlanTimeProfile randomPlanTimeProfile() {
-        return randomBoolean() ? null : new PlanTimeProfile(randomNonNegativeLong(), randomNonNegativeLong());
+        return randomBoolean() ? null : new PlanTimeProfile(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong());
     }
 
     public void testEmpty() {