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

Example 26 with Variable

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

the class PruneTableScanColumns method pruneColumns.

public static Optional<PlanNode> pruneColumns(Metadata metadata, TypeProvider types, Session session, TableScanNode node, Set<Symbol> referencedOutputs) {
    List<Symbol> newOutputs = filteredCopy(node.getOutputSymbols(), referencedOutputs::contains);
    if (newOutputs.size() == node.getOutputSymbols().size()) {
        return Optional.empty();
    }
    List<ConnectorExpression> projections = newOutputs.stream().map(symbol -> new Variable(symbol.getName(), types.get(symbol))).collect(toImmutableList());
    TableHandle handle = node.getTable();
    Optional<ProjectionApplicationResult<TableHandle>> result = metadata.applyProjection(session, handle, projections, newOutputs.stream().collect(toImmutableMap(Symbol::getName, node.getAssignments()::get)));
    Map<Symbol, ColumnHandle> newAssignments;
    // Bail out if the connector does anything other than limit the list of columns (e.g., if it synthesizes arbitrary expressions)
    if (result.isPresent() && result.get().getProjections().stream().allMatch(Variable.class::isInstance)) {
        handle = result.get().getHandle();
        Map<String, ColumnHandle> assignments = result.get().getAssignments().stream().collect(toImmutableMap(Assignment::getVariable, Assignment::getColumn));
        ImmutableMap.Builder<Symbol, ColumnHandle> builder = ImmutableMap.builder();
        for (int i = 0; i < newOutputs.size(); i++) {
            Variable variable = (Variable) result.get().getProjections().get(i);
            builder.put(newOutputs.get(i), assignments.get(variable.getName()));
        }
        newAssignments = builder.buildOrThrow();
    } else {
        newAssignments = newOutputs.stream().collect(toImmutableMap(Function.identity(), node.getAssignments()::get));
    }
    Set<ColumnHandle> visibleColumns = ImmutableSet.copyOf(newAssignments.values());
    TupleDomain<ColumnHandle> enforcedConstraint = node.getEnforcedConstraint().filter((columnHandle, domain) -> visibleColumns.contains(columnHandle));
    Optional<PlanNodeStatsEstimate> newStatistics = node.getStatistics().map(statistics -> new PlanNodeStatsEstimate(statistics.getOutputRowCount(), statistics.getSymbolStatistics().entrySet().stream().filter(entry -> newAssignments.containsKey(entry.getKey())).collect(toImmutableMap(Entry::getKey, Entry::getValue))));
    return Optional.of(new TableScanNode(node.getId(), handle, newOutputs, newAssignments, enforcedConstraint, newStatistics, node.isUpdateTarget(), node.getUseConnectorNodePartitioning()));
}

Example 27 with Variable

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

the class PushProjectionIntoTableScan method apply.

