Examples with PlannerContext io.trino.sql.PlannerContext used on opensource projects

Example 16 with PlannerContext

use of io.trino.sql.PlannerContext in project trino by trinodb.

the class AbstractTestType method testLiteralFormRecognized.

@Test
public void testLiteralFormRecognized() {
    PlannerContext plannerContext = createPlannerContext();
    LiteralEncoder literalEncoder = new LiteralEncoder(plannerContext);
    for (int position = 0; position < testBlock.getPositionCount(); position++) {
        Object value = readNativeValue(type, testBlock, position);
        Expression expression = literalEncoder.toExpression(TEST_SESSION, value, type);
        if (!isEffectivelyLiteral(plannerContext, TEST_SESSION, expression)) {
            fail(format("Expression not recognized literal for value %s at position %s (%s): %s", value, position, type.getObjectValue(SESSION, testBlock, position), expression));
        }
    }
}

Example 17 with PlannerContext

use of io.trino.sql.PlannerContext in project trino by trinodb.

the class TestPartialTopNWithPresortedInput method assertPlanWithValidation.

private void assertPlanWithValidation(@Language("SQL") String sql, PlanMatchPattern pattern) {
    LocalQueryRunner queryRunner = getQueryRunner();
    queryRunner.inTransaction(queryRunner.getDefaultSession(), transactionSession -> {
        Plan actualPlan = queryRunner.createPlan(transactionSession, sql, OPTIMIZED_AND_VALIDATED, false, WarningCollector.NOOP);
        PlanAssert.assertPlan(transactionSession, queryRunner.getMetadata(), queryRunner.getFunctionManager(), queryRunner.getStatsCalculator(), actualPlan, pattern);
        PlannerContext plannerContext = queryRunner.getPlannerContext();
        new ValidateLimitWithPresortedInput().validate(actualPlan.getRoot(), transactionSession, plannerContext, createTestingTypeAnalyzer(plannerContext), actualPlan.getTypes(), WarningCollector.NOOP);
        return null;
    });
}

Example 18 with PlannerContext

use of io.trino.sql.PlannerContext in project trino by trinodb.

the class ExtractSpatialJoins method tryCreateSpatialJoin.

private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, Expression filter, PlanNodeId nodeId, List<Symbol> outputSymbols, ComparisonExpression spatialComparison, PlannerContext plannerContext, SplitManager splitManager, PageSourceManager pageSourceManager, TypeAnalyzer typeAnalyzer) {
    PlanNode leftNode = joinNode.getLeft();
    PlanNode rightNode = joinNode.getRight();
    List<Symbol> leftSymbols = leftNode.getOutputSymbols();
    List<Symbol> rightSymbols = rightNode.getOutputSymbols();
    Expression radius;
    Optional<Symbol> newRadiusSymbol;
    ComparisonExpression newComparison;
    if (spatialComparison.getOperator() == LESS_THAN || spatialComparison.getOperator() == LESS_THAN_OR_EQUAL) {
        // ST_Distance(a, b) <= r
        radius = spatialComparison.getRight();
        Set<Symbol> radiusSymbols = extractUnique(radius);
        if (radiusSymbols.isEmpty() || (rightSymbols.containsAll(radiusSymbols) && containsNone(leftSymbols, radiusSymbols))) {
            newRadiusSymbol = newRadiusSymbol(context, radius);
            newComparison = new ComparisonExpression(spatialComparison.getOperator(), spatialComparison.getLeft(), toExpression(newRadiusSymbol, radius));
        } else {
            return Result.empty();
        }
    } else {
        // r >= ST_Distance(a, b)
        radius = spatialComparison.getLeft();
        Set<Symbol> radiusSymbols = extractUnique(radius);
        if (radiusSymbols.isEmpty() || (rightSymbols.containsAll(radiusSymbols) && containsNone(leftSymbols, radiusSymbols))) {
            newRadiusSymbol = newRadiusSymbol(context, radius);
            newComparison = new ComparisonExpression(spatialComparison.getOperator().flip(), spatialComparison.getRight(), toExpression(newRadiusSymbol, radius));
        } else {
            return Result.empty();
        }
    }
    Expression newFilter = replaceExpression(filter, ImmutableMap.of(spatialComparison, newComparison));
    PlanNode newRightNode = newRadiusSymbol.map(symbol -> addProjection(context, rightNode, symbol, radius)).orElse(rightNode);
    JoinNode newJoinNode = new JoinNode(joinNode.getId(), joinNode.getType(), leftNode, newRightNode, joinNode.getCriteria(), joinNode.getLeftOutputSymbols(), joinNode.getRightOutputSymbols(), joinNode.isMaySkipOutputDuplicates(), Optional.of(newFilter), joinNode.getLeftHashSymbol(), joinNode.getRightHashSymbol(), joinNode.getDistributionType(), joinNode.isSpillable(), joinNode.getDynamicFilters(), joinNode.getReorderJoinStatsAndCost());
    return tryCreateSpatialJoin(context, newJoinNode, newFilter, nodeId, outputSymbols, (FunctionCall) newComparison.getLeft(), Optional.of(newComparison.getRight()), plannerContext, splitManager, pageSourceManager, typeAnalyzer);
}

