Simplify parametric aggregation

Consolidate the two annotation processing strategy into a single strategy Simplify handling of function aliases Move deprecated flag to AggregationHeader
trinodb · Oct 9, 2021 · 5f3966f · 5f3966f
1 parent 270e880
commit 5f3966f
Show file tree

Hide file tree

Showing 11 changed files with 222 additions and 201 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/metadata/FunctionListBuilder.java b/core/trino-main/src/main/java/io/trino/metadata/FunctionListBuilder.java
@@ -33,12 +33,6 @@ public FunctionListBuilder window(Class<? extends WindowFunction> clazz)
         return this;
     }
 
-    public FunctionListBuilder aggregate(Class<?> aggregationDefinition)
-    {
-        functions.addAll(SqlAggregationFunction.createFunctionByAnnotations(aggregationDefinition));
-        return this;
-    }
-
     public FunctionListBuilder aggregates(Class<?> aggregationDefinition)
     {
         functions.addAll(SqlAggregationFunction.createFunctionsByAnnotations(aggregationDefinition));

diff --git a/core/trino-main/src/main/java/io/trino/metadata/FunctionRegistry.java b/core/trino-main/src/main/java/io/trino/metadata/FunctionRegistry.java
@@ -423,12 +423,12 @@ public FunctionRegistry(
                 .window(NthValueFunction.class)
                 .window(LagFunction.class)
                 .window(LeadFunction.class)
-                .aggregate(ApproximateCountDistinctAggregation.class)
-                .aggregate(DefaultApproximateCountDistinctAggregation.class)
-                .aggregate(BooleanApproximateCountDistinctAggregation.class)
-                .aggregate(BooleanDefaultApproximateCountDistinctAggregation.class)
-                .aggregate(SumDataSizeForStats.class)
-                .aggregate(MaxDataSizeForStats.class)
+                .aggregates(ApproximateCountDistinctAggregation.class)
+                .aggregates(DefaultApproximateCountDistinctAggregation.class)
+                .aggregates(BooleanApproximateCountDistinctAggregation.class)
+                .aggregates(BooleanDefaultApproximateCountDistinctAggregation.class)
+                .aggregates(SumDataSizeForStats.class)
+                .aggregates(MaxDataSizeForStats.class)
                 .aggregates(CountAggregation.class)
                 .aggregates(VarianceAggregation.class)
                 .aggregates(CentralMomentsAggregation.class)
@@ -457,7 +457,7 @@ public FunctionRegistry(
                 .aggregates(RealGeometricMeanAggregations.class)
                 .aggregates(MergeHyperLogLogAggregation.class)
                 .aggregates(ApproximateSetAggregation.class)
-                .aggregate(ApproximateSetGenericAggregation.class)
+                .aggregates(ApproximateSetGenericAggregation.class)
                 .aggregates(TDigestAggregationFunction.class)
                 .functions(QDIGEST_AGG, QDIGEST_AGG_WITH_WEIGHT, QDIGEST_AGG_WITH_WEIGHT_AND_ERROR)
                 .function(MergeQuantileDigestFunction.MERGE)
@@ -621,13 +621,13 @@ public FunctionRegistry(
                 .function(new GenericLessThanOperator(typeOperators))
                 .function(new GenericLessThanOrEqualOperator(typeOperators))
                 .function(new VersionFunction(nodeVersion))
-                .aggregate(MergeSetDigestAggregation.class)
-                .aggregate(BuildSetDigestAggregation.class)
+                .aggregates(MergeSetDigestAggregation.class)
+                .aggregates(BuildSetDigestAggregation.class)
                 .scalars(SetDigestFunctions.class)
                 .scalars(SetDigestOperators.class)
                 .scalars(WilsonInterval.class)
-                .aggregate(BigintApproximateMostFrequent.class)
-                .aggregate(VarcharApproximateMostFrequent.class);
+                .aggregates(BigintApproximateMostFrequent.class)
+                .aggregates(VarcharApproximateMostFrequent.class);
 
         // timestamp operators and functions
         builder

diff --git a/core/trino-main/src/main/java/io/trino/metadata/SqlAggregationFunction.java b/core/trino-main/src/main/java/io/trino/metadata/SqlAggregationFunction.java
@@ -35,17 +35,9 @@ public abstract class SqlAggregationFunction
     private final boolean orderSensitive;
     private final boolean decomposable;
 
-    public static List<SqlAggregationFunction> createFunctionByAnnotations(Class<?> aggregationDefinition)
-    {
-        return ImmutableList.of(AggregationFromAnnotationsParser.parseFunctionDefinition(aggregationDefinition));
-    }
-
     public static List<SqlAggregationFunction> createFunctionsByAnnotations(Class<?> aggregationDefinition)
     {
-        return AggregationFromAnnotationsParser.parseFunctionDefinitions(aggregationDefinition)
-                .stream()
-                .map(SqlAggregationFunction.class::cast)
-                .collect(toImmutableList());
+        return ImmutableList.copyOf(AggregationFromAnnotationsParser.parseFunctionDefinitions(aggregationDefinition));
     }
 
     protected SqlAggregationFunction(FunctionMetadata functionMetadata, boolean decomposable, boolean orderSensitive)

diff --git a/core/trino-main/src/main/java/io/trino/operator/ParametricImplementation.java b/core/trino-main/src/main/java/io/trino/operator/ParametricImplementation.java
@@ -27,4 +27,6 @@ public interface ParametricImplementation
     boolean isNullable();
 
     List<FunctionArgumentDefinition> getArgumentDefinitions();
+
+    ParametricImplementation withAlias(String alias);
 }
diff --git a/core/trino-main/src/main/java/io/trino/operator/ParametricImplementationsGroup.java b/core/trino-main/src/main/java/io/trino/operator/ParametricImplementationsGroup.java
@@ -22,9 +22,13 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.Iterables.getOnlyElement;
+import static io.trino.operator.ParametricFunctionHelpers.signatureWithName;
 import static io.trino.operator.annotations.FunctionsParserHelper.validateSignaturesCompatibility;
 import static java.util.Objects.requireNonNull;
 
@@ -85,6 +89,7 @@ public List<FunctionArgumentDefinition> getArgumentDefinitions()
         return argumentDefinitions;
     }
 
+    @SafeVarargs
     public static <T extends ParametricImplementation> ParametricImplementationsGroup<T> of(T... implementations)
     {
         ParametricImplementationsGroup.Builder<T> builder = builder();
@@ -114,6 +119,30 @@ public Signature getSignature()
         return signature;
     }
 
+    public ParametricImplementationsGroup<T> withAlias(String alias)
+    {
+        if (alias.equals(signature.getName())) {
+            return this;
+        }
+        return new ParametricImplementationsGroup<>(
+                exactImplementations.values().stream()
+                        .map(implementation -> withAlias(alias, implementation))
+                        .collect(toImmutableMap(T::getSignature, Function.identity())),
+                specializedImplementations.stream()
+                        .map(implementation -> withAlias(alias, implementation))
+                        .collect(toImmutableList()),
+                genericImplementations.stream()
+                        .map(implementation -> withAlias(alias, implementation))
+                        .collect(toImmutableList()),
+                signatureWithName(alias, signature));
+    }
+
+    @SuppressWarnings("unchecked")
+    private static <T extends ParametricImplementation> T withAlias(String name, T implementation)
+    {
+        return (T) implementation.withAlias(name);
+    }
+
     public static <T extends ParametricImplementation> Builder<T> builder()
     {
         return new Builder<>();
@@ -174,7 +203,7 @@ private static <T extends ParametricImplementation> Signature determineGenericSi
                 signature = Optional.of(implementation.getSignature());
             }
 
-            return signature.get();
+            return signature.orElseThrow();
         }
     }
 }