Examples with Variable io.trino.spi.expression.Variable used on opensource projects

Example 46 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class ImplementSum method rewrite.

@Override
public Optional<AggregateExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<Void> context) {
    Variable argument = captures.get(ARGUMENT);
    PinotColumnHandle columnHandle = (PinotColumnHandle) context.getAssignment(argument.getName());
    return Optional.of(new AggregateExpression(aggregateFunction.getFunctionName(), identifierQuote.apply(columnHandle.getColumnName()), true));
}

Example 47 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class PinotMetadata method applyAggregation.

@Override
public Optional<AggregationApplicationResult<ConnectorTableHandle>> applyAggregation(ConnectorSession session, ConnectorTableHandle handle, List<AggregateFunction> aggregates, Map<String, ColumnHandle> assignments, List<List<ColumnHandle>> groupingSets) {
    if (!isAggregationPushdownEnabled(session)) {
        return Optional.empty();
    }
    // Global aggregation is represented by [[]]
    verify(!groupingSets.isEmpty(), "No grouping sets provided");
    // Pinot currently only supports simple GROUP BY clauses with a single grouping set
    if (groupingSets.size() != 1) {
        return Optional.empty();
    }
    // See https://github.com/apache/pinot/issues/8353 for more details.
    if (getOnlyElement(groupingSets).stream().filter(columnHandle -> ((PinotColumnHandle) columnHandle).getDataType() instanceof ArrayType).findFirst().isPresent()) {
        return Optional.empty();
    }
    PinotTableHandle tableHandle = (PinotTableHandle) handle;
    // If there is an offset then do not push the aggregation down as the results will not be correct
    if (tableHandle.getQuery().isPresent() && (!tableHandle.getQuery().get().getAggregateColumns().isEmpty() || tableHandle.getQuery().get().isAggregateInProjections() || tableHandle.getQuery().get().getOffset().isPresent())) {
        return Optional.empty();
    }
    ImmutableList.Builder<ConnectorExpression> projections = ImmutableList.builder();
    ImmutableList.Builder<Assignment> resultAssignments = ImmutableList.builder();
    ImmutableList.Builder<PinotColumnHandle> aggregateColumnsBuilder = ImmutableList.builder();
    for (AggregateFunction aggregate : aggregates) {
        Optional<AggregateExpression> rewriteResult = aggregateFunctionRewriter.rewrite(session, aggregate, assignments);
        rewriteResult = applyCountDistinct(session, aggregate, assignments, tableHandle, rewriteResult);
        if (rewriteResult.isEmpty()) {
            return Optional.empty();
        }
        AggregateExpression aggregateExpression = rewriteResult.get();
        PinotColumnHandle pinotColumnHandle = new PinotColumnHandle(aggregateExpression.toFieldName(), aggregate.getOutputType(), aggregateExpression.toExpression(), false, true, aggregateExpression.isReturnNullOnEmptyGroup(), Optional.of(aggregateExpression.getFunction()), Optional.of(aggregateExpression.getArgument()));
        aggregateColumnsBuilder.add(pinotColumnHandle);
        projections.add(new Variable(pinotColumnHandle.getColumnName(), pinotColumnHandle.getDataType()));
        resultAssignments.add(new Assignment(pinotColumnHandle.getColumnName(), pinotColumnHandle, pinotColumnHandle.getDataType()));
    }
    List<PinotColumnHandle> groupingColumns = getOnlyElement(groupingSets).stream().map(PinotColumnHandle.class::cast).map(PinotColumnHandle::fromNonAggregateColumnHandle).collect(toImmutableList());
    OptionalLong limitForDynamicTable = OptionalLong.empty();
    // know when the limit was exceeded and throw an error
    if (tableHandle.getLimit().isEmpty() && !groupingColumns.isEmpty()) {
        limitForDynamicTable = OptionalLong.of(maxRowsPerBrokerQuery + 1);
    }
    List<PinotColumnHandle> aggregationColumns = aggregateColumnsBuilder.build();
    String newQuery = "";
    List<PinotColumnHandle> newSelections = groupingColumns;
    if (tableHandle.getQuery().isPresent()) {
        newQuery = tableHandle.getQuery().get().getQuery();
        Map<String, PinotColumnHandle> projectionsMap = tableHandle.getQuery().get().getProjections().stream().collect(toImmutableMap(PinotColumnHandle::getColumnName, identity()));
        groupingColumns = groupingColumns.stream().map(groupIngColumn -> projectionsMap.getOrDefault(groupIngColumn.getColumnName(), groupIngColumn)).collect(toImmutableList());
        ImmutableList.Builder<PinotColumnHandle> newSelectionsBuilder = ImmutableList.<PinotColumnHandle>builder().addAll(groupingColumns);
        aggregationColumns = aggregationColumns.stream().map(aggregateExpression -> resolveAggregateExpressionWithAlias(aggregateExpression, projectionsMap)).collect(toImmutableList());
        newSelections = newSelectionsBuilder.build();
    }
    DynamicTable dynamicTable = new DynamicTable(tableHandle.getTableName(), Optional.empty(), newSelections, tableHandle.getQuery().flatMap(DynamicTable::getFilter), groupingColumns, aggregationColumns, ImmutableList.of(), limitForDynamicTable, OptionalLong.empty(), newQuery);
    tableHandle = new PinotTableHandle(tableHandle.getSchemaName(), tableHandle.getTableName(), tableHandle.getConstraint(), tableHandle.getLimit(), Optional.of(dynamicTable));
    return Optional.of(new AggregationApplicationResult<>(tableHandle, projections.build(), resultAssignments.build(), ImmutableMap.of(), false));
}

