Migrate Real Average Aggregation to use multiple aggregation state

presto-main/src/main/java/com/facebook/presto/metadata/FunctionRegistry.java

@@ -45,7 +45,6 @@
 import com.facebook.presto.operator.aggregation.LongSumAggregation;
 import com.facebook.presto.operator.aggregation.MaxDataSizeForStats;
 import com.facebook.presto.operator.aggregation.MergeHyperLogLogAggregation;
-import com.facebook.presto.operator.aggregation.RealAverageAggregation;
 import com.facebook.presto.operator.aggregation.RealCorrelationAggregation;
 import com.facebook.presto.operator.aggregation.RealCovarianceAggregation;
 import com.facebook.presto.operator.aggregation.RealGeometricMeanAggregations;
@@ -225,6 +224,7 @@
 import static com.facebook.presto.operator.aggregation.MaxNAggregationFunction.MAX_N_AGGREGATION;
 import static com.facebook.presto.operator.aggregation.MinAggregationFunction.MIN_AGGREGATION;
 import static com.facebook.presto.operator.aggregation.MinNAggregationFunction.MIN_N_AGGREGATION;
+import static com.facebook.presto.operator.aggregation.RealAverageAggregation.REAL_AVERAGE_AGGREGATION;
 import static com.facebook.presto.operator.aggregation.minmaxby.MaxByAggregationFunction.MAX_BY;
 import static com.facebook.presto.operator.aggregation.minmaxby.MaxByNAggregationFunction.MAX_BY_N_AGGREGATION;
 import static com.facebook.presto.operator.aggregation.minmaxby.MinByAggregationFunction.MIN_BY;
