|
15 | 15 | import org.elasticsearch.compute.data.Block;
|
16 | 16 | import org.elasticsearch.compute.data.BlockFactory;
|
17 | 17 | import org.elasticsearch.compute.data.DoubleBlock;
|
| 18 | +import org.elasticsearch.compute.data.DoubleVector; |
| 19 | +import org.elasticsearch.compute.data.IntArrayBlock; |
| 20 | +import org.elasticsearch.compute.data.IntBigArrayBlock; |
18 | 21 | import org.elasticsearch.compute.data.IntVector;
|
19 | 22 | import org.elasticsearch.compute.operator.DriverContext;
|
20 | 23 | import org.elasticsearch.core.Releasables;
|
@@ -121,6 +124,43 @@ public static Block evaluateFinal(GroupingSumState state, IntVector selected, Gr
|
121 | 124 | }
|
122 | 125 | }
|
123 | 126 |
|
| 127 | + public static GroupingAggregatorFunction.AddInput wrapAddInput( |
| 128 | + GroupingAggregatorFunction.AddInput delegate, |
| 129 | + GroupingSumState state, |
| 130 | + DoubleVector values |
| 131 | + ) { |
| 132 | + return new GroupingAggregatorFunction.AddInput() { |
| 133 | + @Override |
| 134 | + public void add(int positionOffset, IntArrayBlock groupIds) { |
| 135 | + delegate.add(positionOffset, groupIds); |
| 136 | + } |
| 137 | + |
| 138 | + @Override |
| 139 | + public void add(int positionOffset, IntBigArrayBlock groupIds) { |
| 140 | + delegate.add(positionOffset, groupIds); |
| 141 | + } |
| 142 | + |
| 143 | + @Override |
| 144 | + public void add(int positionOffset, IntVector groupIds) { |
| 145 | + if (groupIds.isConstant()) { |
| 146 | + double sum = 0.0; |
| 147 | + int positionCount = groupIds.getPositionCount(); |
| 148 | + for (int i = 0; i < positionCount; i++) { |
| 149 | + sum += values.getDouble(i); |
| 150 | + } |
| 151 | + state.add(sum, groupIds.getInt(0)); |
| 152 | + } else { |
| 153 | + delegate.add(positionOffset, groupIds); |
| 154 | + } |
| 155 | + } |
| 156 | + |
| 157 | + @Override |
| 158 | + public void close() { |
| 159 | + Releasables.close(delegate); |
| 160 | + } |
| 161 | + }; |
| 162 | + } |
| 163 | + |
124 | 164 | static final class SumState implements AggregatorState {
|
125 | 165 | private boolean seen;
|
126 | 166 | double value;
|
@@ -149,7 +189,7 @@ static final class GroupingSumState extends AbstractArrayState implements Groupi
|
149 | 189 | super(bigArrays);
|
150 | 190 | boolean success = false;
|
151 | 191 | try {
|
152 |
| - this.values = bigArrays.newDoubleArray(1); |
| 192 | + this.values = bigArrays.newDoubleArray(128); |
153 | 193 | success = true;
|
154 | 194 | } finally {
|
155 | 195 | if (success == false) {
|
|
0 commit comments