Examples with ProjectionApplicationResult io.trino.spi.connector.ProjectionApplicationResult used on opensource projects

Example 6 with ProjectionApplicationResult

use of io.trino.spi.connector.ProjectionApplicationResult 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 7 with ProjectionApplicationResult

use of io.trino.spi.connector.ProjectionApplicationResult in project trino by trinodb.

the class HiveMetadata 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));
    HiveTableHandle hiveTableHandle = (HiveTableHandle) handle;
    // all references are simple variables
    if (columnProjections.values().stream().allMatch(ProjectedColumnRepresentation::isVariable)) {
        Set<ColumnHandle> projectedColumns = ImmutableSet.copyOf(assignments.values());
        if (hiveTableHandle.getProjectedColumns().equals(projectedColumns)) {
            return Optional.empty();
        }
        List<Assignment> assignmentsList = assignments.entrySet().stream().map(assignment -> new Assignment(assignment.getKey(), assignment.getValue(), ((HiveColumnHandle) assignment.getValue()).getType())).collect(toImmutableList());
        return Optional.of(new ProjectionApplicationResult<>(hiveTableHandle.withProjectedColumns(projectedColumns), projections, assignmentsList, false));
    }
    Map<String, Assignment> newAssignments = new HashMap<>();
    ImmutableMap.Builder<ConnectorExpression, Variable> newVariablesBuilder = ImmutableMap.builder();
    ImmutableSet.Builder<ColumnHandle> projectedColumnsBuilder = ImmutableSet.builder();
    for (Map.Entry<ConnectorExpression, ProjectedColumnRepresentation> entry : columnProjections.entrySet()) {
        ConnectorExpression expression = entry.getKey();
        ProjectedColumnRepresentation projectedColumn = entry.getValue();
        ColumnHandle projectedColumnHandle;
        String projectedColumnName;
        // See if input already contains a columnhandle for this projected column, avoid creating duplicates.
        Optional<String> existingColumn = find(assignments, projectedColumn);
        if (existingColumn.isPresent()) {
            projectedColumnName = existingColumn.get();
            projectedColumnHandle = assignments.get(projectedColumnName);
        } else {
            // Create a new column handle
            HiveColumnHandle oldColumnHandle = (HiveColumnHandle) assignments.get(projectedColumn.getVariable().getName());
            projectedColumnHandle = createProjectedColumnHandle(oldColumnHandle, projectedColumn.getDereferenceIndices());
            projectedColumnName = ((HiveColumnHandle) projectedColumnHandle).getName();
        }
        Variable projectedColumnVariable = new Variable(projectedColumnName, expression.getType());
        Assignment newAssignment = new Assignment(projectedColumnName, projectedColumnHandle, expression.getType());
        newAssignments.put(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<>(hiveTableHandle.withProjectedColumns(projectedColumnsBuilder.build()), newProjections, outputAssignments, false));
}