Track blocks of intermediate state of aggs (#102562)

This PR tracks blocks created when evaluating the intermediate state of aggregation.
elastic · Nov 24, 2023 · e0df305 · e0df305
1 parent 391f291
commit e0df305
Show file tree

Hide file tree

Showing 35 changed files with 98 additions and 98 deletions.
diff --git a/docs/changelog/102562.yaml b/docs/changelog/102562.yaml
@@ -0,0 +1,5 @@
+pr: 102562
+summary: Track blocks of intermediate state of aggs
+area: ES|QL
+type: enhancement
+issues: []
diff --git a/...n/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/AggregatorImplementer.java b/...n/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/AggregatorImplementer.java
@@ -478,8 +478,9 @@ private MethodSpec evaluateIntermediate() {
         builder.addAnnotation(Override.class)
             .addModifiers(Modifier.PUBLIC)
             .addParameter(BLOCK_ARRAY, "blocks")
-            .addParameter(TypeName.INT, "offset");
-        builder.addStatement("state.toIntermediate(blocks, offset)");
+            .addParameter(TypeName.INT, "offset")
+            .addParameter(DRIVER_CONTEXT, "driverContext");
+        builder.addStatement("state.toIntermediate(blocks, offset, driverContext)");
         return builder.build();
     }
 

diff --git a/...sql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/DoubleState.java b/...sql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/DoubleState.java
@@ -8,8 +8,7 @@
 package org.elasticsearch.compute.aggregation;
 
 import org.elasticsearch.compute.data.Block;
-import org.elasticsearch.compute.data.ConstantBooleanVector;
-import org.elasticsearch.compute.data.ConstantDoubleVector;
+import org.elasticsearch.compute.operator.DriverContext;
 
 /**
  * Aggregator state for a single double.
@@ -45,10 +44,10 @@ void seen(boolean seen) {
 
     /** Extracts an intermediate view of the contents of this state.  */
     @Override
-    public void toIntermediate(Block[] blocks, int offset) {
+    public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
         assert blocks.length >= offset + 2;
-        blocks[offset + 0] = new ConstantDoubleVector(value, 1).asBlock();
-        blocks[offset + 1] = new ConstantBooleanVector(seen, 1).asBlock();
+        blocks[offset + 0] = driverContext.blockFactory().newConstantDoubleBlockWith(value, 1);
+        blocks[offset + 1] = driverContext.blockFactory().newConstantBooleanBlockWith(seen, 1);
     }
 
     @Override

diff --git a/...n/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/IntState.java b/...n/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/IntState.java
@@ -8,8 +8,7 @@
 package org.elasticsearch.compute.aggregation;
 
 import org.elasticsearch.compute.data.Block;
-import org.elasticsearch.compute.data.ConstantBooleanVector;
-import org.elasticsearch.compute.data.ConstantIntVector;
+import org.elasticsearch.compute.operator.DriverContext;
 
 /**
  * Aggregator state for a single int.
@@ -45,10 +44,10 @@ void seen(boolean seen) {
 
     /** Extracts an intermediate view of the contents of this state.  */
     @Override
-    public void toIntermediate(Block[] blocks, int offset) {
+    public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
         assert blocks.length >= offset + 2;
-        blocks[offset + 0] = new ConstantIntVector(value, 1).asBlock();
-        blocks[offset + 1] = new ConstantBooleanVector(seen, 1).asBlock();
+        blocks[offset + 0] = driverContext.blockFactory().newConstantIntBlockWith(value, 1);
+        blocks[offset + 1] = driverContext.blockFactory().newConstantBooleanBlockWith(seen, 1);
     }
 
     @Override

diff --git a/.../esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/LongState.java b/.../esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/LongState.java
@@ -8,8 +8,7 @@
 package org.elasticsearch.compute.aggregation;
 
 import org.elasticsearch.compute.data.Block;
-import org.elasticsearch.compute.data.ConstantBooleanVector;
-import org.elasticsearch.compute.data.ConstantLongVector;
+import org.elasticsearch.compute.operator.DriverContext;
 
 /**
  * Aggregator state for a single long.
@@ -45,10 +44,10 @@ void seen(boolean seen) {
 
     /** Extracts an intermediate view of the contents of this state.  */
     @Override
-    public void toIntermediate(Block[] blocks, int offset) {
+    public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
         assert blocks.length >= offset + 2;
-        blocks[offset + 0] = new ConstantLongVector(value, 1).asBlock();
-        blocks[offset + 1] = new ConstantBooleanVector(seen, 1).asBlock();
+        blocks[offset + 0] = driverContext.blockFactory().newConstantLongBlockWith(value, 1);
+        blocks[offset + 1] = driverContext.blockFactory().newConstantBooleanBlockWith(seen, 1);
     }
 
     @Override

diff --git a/...nerated/org/elasticsearch/compute/aggregation/CountDistinctBooleanAggregatorFunction.java b/...nerated/org/elasticsearch/compute/aggregation/CountDistinctBooleanAggregatorFunction.java
@@ -98,8 +98,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...erated/org/elasticsearch/compute/aggregation/CountDistinctBytesRefAggregatorFunction.java b/...erated/org/elasticsearch/compute/aggregation/CountDistinctBytesRefAggregatorFunction.java
@@ -106,8 +106,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...enerated/org/elasticsearch/compute/aggregation/CountDistinctDoubleAggregatorFunction.java b/...enerated/org/elasticsearch/compute/aggregation/CountDistinctDoubleAggregatorFunction.java
@@ -106,8 +106,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...n/generated/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunction.java b/...n/generated/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunction.java
@@ -106,8 +106,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/.../generated/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunction.java b/.../generated/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunction.java
@@ -106,8 +106,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...src/main/generated/org/elasticsearch/compute/aggregation/MaxDoubleAggregatorFunction.java b/...src/main/generated/org/elasticsearch/compute/aggregation/MaxDoubleAggregatorFunction.java
@@ -105,8 +105,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...te/src/main/generated/org/elasticsearch/compute/aggregation/MaxIntAggregatorFunction.java b/...te/src/main/generated/org/elasticsearch/compute/aggregation/MaxIntAggregatorFunction.java
@@ -105,8 +105,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...e/src/main/generated/org/elasticsearch/compute/aggregation/MaxLongAggregatorFunction.java b/...e/src/main/generated/org/elasticsearch/compute/aggregation/MaxLongAggregatorFunction.java
@@ -105,8 +105,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...rg/elasticsearch/compute/aggregation/MedianAbsoluteDeviationDoubleAggregatorFunction.java b/...rg/elasticsearch/compute/aggregation/MedianAbsoluteDeviationDoubleAggregatorFunction.java
@@ -99,8 +99,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...d/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationIntAggregatorFunction.java b/...d/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationIntAggregatorFunction.java
@@ -99,8 +99,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/.../org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationLongAggregatorFunction.java b/.../org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationLongAggregatorFunction.java
@@ -99,8 +99,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...src/main/generated/org/elasticsearch/compute/aggregation/MinDoubleAggregatorFunction.java b/...src/main/generated/org/elasticsearch/compute/aggregation/MinDoubleAggregatorFunction.java
@@ -105,8 +105,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...te/src/main/generated/org/elasticsearch/compute/aggregation/MinIntAggregatorFunction.java b/...te/src/main/generated/org/elasticsearch/compute/aggregation/MinIntAggregatorFunction.java
@@ -105,8 +105,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...e/src/main/generated/org/elasticsearch/compute/aggregation/MinLongAggregatorFunction.java b/...e/src/main/generated/org/elasticsearch/compute/aggregation/MinLongAggregatorFunction.java
@@ -105,8 +105,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...n/generated/org/elasticsearch/compute/aggregation/PercentileDoubleAggregatorFunction.java b/...n/generated/org/elasticsearch/compute/aggregation/PercentileDoubleAggregatorFunction.java
@@ -102,8 +102,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...main/generated/org/elasticsearch/compute/aggregation/PercentileIntAggregatorFunction.java b/...main/generated/org/elasticsearch/compute/aggregation/PercentileIntAggregatorFunction.java
@@ -102,8 +102,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...ain/generated/org/elasticsearch/compute/aggregation/PercentileLongAggregatorFunction.java b/...ain/generated/org/elasticsearch/compute/aggregation/PercentileLongAggregatorFunction.java
@@ -102,8 +102,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...src/main/generated/org/elasticsearch/compute/aggregation/SumDoubleAggregatorFunction.java b/...src/main/generated/org/elasticsearch/compute/aggregation/SumDoubleAggregatorFunction.java
@@ -104,8 +104,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...te/src/main/generated/org/elasticsearch/compute/aggregation/SumIntAggregatorFunction.java b/...te/src/main/generated/org/elasticsearch/compute/aggregation/SumIntAggregatorFunction.java
@@ -107,8 +107,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...e/src/main/generated/org/elasticsearch/compute/aggregation/SumLongAggregatorFunction.java b/...e/src/main/generated/org/elasticsearch/compute/aggregation/SumLongAggregatorFunction.java
@@ -105,8 +105,8 @@ public void addIntermediateInput(Page page) {
   }
 
   @Override
-  public void evaluateIntermediate(Block[] blocks, int offset) {
-    state.toIntermediate(blocks, offset);
+  public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+    state.toIntermediate(blocks, offset, driverContext);
   }
 
   @Override

diff --git a/...k/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/Aggregator.java b/...k/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/Aggregator.java
@@ -45,7 +45,7 @@ public void processPage(Page page) {
 
     public void evaluate(Block[] blocks, int offset, DriverContext driverContext) {
         if (mode.isOutputPartial()) {
-            aggregatorFunction.evaluateIntermediate(blocks, offset);
+            aggregatorFunction.evaluateIntermediate(blocks, offset, driverContext);
         } else {
             aggregatorFunction.evaluateFinal(blocks, offset, driverContext);
         }

diff --git a/.../esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AggregatorFunction.java b/.../esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AggregatorFunction.java
@@ -18,7 +18,7 @@ public interface AggregatorFunction extends Releasable {
 
     void addIntermediateInput(Page page);
 
-    void evaluateIntermediate(Block[] blocks, int offset);
+    void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext);
 
     void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext);
 

diff --git a/...gin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AggregatorState.java b/...gin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AggregatorState.java
@@ -8,10 +8,11 @@
 package org.elasticsearch.compute.aggregation;
 
 import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.core.Releasable;
 
 public interface AggregatorState extends Releasable {
 
     /** Extracts an intermediate view of the contents of this state.  */
-    void toIntermediate(Block[] blocks, int offset);
+    void toIntermediate(Block[] blocks, int offset, DriverContext driverContext);
 }
diff --git a/.../compute/src/main/java/org/elasticsearch/compute/aggregation/CountAggregatorFunction.java b/.../compute/src/main/java/org/elasticsearch/compute/aggregation/CountAggregatorFunction.java
@@ -97,13 +97,13 @@ public void addIntermediateInput(Page page) {
     }
 
     @Override
-    public void evaluateIntermediate(Block[] blocks, int offset) {
-        state.toIntermediate(blocks, offset);
+    public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
+        state.toIntermediate(blocks, offset, driverContext);
     }
 
     @Override
     public void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext) {
-        blocks[offset] = LongBlock.newConstantBlockWith(state.longValue(), 1, driverContext.blockFactory());
+        blocks[offset] = driverContext.blockFactory().newConstantLongBlockWith(state.longValue(), 1);
     }
 
     @Override

diff --git a/...e/src/main/java/org/elasticsearch/compute/aggregation/CountDistinctBooleanAggregator.java b/...e/src/main/java/org/elasticsearch/compute/aggregation/CountDistinctBooleanAggregator.java
@@ -86,10 +86,10 @@ static class SingleState implements AggregatorState {
 
         /** Extracts an intermediate view of the contents of this state.  */
         @Override
-        public void toIntermediate(Block[] blocks, int offset) {
+        public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
             assert blocks.length >= offset + 2;
-            blocks[offset + 0] = BooleanBlock.newConstantBlockWith((bits & BIT_FALSE) != 0, 1);
-            blocks[offset + 1] = BooleanBlock.newConstantBlockWith((bits & BIT_TRUE) != 0, 1);
+            blocks[offset + 0] = driverContext.blockFactory().newConstantBooleanBlockWith((bits & BIT_FALSE) != 0, 1);
+            blocks[offset + 1] = driverContext.blockFactory().newConstantBooleanBlockWith((bits & BIT_TRUE) != 0, 1);
         }
 
         @Override

diff --git a/...ck/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/HllStates.java b/...ck/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/HllStates.java
@@ -17,7 +17,6 @@
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BytesRefBlock;
-import org.elasticsearch.compute.data.ConstantBytesRefVector;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.core.Releasables;
@@ -118,9 +117,9 @@ void merge(int groupId, BytesRef other, int otherGroup) {
 
         /** Extracts an intermediate view of the contents of this state.  */
         @Override
-        public void toIntermediate(Block[] blocks, int offset) {
+        public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
             assert blocks.length >= offset + 1;
-            blocks[offset] = new ConstantBytesRefVector(serializeHLL(SINGLE_BUCKET_ORD, hll), 1).asBlock();
+            blocks[offset] = driverContext.blockFactory().newConstantBytesRefBlockWith(serializeHLL(SINGLE_BUCKET_ORD, hll), 1);
         }
 
         @Override