diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/LossySumDoubleGroupingAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/LossySumDoubleGroupingAggregatorFunction.java index bfe214035260e..54127f2b05b59 100644 --- a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/LossySumDoubleGroupingAggregatorFunction.java +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/LossySumDoubleGroupingAggregatorFunction.java @@ -86,7 +86,7 @@ public void close() { } }; } - return new GroupingAggregatorFunction.AddInput() { + var addInput = new GroupingAggregatorFunction.AddInput() { @Override public void add(int positionOffset, IntArrayBlock groupIds) { addRawInput(positionOffset, groupIds, vVector); @@ -106,6 +106,7 @@ public void add(int positionOffset, IntVector groupIds) { public void close() { } }; + return LossySumDoubleAggregator.wrapAddInput(addInput, state, vVector); } private void addRawInput(int positionOffset, IntArrayBlock groups, DoubleBlock vBlock) { diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/LossySumDoubleAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/LossySumDoubleAggregator.java index 7a041169d4c0f..b1031ec2759ee 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/LossySumDoubleAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/LossySumDoubleAggregator.java @@ -15,6 +15,9 @@ import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.compute.data.DoubleBlock; +import org.elasticsearch.compute.data.DoubleVector; +import org.elasticsearch.compute.data.IntArrayBlock; +import org.elasticsearch.compute.data.IntBigArrayBlock; import org.elasticsearch.compute.data.IntVector; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.core.Releasables; @@ -121,6 +124,43 @@ public static Block evaluateFinal(GroupingSumState state, IntVector selected, Gr } } + public static GroupingAggregatorFunction.AddInput wrapAddInput( + GroupingAggregatorFunction.AddInput delegate, + GroupingSumState state, + DoubleVector values + ) { + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntArrayBlock groupIds) { + delegate.add(positionOffset, groupIds); + } + + @Override + public void add(int positionOffset, IntBigArrayBlock groupIds) { + delegate.add(positionOffset, groupIds); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + if (groupIds.isConstant()) { + double sum = 0.0; + int positionCount = groupIds.getPositionCount(); + for (int i = 0; i < positionCount; i++) { + sum += values.getDouble(i); + } + state.add(sum, groupIds.getInt(0)); + } else { + delegate.add(positionOffset, groupIds); + } + } + + @Override + public void close() { + Releasables.close(delegate); + } + }; + } + static final class SumState implements AggregatorState { private boolean seen; double value; @@ -149,7 +189,7 @@ static final class GroupingSumState extends AbstractArrayState implements Groupi super(bigArrays); boolean success = false; try { - this.values = bigArrays.newDoubleArray(1); + this.values = bigArrays.newDoubleArray(128); success = true; } finally { if (success == false) {