@Override
public Result apply(ProjectNode project, Captures captures, Context context) {
    TableScanNode tableScan = captures.get(TABLE_SCAN);
    // Extract translatable components from projection expressions. Prepare a mapping from these internal
    // expression nodes to corresponding ConnectorExpression translations.
    Map<NodeRef<Expression>, ConnectorExpression> partialTranslations = project.getAssignments().getMap().entrySet().stream().flatMap(expression -> extractPartialTranslations(expression.getValue(), context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes(), plannerContext).entrySet().stream()).collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (first, ignore) -> first));
    List<NodeRef<Expression>> nodesForPartialProjections = ImmutableList.copyOf(partialTranslations.keySet());
    List<ConnectorExpression> connectorPartialProjections = ImmutableList.copyOf(partialTranslations.values());
    Map<String, Symbol> inputVariableMappings = tableScan.getAssignments().keySet().stream().collect(toImmutableMap(Symbol::getName, identity()));
    Map<String, ColumnHandle> assignments = inputVariableMappings.entrySet().stream().collect(toImmutableMap(Entry::getKey, entry -> tableScan.getAssignments().get(entry.getValue())));
    Optional<ProjectionApplicationResult<TableHandle>> result = plannerContext.getMetadata().applyProjection(context.getSession(), tableScan.getTable(), connectorPartialProjections, assignments);
    if (result.isEmpty()) {
        return Result.empty();
    }
    List<ConnectorExpression> newConnectorPartialProjections = result.get().getProjections();
    checkState(newConnectorPartialProjections.size() == connectorPartialProjections.size(), "Mismatch between input and output projections from the connector: expected %s but got %s", connectorPartialProjections.size(), newConnectorPartialProjections.size());
    List<Symbol> newScanOutputs = new ArrayList<>();
    Map<Symbol, ColumnHandle> newScanAssignments = new HashMap<>();
    Map<String, Symbol> variableMappings = new HashMap<>();
    for (Assignment assignment : result.get().getAssignments()) {
        Symbol symbol = context.getSymbolAllocator().newSymbol(assignment.getVariable(), assignment.getType());
        newScanOutputs.add(symbol);
        newScanAssignments.put(symbol, assignment.getColumn());
        variableMappings.put(assignment.getVariable(), symbol);
    }
    // Translate partial connector projections back to new partial projections
    List<Expression> newPartialProjections = newConnectorPartialProjections.stream().map(expression -> ConnectorExpressionTranslator.translate(context.getSession(), expression, plannerContext, variableMappings, new LiteralEncoder(plannerContext))).collect(toImmutableList());
    // Map internal node references to new partial projections
    ImmutableMap.Builder<NodeRef<Expression>, Expression> nodesToNewPartialProjectionsBuilder = ImmutableMap.builder();
    for (int i = 0; i < nodesForPartialProjections.size(); i++) {
        nodesToNewPartialProjectionsBuilder.put(nodesForPartialProjections.get(i), newPartialProjections.get(i));
    }
    Map<NodeRef<Expression>, Expression> nodesToNewPartialProjections = nodesToNewPartialProjectionsBuilder.buildOrThrow();
    // Stitch partial translations to form new complete projections
    Assignments.Builder newProjectionAssignments = Assignments.builder();
    project.getAssignments().entrySet().forEach(entry -> {
        newProjectionAssignments.put(entry.getKey(), replaceExpression(entry.getValue(), nodesToNewPartialProjections));
    });
    Optional<PlanNodeStatsEstimate> newStatistics = tableScan.getStatistics().map(statistics -> {
        PlanNodeStatsEstimate.Builder builder = PlanNodeStatsEstimate.builder();
        builder.setOutputRowCount(statistics.getOutputRowCount());
        for (int i = 0; i < connectorPartialProjections.size(); i++) {
            ConnectorExpression inputConnectorExpression = connectorPartialProjections.get(i);
            ConnectorExpression resultConnectorExpression = newConnectorPartialProjections.get(i);
            if (!(resultConnectorExpression instanceof Variable)) {
                continue;
            }
            String resultVariableName = ((Variable) resultConnectorExpression).getName();
            Expression inputExpression = ConnectorExpressionTranslator.translate(context.getSession(), inputConnectorExpression, plannerContext, inputVariableMappings, new LiteralEncoder(plannerContext));
            SymbolStatsEstimate symbolStatistics = scalarStatsCalculator.calculate(inputExpression, statistics, context.getSession(), context.getSymbolAllocator().getTypes());
            builder.addSymbolStatistics(variableMappings.get(resultVariableName), symbolStatistics);
        }
        return builder.build();
    });
    verifyTablePartitioning(context, tableScan, result.get().getHandle());
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new TableScanNode(tableScan.getId(), result.get().getHandle(), newScanOutputs, newScanAssignments, TupleDomain.all(), newStatistics, tableScan.isUpdateTarget(), tableScan.getUseConnectorNodePartitioning()), newProjectionAssignments.build()));
}

Example 28 with Variable

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

the class DefaultJdbcMetadata method applyAggregation.

