Examples with AggregationFunctionMetadata io.trino.metadata.AggregationFunctionMetadata used on opensource projects

Example 1 with AggregationFunctionMetadata

use of io.trino.metadata.AggregationFunctionMetadata in project trino by trinodb.

the class PushPartialAggregationThroughExchange method split.

private PlanNode split(AggregationNode node, Context context) {
    // otherwise, add a partial and final with an exchange in between
    Map<Symbol, AggregationNode.Aggregation> intermediateAggregation = new HashMap<>();
    Map<Symbol, AggregationNode.Aggregation> finalAggregation = new HashMap<>();
    for (Map.Entry<Symbol, AggregationNode.Aggregation> entry : node.getAggregations().entrySet()) {
        AggregationNode.Aggregation originalAggregation = entry.getValue();
        ResolvedFunction resolvedFunction = originalAggregation.getResolvedFunction();
        AggregationFunctionMetadata functionMetadata = plannerContext.getMetadata().getAggregationFunctionMetadata(resolvedFunction);
        List<Type> intermediateTypes = functionMetadata.getIntermediateTypes().stream().map(plannerContext.getTypeManager()::getType).collect(toImmutableList());
        Type intermediateType = intermediateTypes.size() == 1 ? intermediateTypes.get(0) : RowType.anonymous(intermediateTypes);
        Symbol intermediateSymbol = context.getSymbolAllocator().newSymbol(resolvedFunction.getSignature().getName(), intermediateType);
        checkState(originalAggregation.getOrderingScheme().isEmpty(), "Aggregate with ORDER BY does not support partial aggregation");
        intermediateAggregation.put(intermediateSymbol, new AggregationNode.Aggregation(resolvedFunction, originalAggregation.getArguments(), originalAggregation.isDistinct(), originalAggregation.getFilter(), originalAggregation.getOrderingScheme(), originalAggregation.getMask()));
        // rewrite final aggregation in terms of intermediate function
        finalAggregation.put(entry.getKey(), new AggregationNode.Aggregation(resolvedFunction, ImmutableList.<Expression>builder().add(intermediateSymbol.toSymbolReference()).addAll(originalAggregation.getArguments().stream().filter(LambdaExpression.class::isInstance).collect(toImmutableList())).build(), false, Optional.empty(), Optional.empty(), Optional.empty()));
    }
    PlanNode partial = new AggregationNode(context.getIdAllocator().getNextId(), node.getSource(), intermediateAggregation, node.getGroupingSets(), // through the exchange may or may not preserve these properties. Hence, it is safest to drop preGroupedSymbols here.
    ImmutableList.of(), PARTIAL, node.getHashSymbol(), node.getGroupIdSymbol());
    return new AggregationNode(node.getId(), partial, finalAggregation, node.getGroupingSets(), // through the exchange may or may not preserve these properties. Hence, it is safest to drop preGroupedSymbols here.
    ImmutableList.of(), FINAL, node.getHashSymbol(), node.getGroupIdSymbol());
}

Example 2 with AggregationFunctionMetadata

use of io.trino.metadata.AggregationFunctionMetadata in project trino by trinodb.

the class TestAnnotationEngineForAggregates method testNotDecomposableAggregationParse.

@Test
public void testNotDecomposableAggregationParse() {
    Signature expectedSignature = new Signature("custom_decomposable_aggregate", DoubleType.DOUBLE.getTypeSignature(), ImmutableList.of(DoubleType.DOUBLE.getTypeSignature()));
    ParametricAggregation aggregation = getOnlyElement(parseFunctionDefinitions(NotDecomposableAggregationFunction.class));
    assertEquals(aggregation.getFunctionMetadata().getDescription(), "Aggregate with Decomposable=false");
    assertTrue(aggregation.getFunctionMetadata().isDeterministic());
    assertEquals(aggregation.getFunctionMetadata().getSignature(), expectedSignature);
    BoundSignature boundSignature = new BoundSignature(aggregation.getFunctionMetadata().getSignature().getName(), DoubleType.DOUBLE, ImmutableList.of(DoubleType.DOUBLE));
    AggregationFunctionMetadata aggregationMetadata = aggregation.getAggregationMetadata();
    assertFalse(aggregationMetadata.isOrderSensitive());
    assertTrue(aggregationMetadata.getIntermediateTypes().isEmpty());
    aggregation.specialize(boundSignature, NO_FUNCTION_DEPENDENCIES);
}

Example 3 with AggregationFunctionMetadata

use of io.trino.metadata.AggregationFunctionMetadata in project trino by trinodb.

the class TestAnnotationEngineForAggregates method testSimpleExplicitSpecializedAggregationParse.

