feat(function): Add array_union_sum aggregation function

presto-docs/src/main/sphinx/functions/aggregate.rst

@@ -71,6 +71,23 @@ General Aggregate Functions
         ) AS t(name, age, gender);
         --['Alice','Bob','Charlie','Lucy']
 
+.. function:: array_union_sum(array(T)) -> array(T)
+
+    Returns the element-wise sum of all input arrays. Arrays of different lengths
+    are padded with zeros for missing positions. All null values in the original
+    arrays are coalesced to 0.
+
+    ::
+
+        SELECT array_union_sum(arrays) AS merged_sum_array
+        FROM (
+        VALUES
+            (ARRAY[1, 2, 3]),
+            (ARRAY[10, 5, 4, 1]),
+            (ARRAY[9, 0, 5, 4])
+        ) AS t(arrays);
+        --[20, 7, 12, 5]
+
 .. function:: avg(x) -> double
 
     Returns the average (arithmetic mean) of all input values.

presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java

@@ -345,6 +345,7 @@
 import static com.facebook.presto.operator.aggregation.AlternativeMinAggregationFunction.ALTERNATIVE_MIN;
 import static com.facebook.presto.operator.aggregation.ArbitraryAggregationFunction.ANY_VALUE_AGGREGATION;
 import static com.facebook.presto.operator.aggregation.ArbitraryAggregationFunction.ARBITRARY_AGGREGATION;
+import static com.facebook.presto.operator.aggregation.ArrayUnionSumAggregation.ARRAY_UNION_SUM;
 import static com.facebook.presto.operator.aggregation.ChecksumAggregationFunction.CHECKSUM_AGGREGATION;
 import static com.facebook.presto.operator.aggregation.CountColumn.COUNT_COLUMN;
 import static com.facebook.presto.operator.aggregation.DecimalAverageAggregation.DECIMAL_AVERAGE_AGGREGATION;
