Add optimized path for intermediate values aggregator

Author: dnhatn

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesAggregatorBenchmark.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.compute.aggregation.AggregatorFunction;
 import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier;
 import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.compute.aggregation.ValuesBytesRefAggregatorFunctionSupplier;
@@ -57,8 +58,8 @@
  * Benchmark for the {@code VALUES} aggregator that supports grouping by many many
  * many values.
  */
-@Warmup(iterations = 5)
-@Measurement(iterations = 7)
+@Warmup(iterations = 1)
+@Measurement(iterations = 2)
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 @State(Scope.Thread)
@@ -107,22 +108,22 @@ static void selfTest() {
  private static final String INT = "int";
  private static final String LONG = "long";
 
- @Param({ "1", "1000", /*"1000000"*/ })
+ @Param({ "1", "1000", "1000000" })
  public int groups;
 
- @Param({ BYTES_REF, INT, LONG })
+ @Param({ BYTES_REF })
  public String dataType;
 
- private static Operator operator(DriverContext driverContext, int groups, String dataType) {
+ private static Operator operator(DriverContext driverContext, int groups, String dataType, AggregatorMode mode) {
  if (groups == 1) {
  return new AggregationOperator(
- List.of(supplier(dataType).aggregatorFactory(AggregatorMode.SINGLE, List.of(0)).apply(driverContext)),
+ List.of(supplier(dataType).aggregatorFactory(mode, List.of(0)).apply(driverContext)),
  driverContext
  );
  }
  List<BlockHash.GroupSpec> groupSpec = List.of(new BlockHash.GroupSpec(0, ElementType.LONG));
  return new HashAggregationOperator(
- List.of(supplier(dataType).groupingAggregatorFactory(AggregatorMode.SINGLE, List.of(1))),
+ List.of(supplier(dataType).groupingAggregatorFactory(mode, List.of(1))),
  () -> BlockHash.build(groupSpec, driverContext.blockFactory(), 16 * 1024, false),
  driverContext
  ) {
@@ -341,13 +342,21 @@ public void run() {
 
  private static void run(int groups, String dataType, int opCount) {
  DriverContext driverContext = driverContext();
- try (Operator operator = operator(driverContext, groups, dataType)) {
- Page page = page(groups, dataType);
- for (int i = 0; i < opCount; i++) {
- operator.addInput(page.shallowCopy());
+ try (Operator finalAggregator = operator(driverContext, groups, dataType, AggregatorMode.FINAL)) {
+ try (Operator initialAggregator = operator(driverContext, groups, dataType, AggregatorMode.INITIAL)) {
+ Page rawPage = page(groups, dataType);
+ for (int i = 0; i < opCount; i++) {
+ initialAggregator.addInput(rawPage.shallowCopy());
+ }
+ initialAggregator.finish();
+ Page intermediatePage = initialAggregator.getOutput();
+ for (int i = 0; i < opCount; i++) {
+ finalAggregator.addInput(intermediatePage.shallowCopy());
+ }
+ finalAggregator.finish();
+ Page outputPage = finalAggregator.getOutput();
+ checkExpected(groups, dataType, outputPage);
  }
- operator.finish();
- checkExpected(groups, dataType, operator.getOutput());
  }
  }
 

x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/GroupingAggregatorImplementer.java

@@ -609,63 +609,80 @@ private MethodSpec addIntermediateInput() {
  .collect(joining(" && "))
  );
  }
- if (intermediateState.stream().map(AggregatorImplementer.IntermediateStateDesc::elementType).anyMatch(n -> n.equals("BYTES_REF"))) {
- builder.addStatement("$T scratch = new $T()", BYTES_REF, BYTES_REF);
- }
- builder.beginControlFlow("for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++)");
- {
- builder.addStatement("int groupId = groups.getInt(groupPosition)");
- if (aggState.declaredType().isPrimitive()) {
- if (warnExceptions.isEmpty()) {
- assert intermediateState.size() == 2;
- assert intermediateState.get(1).name().equals("seen");
- builder.beginControlFlow("if (seen.getBoolean(groupPosition + positionOffset))");
- } else {
- assert intermediateState.size() == 3;
- assert intermediateState.get(1).name().equals("seen");
- assert intermediateState.get(2).name().equals("failed");
- builder.beginControlFlow("if (failed.getBoolean(groupPosition + positionOffset))");
- {
- builder.addStatement("state.setFailed(groupId)");
+ var bulkCombineIntermediateMethod = optionalStaticMethod(
+ declarationType,
+ requireVoidType(),
+ requireName("combineIntermediate"),
+ requireArgs(
+ Stream.of(
+ Stream.of(aggState.declaredType(), TypeName.INT, INT_VECTOR), // aggState, positionOffset, groupIds
+ intermediateState.stream().map(AggregatorImplementer.IntermediateStateDesc::combineArgType)
+ ).flatMap(Function.identity()).map(Methods::requireType).toArray(Methods.TypeMatcher[]::new)
+ )
+ );
+ if (bulkCombineIntermediateMethod != null) {
+ var states = intermediateState.stream().map(AggregatorImplementer.IntermediateStateDesc::name).collect(Collectors.joining(","));
+ builder.addStatement("$T.combineIntermediate(state, positionOffset, groups," + states + ")", declarationType);
+ } else {
+ if (intermediateState.stream()
+ .map(AggregatorImplementer.IntermediateStateDesc::elementType)
+ .anyMatch(n -> n.equals("BYTES_REF"))) {
+ builder.addStatement("$T scratch = new $T()", BYTES_REF, BYTES_REF);
+ }
+ builder.beginControlFlow("for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++)");
+ {
+ builder.addStatement("int groupId = groups.getInt(groupPosition)");
+ if (aggState.declaredType().isPrimitive()) {
+ if (warnExceptions.isEmpty()) {
+ assert intermediateState.size() == 2;
+ assert intermediateState.get(1).name().equals("seen");
+ builder.beginControlFlow("if (seen.getBoolean(groupPosition + positionOffset))");
+ } else {
+ assert intermediateState.size() == 3;
+ assert intermediateState.get(1).name().equals("seen");
+ assert intermediateState.get(2).name().equals("failed");
+ builder.beginControlFlow("if (failed.getBoolean(groupPosition + positionOffset))");
+ {
+ builder.addStatement("state.setFailed(groupId)");
+ }
+ builder.nextControlFlow("else if (seen.getBoolean(groupPosition + positionOffset))");
  }
- builder.nextControlFlow("else if (seen.getBoolean(groupPosition + positionOffset))");
- }
 
- warningsBlock(builder, () -> {
- var name = intermediateState.get(0).name();
- var vectorAccessor = vectorAccessorName(intermediateState.get(0).elementType());
- builder.addStatement(
- "state.set(groupId, $T.combine(state.getOrDefault(groupId), $L.$L(groupPosition + positionOffset)))",
+ warningsBlock(builder, () -> {
+ var name = intermediateState.get(0).name();
+ var vectorAccessor = vectorAccessorName(intermediateState.get(0).elementType());
+ builder.addStatement(
+ "state.set(groupId, $T.combine(state.getOrDefault(groupId), $L.$L(groupPosition + positionOffset)))",
+ declarationType,
+ name,
+ vectorAccessor
+ );
+ });
+ builder.endControlFlow();
+ } else {
+ var stateHasBlock = intermediateState.stream().anyMatch(AggregatorImplementer.IntermediateStateDesc::block);
+ requireStaticMethod(
  declarationType,
- name,
- vectorAccessor
+ requireVoidType(),
+ requireName("combineIntermediate"),
+ requireArgs(
+ Stream.of(
+ Stream.of(aggState.declaredType(), TypeName.INT), // aggState and groupId
+ intermediateState.stream().map(AggregatorImplementer.IntermediateStateDesc::combineArgType),
+ Stream.of(TypeName.INT).filter(p -> stateHasBlock) // position
+ ).flatMap(Function.identity()).map(Methods::requireType).toArray(Methods.TypeMatcher[]::new)
+ )
+ );
+ builder.addStatement(
+ "$T.combineIntermediate(state, groupId, "
+ + intermediateState.stream().map(desc -> desc.access("groupPosition + positionOffset")).collect(joining(", "))
+ + (stateHasBlock ? ", groupPosition + positionOffset" : "")
+ + ")",
+ declarationType
  );
- });
+ }
  builder.endControlFlow();
- } else {
- var stateHasBlock = intermediateState.stream().anyMatch(AggregatorImplementer.IntermediateStateDesc::block);
- requireStaticMethod(
- declarationType,
- requireVoidType(),
- requireName("combineIntermediate"),
- requireArgs(
- Stream.of(
- Stream.of(aggState.declaredType(), TypeName.INT), // aggState and groupId
- intermediateState.stream().map(AggregatorImplementer.IntermediateStateDesc::combineArgType),
- Stream.of(TypeName.INT).filter(p -> stateHasBlock) // position
- ).flatMap(Function.identity()).map(Methods::requireType).toArray(Methods.TypeMatcher[]::new)
- )
- );
-
- builder.addStatement(
- "$T.combineIntermediate(state, groupId, "
- + intermediateState.stream().map(desc -> desc.access("groupPosition + positionOffset")).collect(joining(", "))
- + (stateHasBlock ? ", groupPosition + positionOffset" : "")
- + ")",
- declarationType
- );
  }
- builder.endControlFlow();
  }
  return builder.build();
  }

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesBytesRefAggregator.java

@@ -28,15 +28,7 @@ static GroupingAggregatorFunction.AddInput wrapAddInput(
  if (valuesOrdinal == null) {
  return delegate;
  }
- BytesRefVector dict = valuesOrdinal.getDictionaryVector();
- final IntVector hashIds;
- BytesRef spare = new BytesRef();
- try (var hashIdsBuilder = values.blockFactory().newIntVectorFixedBuilder(dict.getPositionCount())) {
- for (int p = 0; p < dict.getPositionCount(); p++) {
- hashIdsBuilder.appendInt(Math.toIntExact(BlockHash.hashOrdToGroup(state.bytes.add(dict.getBytesRef(p, spare)))));
- }
- hashIds = hashIdsBuilder.build();
- }
+ final IntVector hashIds = hashDict(state, valuesOrdinal.getDictionaryVector());
  IntBlock ordinalIds = valuesOrdinal.getOrdinalsBlock();
  return new GroupingAggregatorFunction.AddInput() {
  @Override
@@ -85,17 +77,7 @@ public void add(int positionOffset, IntBigArrayBlock groupIds) {
 
  @Override
  public void add(int positionOffset, IntVector groupIds) {
- for (int groupPosition = 0; groupPosition < groupIds.getPositionCount(); groupPosition++) {
- int groupId = groupIds.getInt(groupPosition);
- if (ordinalIds.isNull(groupPosition + positionOffset)) {
- continue;
- }
- int valuesStart = ordinalIds.getFirstValueIndex(groupPosition + positionOffset);
- int valuesEnd = valuesStart + ordinalIds.getValueCount(groupPosition + positionOffset);
- for (int v = valuesStart; v < valuesEnd; v++) {
- state.addValueOrdinal(groupId, hashIds.getInt(ordinalIds.getInt(v)));
- }
- }
+ addOrdinalInputBlock(state, positionOffset, groupIds, ordinalIds, hashIds);
  }
 
  @Override
@@ -114,15 +96,7 @@ static GroupingAggregatorFunction.AddInput wrapAddInput(
  if (valuesOrdinal == null) {
  return delegate;
  }
- BytesRefVector dict = valuesOrdinal.getDictionaryVector();
- final IntVector hashIds;
- BytesRef spare = new BytesRef();
- try (var hashIdsBuilder = values.blockFactory().newIntVectorFixedBuilder(dict.getPositionCount())) {
- for (int p = 0; p < dict.getPositionCount(); p++) {
- hashIdsBuilder.appendInt(Math.toIntExact(BlockHash.hashOrdToGroup(state.bytes.add(dict.getBytesRef(p, spare)))));
- }
- hashIds = hashIdsBuilder.build();
- }
+ final IntVector hashIds = hashDict(state, valuesOrdinal.getDictionaryVector());
  var ordinalIds = valuesOrdinal.getOrdinalsVector();
  return new GroupingAggregatorFunction.AddInput() {
  @Override
@@ -157,10 +131,7 @@ public void add(int positionOffset, IntBigArrayBlock groupIds) {
 
  @Override
  public void add(int positionOffset, IntVector groupIds) {
- for (int groupPosition = 0; groupPosition < groupIds.getPositionCount(); groupPosition++) {
- int groupId = groupIds.getInt(groupPosition);
- state.addValueOrdinal(groupId, hashIds.getInt(ordinalIds.getInt(groupPosition + positionOffset)));
- }
+ addOrdinalInputVector(state, positionOffset, groupIds, ordinalIds, hashIds);
  }
 
  @Override
@@ -169,4 +140,86 @@ public void close() {
  }
  };
  }
+
+ static IntVector hashDict(ValuesBytesRefAggregator.GroupingState state, BytesRefVector dict) {
+ BytesRef scratch = new BytesRef();
+ try (var hashIdsBuilder = dict.blockFactory().newIntVectorFixedBuilder(dict.getPositionCount())) {
+ for (int p = 0; p < dict.getPositionCount(); p++) {
+ final long hashId = BlockHash.hashOrdToGroup(state.bytes.add(dict.getBytesRef(p, scratch)));
+ hashIdsBuilder.appendInt(Math.toIntExact(hashId));
+ }
+ return hashIdsBuilder.build();
+ }
+ }
+
+ static void addOrdinalInputBlock(
+ ValuesBytesRefAggregator.GroupingState state,
+ int positionOffset,
+ IntVector groupIds,
+ IntBlock ordinalIds,
+ IntVector hashIds
+ ) {
+ for (int p = 0; p < groupIds.getPositionCount(); p++) {
+ final int valuePosition = p + positionOffset;
+ final int groupId = groupIds.getInt(valuePosition);
+ final int start = ordinalIds.getFirstValueIndex(valuePosition);
+ final int end = start + ordinalIds.getValueCount(valuePosition);
+ for (int i = start; i < end; i++) {
+ int ord = ordinalIds.getInt(i);
+ state.addValueOrdinal(groupId, hashIds.getInt(ord));
+ }
+ }
+ }
+
+ static void addOrdinalInputVector(
+ ValuesBytesRefAggregator.GroupingState state,
+ int positionOffset,
+ IntVector groupIds,
+ IntVector ordinalIds,
+ IntVector hashIds
+ ) {
+ for (int p = 0; p < groupIds.getPositionCount(); p++) {
+ int groupId = groupIds.getInt(p);
+ int ord = ordinalIds.getInt(p + positionOffset);
+ state.addValueOrdinal(groupId, hashIds.getInt(ord));
+ }
+ }
+
+ static void combineIntermediateInputValues(
+ ValuesBytesRefAggregator.GroupingState state,
+ int positionOffset,
+ IntVector groupIds,
+ BytesRefBlock values
+ ) {
+ BytesRefVector dict = null;
+ IntBlock ordinals = null;
+ {
+ final OrdinalBytesRefBlock asOrdinals = values.asOrdinals();
+ if (asOrdinals != null) {
+ dict = asOrdinals.getDictionaryVector();
+ ordinals = asOrdinals.getOrdinalsBlock();
+ }
+ }
+ if (dict != null && dict.getPositionCount() < groupIds.getPositionCount()) {
+ try (var hashIds = hashDict(state, dict)) {
+ IntVector ordinalsVector = ordinals.asVector();
+ if (ordinalsVector != null) {
+ addOrdinalInputVector(state, positionOffset, groupIds, ordinalsVector, hashIds);
+ } else {
+ addOrdinalInputBlock(state, positionOffset, groupIds, ordinals, hashIds);
+ }
+ }
+ } else {
+ final BytesRef scratch = new BytesRef();
+ for (int p = 0; p < groupIds.getPositionCount(); p++) {
+ final int valuePosition = p + positionOffset;
+ final int groupId = groupIds.getInt(valuePosition);
+ final int start = values.getFirstValueIndex(valuePosition);
+ final int end = start + values.getValueCount(valuePosition);
+ for (int i = start; i < end; i++) {
+ state.addValue(groupId, values.getBytesRef(i, scratch));
+ }
+ }
+ }
+ }
 }