@Override
public Optional<AggregationApplicationResult<ConnectorTableHandle>> applyAggregation(ConnectorSession session, ConnectorTableHandle table, List<AggregateFunction> aggregates, Map<String, ColumnHandle> assignments, List<List<ColumnHandle>> groupingSets) {
    if (!isAggregationPushdownEnabled(session)) {
        return Optional.empty();
    }
    JdbcTableHandle handle = (JdbcTableHandle) table;
    // Global aggregation is represented by [[]]
    verify(!groupingSets.isEmpty(), "No grouping sets provided");
    if (!jdbcClient.supportsAggregationPushdown(session, handle, aggregates, assignments, groupingSets)) {
        // JDBC client implementation prevents pushdown for the given table
        return Optional.empty();
    }
    if (handle.getLimit().isPresent()) {
        handle = flushAttributesAsQuery(session, handle);
    }
    int nextSyntheticColumnId = handle.getNextSyntheticColumnId();
    ImmutableList.Builder<JdbcColumnHandle> newColumns = ImmutableList.builder();
    ImmutableList.Builder<ConnectorExpression> projections = ImmutableList.builder();
    ImmutableList.Builder<Assignment> resultAssignments = ImmutableList.builder();
    ImmutableMap.Builder<String, String> expressions = ImmutableMap.builder();
    List<List<JdbcColumnHandle>> groupingSetsAsJdbcColumnHandles = groupingSets.stream().map(groupingSet -> groupingSet.stream().map(JdbcColumnHandle.class::cast).collect(toImmutableList())).collect(toImmutableList());
    Optional<List<JdbcColumnHandle>> tableColumns = handle.getColumns();
    groupingSetsAsJdbcColumnHandles.stream().flatMap(List::stream).distinct().peek(handle.getColumns().<Consumer<JdbcColumnHandle>>map(columns -> groupKey -> verify(columns.contains(groupKey), "applyAggregation called with a grouping column %s which was not included in the table columns: %s", groupKey, tableColumns)).orElse(groupKey -> {
    })).forEach(newColumns::add);
    for (AggregateFunction aggregate : aggregates) {
        Optional<JdbcExpression> expression = jdbcClient.implementAggregation(session, aggregate, assignments);
        if (expression.isEmpty()) {
            return Optional.empty();
        }
        String columnName = SYNTHETIC_COLUMN_NAME_PREFIX + nextSyntheticColumnId;
        nextSyntheticColumnId++;
        JdbcColumnHandle newColumn = JdbcColumnHandle.builder().setColumnName(columnName).setJdbcTypeHandle(expression.get().getJdbcTypeHandle()).setColumnType(aggregate.getOutputType()).setComment(Optional.of("synthetic")).build();
        newColumns.add(newColumn);
        projections.add(new Variable(newColumn.getColumnName(), aggregate.getOutputType()));
        resultAssignments.add(new Assignment(newColumn.getColumnName(), newColumn, aggregate.getOutputType()));
        expressions.put(columnName, expression.get().getExpression());
    }
    List<JdbcColumnHandle> newColumnsList = newColumns.build();
    // We need to have matching column handles in JdbcTableHandle constructed below, as columns read via JDBC must match column handles list.
    // For more context see assertion in JdbcRecordSetProvider.getRecordSet
    PreparedQuery preparedQuery = jdbcClient.prepareQuery(session, handle, Optional.of(groupingSetsAsJdbcColumnHandles), newColumnsList, expressions.buildOrThrow());
    handle = new JdbcTableHandle(new JdbcQueryRelationHandle(preparedQuery), TupleDomain.all(), ImmutableList.of(), Optional.empty(), OptionalLong.empty(), Optional.of(newColumnsList), handle.getAllReferencedTables(), nextSyntheticColumnId);
    return Optional.of(new AggregationApplicationResult<>(handle, projections.build(), resultAssignments.build(), ImmutableMap.of(), false));
}

Example 29 with Variable

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

the class DefaultJdbcMetadata method getVariableColumnHandle.

private static Optional<JdbcColumnHandle> getVariableColumnHandle(Map<String, ColumnHandle> assignments, ConnectorExpression expression) {
    requireNonNull(assignments, "assignments is null");
    requireNonNull(expression, "expression is null");
    if (!(expression instanceof Variable)) {
        return Optional.empty();
    }
    String name = ((Variable) expression).getName();
    ColumnHandle columnHandle = assignments.get(name);
    verifyNotNull(columnHandle, "No assignment for %s", name);
    return Optional.of(((JdbcColumnHandle) columnHandle));
}

Example 30 with Variable

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

the class ImplementAvgFloatingPoint 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(aggregateFunction.getOutputType() == columnHandle.getColumnType());
    return Optional.of(new JdbcExpression(format("avg(%s)", context.rewriteExpression(argument).orElseThrow()), columnHandle.getJdbcTypeHandle()));
}