@@ -942,7 +943,7 @@ private List<? extends SqlFunction> getBuiltInFunctions(FunctionsConfig function
                 .function(new ArrayAggregationFunction(functionsConfig.isLegacyArrayAgg(), functionsConfig.getArrayAggGroupImplementation()))
                 .functions(new MapSubscriptOperator(functionsConfig.isLegacyMapSubscript()))
                 .functions(MAP_CONSTRUCTOR, MAP_TO_JSON, JSON_TO_MAP, JSON_STRING_TO_MAP)
-                .functions(MAP_AGG, MAP_UNION, MAP_UNION_SUM)
+                .functions(MAP_AGG, MAP_UNION, MAP_UNION_SUM, ARRAY_UNION_SUM)
                 .function(new ReduceAggregationFunction(functionsConfig.isReduceAggForComplexTypesEnabled()))
                 .functions(DECIMAL_TO_VARCHAR_CAST, DECIMAL_TO_INTEGER_CAST, DECIMAL_TO_BIGINT_CAST, DECIMAL_TO_DOUBLE_CAST, DECIMAL_TO_REAL_CAST, DECIMAL_TO_BOOLEAN_CAST, DECIMAL_TO_TINYINT_CAST, DECIMAL_TO_SMALLINT_CAST)
                 .functions(VARCHAR_TO_DECIMAL_CAST, INTEGER_TO_DECIMAL_CAST, BIGINT_TO_DECIMAL_CAST, DOUBLE_TO_DECIMAL_CAST, REAL_TO_DECIMAL_CAST, BOOLEAN_TO_DECIMAL_CAST, TINYINT_TO_DECIMAL_CAST, SMALLINT_TO_DECIMAL_CAST)

presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/ArrayUnionSumAggregation.java

@@ -0,0 +1,168 @@
+/*
+ * 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;
+
+import com.facebook.presto.bytecode.DynamicClassLoader;
+import com.facebook.presto.common.block.Block;
+import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.type.ArrayType;
+import com.facebook.presto.common.type.Type;
+import com.facebook.presto.common.type.TypeSignatureParameter;
+import com.facebook.presto.metadata.BoundVariables;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
+import com.facebook.presto.metadata.SqlAggregationFunction;
+import com.facebook.presto.operator.aggregation.state.ArrayUnionSumState;
+import com.facebook.presto.operator.aggregation.state.ArrayUnionSumStateFactory;
+import com.facebook.presto.operator.aggregation.state.ArrayUnionSumStateSerializer;
+import com.facebook.presto.spi.function.aggregation.Accumulator;
+import com.facebook.presto.spi.function.aggregation.AggregationMetadata;
+import com.facebook.presto.spi.function.aggregation.AggregationMetadata.AccumulatorStateDescriptor;
+import com.facebook.presto.spi.function.aggregation.GroupedAccumulator;
+import com.google.common.collect.ImmutableList;
+
+import java.lang.invoke.MethodHandle;
+import java.util.List;
+
+import static com.facebook.presto.common.type.StandardTypes.ARRAY;
+import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.generateAggregationName;
+import static com.facebook.presto.spi.function.Signature.nonDecimalNumericTypeParameter;
+import static com.facebook.presto.spi.function.aggregation.AggregationMetadata.ParameterMetadata;
+import static com.facebook.presto.spi.function.aggregation.AggregationMetadata.ParameterMetadata.ParameterType.INPUT_CHANNEL;
+import static com.facebook.presto.spi.function.aggregation.AggregationMetadata.ParameterMetadata.ParameterType.STATE;
+import static com.facebook.presto.util.Reflection.methodHandle;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
+public class ArrayUnionSumAggregation
+        extends SqlAggregationFunction
+{
+    public static final String NAME = "array_union_sum";
+    public static final ArrayUnionSumAggregation ARRAY_UNION_SUM = new ArrayUnionSumAggregation();
+
+    private static final MethodHandle INPUT_FUNCTION = methodHandle(ArrayUnionSumAggregation.class, "input", Type.class, ArrayUnionSumState.class, Block.class);
+    private static final MethodHandle COMBINE_FUNCTION = methodHandle(ArrayUnionSumAggregation.class, "combine", ArrayUnionSumState.class, ArrayUnionSumState.class);
+    private static final MethodHandle OUTPUT_FUNCTION = methodHandle(ArrayUnionSumAggregation.class, "output", ArrayUnionSumState.class, BlockBuilder.class);
+
+    public ArrayUnionSumAggregation()
+    {
+        super(NAME,
+                ImmutableList.of(nonDecimalNumericTypeParameter("T")),
+                ImmutableList.of(),
+                parseTypeSignature("array<T>"),
+                ImmutableList.of(parseTypeSignature("array<T>")));
+    }
+
+    @Override
+    public String getDescription()
+    {
+        return "Aggregate all the arrays into a single array summing the values at each index";
+    }
+
+    @Override
+    public BuiltInAggregationFunctionImplementation specialize(BoundVariables boundVariables, int arity, FunctionAndTypeManager functionAndTypeManager)
+    {
+        Type elementType = boundVariables.getTypeVariable("T");
+        ArrayType outputType = (ArrayType) functionAndTypeManager.getParameterizedType(ARRAY, ImmutableList.of(
+                TypeSignatureParameter.of(elementType.getTypeSignature())));
+
+        return generateAggregation(elementType, outputType);
+    }
+
+    private static BuiltInAggregationFunctionImplementation generateAggregation(Type elementType, ArrayType outputType)
+    {
+        DynamicClassLoader classLoader = new DynamicClassLoader(ArrayUnionSumAggregation.class.getClassLoader());
+        List<Type> inputTypes = ImmutableList.of(outputType);
+        ArrayUnionSumStateSerializer stateSerializer = new ArrayUnionSumStateSerializer(outputType);
+        Type intermediateType = stateSerializer.getSerializedType();
+
+        AggregationMetadata metadata = new AggregationMetadata(
+                generateAggregationName(NAME, outputType.getTypeSignature(), inputTypes.stream().map(Type::getTypeSignature).collect(toImmutableList())),
+                createInputParameterMetadata(outputType),
+                INPUT_FUNCTION.bindTo(elementType),
+                COMBINE_FUNCTION,
+                OUTPUT_FUNCTION,
+                ImmutableList.of(new AccumulatorStateDescriptor(
+                        ArrayUnionSumState.class,
+                        stateSerializer,
+                        new ArrayUnionSumStateFactory(elementType))),
+                outputType);
+
+        Class<? extends Accumulator> accumulatorClass = AccumulatorCompiler.generateAccumulatorClass(
+                Accumulator.class,
+                metadata,
+                classLoader);
+        Class<? extends GroupedAccumulator> groupedAccumulatorClass = AccumulatorCompiler.generateAccumulatorClass(
+                GroupedAccumulator.class,
+                metadata,
+                classLoader);
+        return new BuiltInAggregationFunctionImplementation(NAME, inputTypes, ImmutableList.of(intermediateType), outputType,
+                true, false, metadata, accumulatorClass, groupedAccumulatorClass);
+    }
+
+    private static List<ParameterMetadata> createInputParameterMetadata(Type inputType)
+    {
+        return ImmutableList.of(
+                new ParameterMetadata(STATE),
+                new ParameterMetadata(INPUT_CHANNEL, inputType));
+    }
+
+    public static void input(Type elementType, ArrayUnionSumState state, Block arrayBlock)
+    {
+        if (arrayBlock == null) {
+            return;
+        }
+
+        ArrayUnionSumResult arrayUnionSumResult = state.get();
+        long startSize;
+
+        if (arrayUnionSumResult == null) {
+            startSize = 0;
+            arrayUnionSumResult = ArrayUnionSumResult.create(elementType, state.getAdder(), arrayBlock);
+            state.set(arrayUnionSumResult);
+        }
+        else {
+            startSize = arrayUnionSumResult.getRetainedSizeInBytes();
+            arrayUnionSumResult = arrayUnionSumResult.unionSum(arrayBlock);
+            state.set(arrayUnionSumResult);
+        }
+
+        state.addMemoryUsage(arrayUnionSumResult.getRetainedSizeInBytes() - startSize);
+    }
+
+    public static void combine(ArrayUnionSumState state, ArrayUnionSumState otherState)
+    {
+        if (state.get() == null) {
+            state.set(otherState.get());
+            if (otherState.get() != null) {
+                state.addMemoryUsage(otherState.get().getRetainedSizeInBytes());
+            }
+            return;
+        }
+
+        long startSize = state.get().getRetainedSizeInBytes();
+        state.set(state.get().unionSum(otherState.get()));
+        state.addMemoryUsage(state.get().getRetainedSizeInBytes() - startSize);
+    }
+
+    public static void output(ArrayUnionSumState state, BlockBuilder out)
+    {
+        ArrayUnionSumResult arrayUnionSumResult = state.get();
+        if (arrayUnionSumResult == null) {
+            out.appendNull();
+        }
+        else {
+            arrayUnionSumResult.serialize(out);
+        }
+    }
+}