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

Example 16 with Variable

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

the class TestSqlServerClient method testImplementSum.

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

Example 17 with Variable

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

the class ImplementSqlServerCountBig method rewrite.

@Override
public Optional<JdbcExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<String> context) {
    Variable argument = captures.get(ARGUMENT);
    verify(aggregateFunction.getOutputType() == BIGINT);
    return Optional.of(new JdbcExpression(format("count_big(%s)", context.rewriteExpression(argument).orElseThrow()), BIGINT_TYPE));
}

Example 18 with Variable

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

the class ImplementSqlServerStddevPop 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("STDEVP(%s)", context.rewriteExpression(argument).orElseThrow()), columnHandle.getJdbcTypeHandle()));
}

Example 19 with Variable

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

the class ImplementSqlServerVariancePop 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("VARP(%s)", context.rewriteExpression(argument).orElseThrow()), columnHandle.getJdbcTypeHandle()));
}

Example 20 with Variable

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

the class IcebergMetadata method applyProjection.

@Override
public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjection(ConnectorSession session, ConnectorTableHandle handle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
    if (!isProjectionPushdownEnabled(session)) {
        return Optional.empty();
    }
    // Create projected column representations for supported sub expressions. Simple column references and chain of
    // dereferences on a variable are supported right now.
    Set<ConnectorExpression> projectedExpressions = projections.stream().flatMap(expression -> extractSupportedProjectedColumns(expression).stream()).collect(toImmutableSet());
    Map<ConnectorExpression, ProjectedColumnRepresentation> columnProjections = projectedExpressions.stream().collect(toImmutableMap(Function.identity(), HiveApplyProjectionUtil::createProjectedColumnRepresentation));
    IcebergTableHandle icebergTableHandle = (IcebergTableHandle) handle;
    // all references are simple variables
    if (columnProjections.values().stream().allMatch(ProjectedColumnRepresentation::isVariable)) {
        Set<IcebergColumnHandle> projectedColumns = assignments.values().stream().map(IcebergColumnHandle.class::cast).collect(toImmutableSet());
        if (icebergTableHandle.getProjectedColumns().equals(projectedColumns)) {
            return Optional.empty();
        }
        List<Assignment> assignmentsList = assignments.entrySet().stream().map(assignment -> new Assignment(assignment.getKey(), assignment.getValue(), ((IcebergColumnHandle) assignment.getValue()).getType())).collect(toImmutableList());
        return Optional.of(new ProjectionApplicationResult<>(icebergTableHandle.withProjectedColumns(projectedColumns), projections, assignmentsList, false));
    }
    Map<String, Assignment> newAssignments = new HashMap<>();
    ImmutableMap.Builder<ConnectorExpression, Variable> newVariablesBuilder = ImmutableMap.builder();
    ImmutableSet.Builder<IcebergColumnHandle> projectedColumnsBuilder = ImmutableSet.builder();
    for (Map.Entry<ConnectorExpression, ProjectedColumnRepresentation> entry : columnProjections.entrySet()) {
        ConnectorExpression expression = entry.getKey();
        ProjectedColumnRepresentation projectedColumn = entry.getValue();
        IcebergColumnHandle baseColumnHandle = (IcebergColumnHandle) assignments.get(projectedColumn.getVariable().getName());
        IcebergColumnHandle projectedColumnHandle = createProjectedColumnHandle(baseColumnHandle, projectedColumn.getDereferenceIndices(), expression.getType());
        String projectedColumnName = projectedColumnHandle.getQualifiedName();
        Variable projectedColumnVariable = new Variable(projectedColumnName, expression.getType());
        Assignment newAssignment = new Assignment(projectedColumnName, projectedColumnHandle, expression.getType());
        newAssignments.putIfAbsent(projectedColumnName, newAssignment);
        newVariablesBuilder.put(expression, projectedColumnVariable);
        projectedColumnsBuilder.add(projectedColumnHandle);
    }
    // Modify projections to refer to new variables
    Map<ConnectorExpression, Variable> newVariables = newVariablesBuilder.buildOrThrow();
    List<ConnectorExpression> newProjections = projections.stream().map(expression -> replaceWithNewVariables(expression, newVariables)).collect(toImmutableList());
    List<Assignment> outputAssignments = newAssignments.values().stream().collect(toImmutableList());
    return Optional.of(new ProjectionApplicationResult<>(icebergTableHandle.withProjectedColumns(projectedColumnsBuilder.build()), newProjections, outputAssignments, false));
}