Example 19 with PlannerContext

use of io.trino.sql.PlannerContext in project trino by trinodb.

the class PushAggregationIntoTableScan method pushAggregationIntoTableScan.

public static Optional<PlanNode> pushAggregationIntoTableScan(PlannerContext plannerContext, Context context, PlanNode aggregationNode, TableScanNode tableScan, Map<Symbol, AggregationNode.Aggregation> aggregations, List<Symbol> groupingKeys) {
    Map<String, ColumnHandle> assignments = tableScan.getAssignments().entrySet().stream().collect(toImmutableMap(entry -> entry.getKey().getName(), Entry::getValue));
    List<Entry<Symbol, AggregationNode.Aggregation>> aggregationsList = aggregations.entrySet().stream().collect(toImmutableList());
    List<AggregateFunction> aggregateFunctions = aggregationsList.stream().map(Entry::getValue).map(aggregation -> toAggregateFunction(plannerContext.getMetadata(), context, aggregation)).collect(toImmutableList());
    List<Symbol> aggregationOutputSymbols = aggregationsList.stream().map(Entry::getKey).collect(toImmutableList());
    List<ColumnHandle> groupByColumns = groupingKeys.stream().map(groupByColumn -> assignments.get(groupByColumn.getName())).collect(toImmutableList());
    Optional<AggregationApplicationResult<TableHandle>> aggregationPushdownResult = plannerContext.getMetadata().applyAggregation(context.getSession(), tableScan.getTable(), aggregateFunctions, assignments, ImmutableList.of(groupByColumns));
    if (aggregationPushdownResult.isEmpty()) {
        return Optional.empty();
    }
    AggregationApplicationResult<TableHandle> result = aggregationPushdownResult.get();
    // The new scan outputs should be the symbols associated with grouping columns plus the symbols associated with aggregations.
    ImmutableList.Builder<Symbol> newScanOutputs = ImmutableList.builder();
    newScanOutputs.addAll(tableScan.getOutputSymbols());
    ImmutableBiMap.Builder<Symbol, ColumnHandle> newScanAssignments = ImmutableBiMap.builder();
    newScanAssignments.putAll(tableScan.getAssignments());
    Map<String, Symbol> variableMappings = new HashMap<>();
    for (Assignment assignment : result.getAssignments()) {
        Symbol symbol = context.getSymbolAllocator().newSymbol(assignment.getVariable(), assignment.getType());
        newScanOutputs.add(symbol);
        newScanAssignments.put(symbol, assignment.getColumn());
        variableMappings.put(assignment.getVariable(), symbol);
    }
    List<Expression> newProjections = result.getProjections().stream().map(expression -> ConnectorExpressionTranslator.translate(context.getSession(), expression, plannerContext, variableMappings, new LiteralEncoder(plannerContext))).collect(toImmutableList());
    verify(aggregationOutputSymbols.size() == newProjections.size());
    Assignments.Builder assignmentBuilder = Assignments.builder();
    IntStream.range(0, aggregationOutputSymbols.size()).forEach(index -> assignmentBuilder.put(aggregationOutputSymbols.get(index), newProjections.get(index)));
    ImmutableBiMap<Symbol, ColumnHandle> scanAssignments = newScanAssignments.build();
    ImmutableBiMap<ColumnHandle, Symbol> columnHandleToSymbol = scanAssignments.inverse();
    // projections assignmentBuilder should have both agg and group by so we add all the group bys as symbol references
    groupingKeys.forEach(groupBySymbol -> {
        // if the connector returned a new mapping from oldColumnHandle to newColumnHandle, groupBy needs to point to
        // new columnHandle's symbol reference, otherwise it will continue pointing at oldColumnHandle.
        ColumnHandle originalColumnHandle = assignments.get(groupBySymbol.getName());
        ColumnHandle groupByColumnHandle = result.getGroupingColumnMapping().getOrDefault(originalColumnHandle, originalColumnHandle);
        assignmentBuilder.put(groupBySymbol, columnHandleToSymbol.get(groupByColumnHandle).toSymbolReference());
    });
    return Optional.of(new ProjectNode(context.getIdAllocator().getNextId(), new TableScanNode(context.getIdAllocator().getNextId(), result.getHandle(), newScanOutputs.build(), scanAssignments, TupleDomain.all(), deriveTableStatisticsForPushdown(context.getStatsProvider(), context.getSession(), result.isPrecalculateStatistics(), aggregationNode), tableScan.isUpdateTarget(), // table scan partitioning might have changed with new table handle
    Optional.empty()), assignmentBuilder.build()));
}

