From d63db3bebfe57ac199eb82cc19cb0193026808a2 Mon Sep 17 00:00:00 2001 From: Dain Sundstrom Date: Wed, 29 Dec 2021 13:40:45 -0800 Subject: [PATCH] Add support for specialized BlockPosition in annotated aggergations --- .../aggregation/AggregationImplementation.java | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/AggregationImplementation.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/AggregationImplementation.java index ce633d00cb4e..20542d890fd8 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/AggregationImplementation.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/AggregationImplementation.java @@ -217,10 +217,12 @@ public boolean areTypesAssignable(BoundSignature boundSignature) Class methodDeclaredType = argumentNativeContainerTypes.get(i).getJavaType(); boolean isCurrentBlockPosition = argumentNativeContainerTypes.get(i).isBlockPosition(); + // block and position works for any type, but if block is annotated with SqlType nativeContainerType, then only types with the + // specified container type match if (isCurrentBlockPosition && methodDeclaredType.isAssignableFrom(Block.class)) { continue; } - if (!isCurrentBlockPosition && methodDeclaredType.isAssignableFrom(argumentType)) { + if (methodDeclaredType.isAssignableFrom(argumentType)) { continue; } return false; @@ -446,7 +448,16 @@ public static List parseSignatureArgumentsTypes(Me continue; } - builder.add(new AggregateNativeContainerType(inputFunction.getParameterTypes()[i], isParameterBlock(annotations))); + Optional> nativeContainerType = Arrays.stream(annotations) + .filter(SqlType.class::isInstance) + .map(SqlType.class::cast) + .findFirst() + .map(SqlType::nativeContainerType); + // Note: this cannot be done as a chain due to strange generic type mismatches + if (nativeContainerType.isPresent() && !nativeContainerType.get().equals(Object.class)) { + parameterType = nativeContainerType.get(); + } + builder.add(new AggregateNativeContainerType(parameterType, isParameterBlock(annotations))); } return builder.build();