// TODO this is not yet supported
@Test(enabled = false)
public void testSimpleExplicitSpecializedAggregationParse() {
    Signature expectedSignature = new Signature("explicit_specialized_aggregate", ImmutableList.of(typeVariable("T")), ImmutableList.of(), new TypeSignature("T"), ImmutableList.of(new TypeSignature(ARRAY, TypeSignatureParameter.typeParameter(new TypeSignature("T")))), false);
    ParametricAggregation aggregation = getOnlyElement(parseFunctionDefinitions(ExplicitSpecializedAggregationFunction.class));
    assertEquals(aggregation.getFunctionMetadata().getDescription(), "Simple explicit specialized aggregate");
    assertTrue(aggregation.getFunctionMetadata().isDeterministic());
    assertEquals(aggregation.getFunctionMetadata().getSignature(), expectedSignature);
    ParametricImplementationsGroup<AggregationImplementation> implementations = aggregation.getImplementations();
    assertImplementationCount(implementations, 0, 1, 1);
    AggregationImplementation implementation1 = implementations.getSpecializedImplementations().get(0);
    assertTrue(implementation1.hasSpecializedTypeParameters());
    assertFalse(implementation1.hasSpecializedTypeParameters());
    assertEquals(implementation1.getInputParameterKinds(), ImmutableList.of(STATE, INPUT_CHANNEL));
    AggregationImplementation implementation2 = implementations.getSpecializedImplementations().get(1);
    assertTrue(implementation2.hasSpecializedTypeParameters());
    assertFalse(implementation2.hasSpecializedTypeParameters());
    assertEquals(implementation2.getInputParameterKinds(), ImmutableList.of(STATE, INPUT_CHANNEL));
    BoundSignature boundSignature = new BoundSignature(aggregation.getFunctionMetadata().getSignature().getName(), DoubleType.DOUBLE, ImmutableList.of(new ArrayType(DoubleType.DOUBLE)));
    AggregationFunctionMetadata aggregationMetadata = aggregation.getAggregationMetadata();
    assertFalse(aggregationMetadata.isOrderSensitive());
    assertFalse(aggregationMetadata.getIntermediateTypes().isEmpty());
    aggregation.specialize(boundSignature, NO_FUNCTION_DEPENDENCIES);
}

Example 4 with AggregationFunctionMetadata

use of io.trino.metadata.AggregationFunctionMetadata in project trino by trinodb.

the class TestAnnotationEngineForAggregates method testSimpleImplicitSpecializedAggregationParse.

// TODO this is not yet supported
@Test(enabled = false)
public void testSimpleImplicitSpecializedAggregationParse() {
    Signature expectedSignature = new Signature("implicit_specialized_aggregate", ImmutableList.of(typeVariable("T")), ImmutableList.of(), new TypeSignature("T"), ImmutableList.of(new TypeSignature(ARRAY, TypeSignatureParameter.typeParameter(new TypeSignature("T"))), new TypeSignature("T")), false);
    ParametricAggregation aggregation = getOnlyElement(parseFunctionDefinitions(ImplicitSpecializedAggregationFunction.class));
    assertEquals(aggregation.getFunctionMetadata().getDescription(), "Simple implicit specialized aggregate");
    assertTrue(aggregation.getFunctionMetadata().isDeterministic());
    assertEquals(aggregation.getFunctionMetadata().getSignature(), expectedSignature);
    ParametricImplementationsGroup<AggregationImplementation> implementations = aggregation.getImplementations();
    assertImplementationCount(implementations, 0, 0, 2);
    AggregationImplementation implementation1 = implementations.getSpecializedImplementations().get(0);
    assertTrue(implementation1.hasSpecializedTypeParameters());
    assertFalse(implementation1.hasSpecializedTypeParameters());
    assertEquals(implementation1.getInputParameterKinds(), ImmutableList.of(STATE, INPUT_CHANNEL, INPUT_CHANNEL));
    AggregationImplementation implementation2 = implementations.getSpecializedImplementations().get(1);
    assertTrue(implementation2.hasSpecializedTypeParameters());
    assertFalse(implementation2.hasSpecializedTypeParameters());
    assertEquals(implementation2.getInputParameterKinds(), ImmutableList.of(STATE, INPUT_CHANNEL, INPUT_CHANNEL));
    BoundSignature boundSignature = new BoundSignature(aggregation.getFunctionMetadata().getSignature().getName(), DoubleType.DOUBLE, ImmutableList.of(new ArrayType(DoubleType.DOUBLE)));
    AggregationFunctionMetadata aggregationMetadata = aggregation.getAggregationMetadata();
    assertFalse(aggregationMetadata.isOrderSensitive());
    assertFalse(aggregationMetadata.getIntermediateTypes().isEmpty());
    aggregation.specialize(boundSignature, NO_FUNCTION_DEPENDENCIES);
}

Example 5 with AggregationFunctionMetadata

use of io.trino.metadata.AggregationFunctionMetadata in project trino by trinodb.

the class TestAnnotationEngineForAggregates method testNotAnnotatedAggregateStateAggregationParse.

@Test
public void testNotAnnotatedAggregateStateAggregationParse() {
    ParametricAggregation aggregation = getOnlyElement(parseFunctionDefinitions(NotAnnotatedAggregateStateAggregationFunction.class));
    AggregationImplementation implementation = getOnlyElement(aggregation.getImplementations().getExactImplementations().values());
    assertEquals(implementation.getInputParameterKinds(), ImmutableList.of(STATE, INPUT_CHANNEL));
    BoundSignature boundSignature = new BoundSignature(aggregation.getFunctionMetadata().getSignature().getName(), DoubleType.DOUBLE, ImmutableList.of(DoubleType.DOUBLE));
    AggregationFunctionMetadata aggregationMetadata = aggregation.getAggregationMetadata();
    assertFalse(aggregationMetadata.isOrderSensitive());
    assertFalse(aggregationMetadata.getIntermediateTypes().isEmpty());
    aggregation.specialize(boundSignature, NO_FUNCTION_DEPENDENCIES);
}