Example 48 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class ImplementSqlServerStdev method rewrite.

@Override
public Optional<JdbcExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<String> context) {
    Variable argument = captures.get(ARGUMENT);
    JdbcColumnHandle columnHandle = (JdbcColumnHandle) context.getAssignment(argument.getName());
    verify(columnHandle.getColumnType().equals(DOUBLE));
    verify(aggregateFunction.getOutputType().equals(DOUBLE));
    return Optional.of(new JdbcExpression(format("STDEV(%s)", context.rewriteExpression(argument).orElseThrow()), columnHandle.getJdbcTypeHandle()));
}

Example 49 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class ImplementSqlServerVariance method rewrite.

@Override
public Optional<JdbcExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<String> context) {
    Variable argument = captures.get(ARGUMENT);
    JdbcColumnHandle columnHandle = (JdbcColumnHandle) context.getAssignment(argument.getName());
    verify(columnHandle.getColumnType().equals(DOUBLE));
    verify(aggregateFunction.getOutputType().equals(DOUBLE));
    return Optional.of(new JdbcExpression(format("VAR(%s)", context.rewriteExpression(argument).orElseThrow()), columnHandle.getJdbcTypeHandle()));
}

Example 50 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class TestSqlServerClient method testImplementCount.

@Test
public void testImplementCount() {
    Variable bigintVariable = new Variable("v_bigint", BIGINT);
    Variable doubleVariable = new Variable("v_double", BIGINT);
    Optional<ConnectorExpression> filter = Optional.of(new Variable("a_filter", BOOLEAN));
    // count(*)
    testImplementAggregation(new AggregateFunction("count", BIGINT, List.of(), List.of(), false, Optional.empty()), Map.of(), Optional.of("count_big(*)"));
    // count(bigint)
    testImplementAggregation(new AggregateFunction("count", BIGINT, List.of(bigintVariable), List.of(), false, Optional.empty()), Map.of(bigintVariable.getName(), BIGINT_COLUMN), Optional.of("count_big(\"c_bigint\")"));
    // count(double)
    testImplementAggregation(new AggregateFunction("count", BIGINT, List.of(doubleVariable), List.of(), false, Optional.empty()), Map.of(doubleVariable.getName(), DOUBLE_COLUMN), Optional.of("count_big(\"c_double\")"));
    // count(DISTINCT bigint)
    testImplementAggregation(new AggregateFunction("count", BIGINT, List.of(bigintVariable), List.of(), true, Optional.empty()), Map.of(bigintVariable.getName(), BIGINT_COLUMN), Optional.empty());
    // count() FILTER (WHERE ...)
    testImplementAggregation(new AggregateFunction("count", BIGINT, List.of(), List.of(), false, filter), Map.of(), Optional.empty());
    // count(bigint) FILTER (WHERE ...)
    testImplementAggregation(new AggregateFunction("count", BIGINT, List.of(bigintVariable), List.of(), false, filter), Map.of(bigintVariable.getName(), BIGINT_COLUMN), Optional.empty());
}