@@ -440,7 +440,7 @@ public FunctionRegistry(TypeManager typeManager, BlockEncodingSerde blockEncodin
                 .aggregates(IntervalDayToSecondSumAggregation.class)
                 .aggregates(IntervalYearToMonthSumAggregation.class)
                 .aggregates(AverageAggregations.class)
-                .aggregates(RealAverageAggregation.class)
+                .function(REAL_AVERAGE_AGGREGATION)
                 .aggregates(IntervalDayToSecondAverageAggregation.class)
                 .aggregates(IntervalYearToMonthAverageAggregation.class)
                 .aggregates(GeometricMeanAggregations.class)

presto-main/src/main/java/com/facebook/presto/operator/aggregation/RealAverageAggregation.java

@@ -13,49 +13,125 @@
  */
 package com.facebook.presto.operator.aggregation;
 
-import com.facebook.presto.operator.aggregation.state.LongAndDoubleState;
+import com.facebook.presto.metadata.BoundVariables;
+import com.facebook.presto.metadata.FunctionRegistry;
+import com.facebook.presto.metadata.SqlAggregationFunction;
+import com.facebook.presto.operator.aggregation.AggregationMetadata.AccumulatorStateDescriptor;
+import com.facebook.presto.operator.aggregation.state.DoubleState;
+import com.facebook.presto.operator.aggregation.state.LongState;
+import com.facebook.presto.operator.aggregation.state.StateCompiler;
 import com.facebook.presto.spi.block.BlockBuilder;
-import com.facebook.presto.spi.function.AggregationFunction;
-import com.facebook.presto.spi.function.AggregationState;
-import com.facebook.presto.spi.function.CombineFunction;
-import com.facebook.presto.spi.function.InputFunction;
-import com.facebook.presto.spi.function.OutputFunction;
-import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.AccumulatorState;
+import com.facebook.presto.spi.function.AccumulatorStateSerializer;
 import com.facebook.presto.spi.type.StandardTypes;
+import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.spi.type.TypeManager;
+import com.google.common.collect.ImmutableList;
+import io.airlift.bytecode.DynamicClassLoader;
 
+import java.lang.invoke.MethodHandle;
+import java.util.List;
+
+import static com.facebook.presto.operator.aggregation.AggregationMetadata.ParameterMetadata;
+import static com.facebook.presto.operator.aggregation.AggregationMetadata.ParameterMetadata.ParameterType.BLOCK_INDEX;
+import static com.facebook.presto.operator.aggregation.AggregationMetadata.ParameterMetadata.ParameterType.BLOCK_INPUT_CHANNEL;
+import static com.facebook.presto.operator.aggregation.AggregationMetadata.ParameterMetadata.ParameterType.INPUT_CHANNEL;
+import static com.facebook.presto.operator.aggregation.AggregationMetadata.ParameterMetadata.ParameterType.STATE;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.generateAggregationName;
 import static com.facebook.presto.spi.type.RealType.REAL;
-import static java.lang.Float.floatToRawIntBits;
+import static com.facebook.presto.spi.type.TypeSignature.parseTypeSignature;
+import static com.facebook.presto.util.Reflection.methodHandle;
+import static java.lang.Float.floatToIntBits;
 import static java.lang.Float.intBitsToFloat;
 
-@AggregationFunction("avg")
-public final class RealAverageAggregation
+public class RealAverageAggregation
+        extends SqlAggregationFunction
 {
-    private RealAverageAggregation() {}
+    public static final RealAverageAggregation REAL_AVERAGE_AGGREGATION = new RealAverageAggregation();
+    private static final String NAME = "avg";
+
+    private static final MethodHandle INPUT_FUNCTION = methodHandle(RealAverageAggregation.class, "input", LongState.class, DoubleState.class, long.class);
+    private static final MethodHandle COMBINE_FUNCTION = methodHandle(RealAverageAggregation.class, "combine", LongState.class, DoubleState.class, LongState.class, DoubleState.class);
+    private static final MethodHandle OUTPUT_FUNCTION = methodHandle(RealAverageAggregation.class, "output", LongState.class, DoubleState.class, BlockBuilder.class);
+
+    protected RealAverageAggregation()
+    {
+        super(NAME,
+                ImmutableList.of(),
+                ImmutableList.of(),
+                parseTypeSignature(StandardTypes.REAL),
+                ImmutableList.of(parseTypeSignature(StandardTypes.REAL)));
+    }
+
+    @Override
+    public String getDescription()
+    {
+        return "Returns the average value of the argument";
+    }
+
+    @Override
+    public InternalAggregationFunction specialize(BoundVariables boundVariables, int arity, TypeManager typeManager, FunctionRegistry functionRegistry)
+    {
+        DynamicClassLoader classLoader = new DynamicClassLoader(AverageAggregations.class.getClassLoader());
+        Class<? extends AccumulatorState> longStateInterface = LongState.class;
+        Class<? extends AccumulatorState> doubleStateInterface = DoubleState.class;
+        AccumulatorStateSerializer<?> longStateSerializer = StateCompiler.generateStateSerializer(longStateInterface, classLoader);
+        AccumulatorStateSerializer<?> doubleStateSerializer = StateCompiler.generateStateSerializer(doubleStateInterface, classLoader);
+
+        AggregationMetadata metadata = new AggregationMetadata(
+                generateAggregationName(NAME, parseTypeSignature(StandardTypes.REAL), ImmutableList.of(parseTypeSignature(StandardTypes.REAL))),
+                ImmutableList.of(new ParameterMetadata(STATE), new ParameterMetadata(STATE), new ParameterMetadata(INPUT_CHANNEL, REAL)),
+                INPUT_FUNCTION,
+                COMBINE_FUNCTION,
+                OUTPUT_FUNCTION,
+                ImmutableList.of(
+                    new AccumulatorStateDescriptor(
+                            longStateInterface,
+                            longStateSerializer,
+                            StateCompiler.generateStateFactory(longStateInterface, classLoader)),
+                    new AccumulatorStateDescriptor(
+                            doubleStateInterface,
+                            doubleStateSerializer,
+                            StateCompiler.generateStateFactory(doubleStateInterface, classLoader))),
+                REAL);
+
+        GenericAccumulatorFactoryBinder factory = AccumulatorCompiler.generateAccumulatorFactoryBinder(metadata, classLoader);
+        return new InternalAggregationFunction(
+                NAME,
+                ImmutableList.of(REAL),
+                ImmutableList.of(
+                        longStateSerializer.getSerializedType(),
+                        doubleStateSerializer.getSerializedType()),
+                REAL,
+                true,
+                false,
+                factory);
+    }
+
+    private static List<ParameterMetadata> createInputParameterMetadata(Type value)
+    {
+        return ImmutableList.of(new ParameterMetadata(STATE), new ParameterMetadata(BLOCK_INPUT_CHANNEL, value), new ParameterMetadata(BLOCK_INDEX));
+    }
 
-    @InputFunction
-    public static void input(@AggregationState LongAndDoubleState state, @SqlType(StandardTypes.REAL) long value)
+    public static void input(LongState count, DoubleState sum, long value)
     {
-        state.setLong(state.getLong() + 1);
-        state.setDouble(state.getDouble() + intBitsToFloat((int) value));
+        count.setLong(count.getLong() + 1);
+        sum.setDouble(sum.getDouble() + intBitsToFloat((int) value));
     }
 
-    @CombineFunction
-    public static void combine(@AggregationState LongAndDoubleState state, LongAndDoubleState otherState)
+    public static void combine(LongState count, DoubleState sum, LongState otherCount, DoubleState otherSum)
     {
-        state.setLong(state.getLong() + otherState.getLong());
-        state.setDouble(state.getDouble() + otherState.getDouble());
+        count.setLong(count.getLong() + otherCount.getLong());
+        sum.setDouble(sum.getDouble() + otherSum.getDouble());
     }
 
-    @OutputFunction(StandardTypes.REAL)
-    public static void output(@AggregationState LongAndDoubleState state, BlockBuilder out)
+    public static void output(LongState count, DoubleState sum, BlockBuilder out)
     {
-        long count = state.getLong();
-        if (count == 0) {
+        if (count.getLong() == 0) {
             out.appendNull();
         }
         else {
-            double average = state.getDouble() / count;
-            REAL.writeLong(out, floatToRawIntBits((float) average));
+            REAL.writeLong(out, floatToIntBits((float) (sum.getDouble() / count.getLong())));
         }
     }
 }

presto-main/src/main/java/com/facebook/presto/operator/aggregation/state/DoubleState.java

@@ -0,0 +1,24 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.operator.aggregation.state;
+
+import com.facebook.presto.spi.function.AccumulatorState;
+
+public interface DoubleState
+        extends AccumulatorState
+{
+    double getDouble();
+
+    void setDouble(double value);
+}