|  | 
| 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