Example 20 with PlannerContext

use of io.trino.sql.PlannerContext in project trino by trinodb.

the class InlineProjections method inlineProjections.

static Optional<ProjectNode> inlineProjections(PlannerContext plannerContext, ProjectNode parent, ProjectNode child, Session session, TypeAnalyzer typeAnalyzer, TypeProvider types) {
    // squash identity projections
    if (parent.isIdentity() && child.isIdentity()) {
        return Optional.of((ProjectNode) parent.replaceChildren(ImmutableList.of(child.getSource())));
    }
    Set<Symbol> targets = extractInliningTargets(plannerContext, parent, child, session, typeAnalyzer, types);
    if (targets.isEmpty()) {
        return Optional.empty();
    }
    // inline the expressions
    Assignments assignments = child.getAssignments().filter(targets::contains);
    Map<Symbol, Expression> parentAssignments = parent.getAssignments().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> inlineReferences(entry.getValue(), assignments)));
    // Synthesize identity assignments for the inputs of expressions that were inlined
    // to place in the child projection.
    Set<Symbol> inputs = child.getAssignments().entrySet().stream().filter(entry -> targets.contains(entry.getKey())).map(Map.Entry::getValue).flatMap(entry -> SymbolsExtractor.extractAll(entry).stream()).collect(toSet());
    Assignments.Builder newChildAssignmentsBuilder = Assignments.builder();
    for (Map.Entry<Symbol, Expression> assignment : child.getAssignments().entrySet()) {
        if (!targets.contains(assignment.getKey())) {
            newChildAssignmentsBuilder.put(assignment);
        }
    }
    for (Symbol input : inputs) {
        newChildAssignmentsBuilder.putIdentity(input);
    }
    Assignments newChildAssignments = newChildAssignmentsBuilder.build();
    PlanNode newChild;
    if (newChildAssignments.isIdentity()) {
        newChild = child.getSource();
    } else {
        newChild = new ProjectNode(child.getId(), child.getSource(), newChildAssignments);
    }
    return Optional.of(new ProjectNode(parent.getId(), newChild, Assignments.copyOf(parentAssignments)));
}