Examples with Assignments io.prestosql.spi.plan.Assignments used on opensource projects

Example 21 with Assignments

use of io.prestosql.spi.plan.Assignments in project boostkit-bigdata by kunpengcompute.

the class TestHivePartialAggregationPushdown method testPartialAggregationAndProjectPushdown.

@Test
public void testPartialAggregationAndProjectPushdown() {
    // select count(x + 5) from table group by x
    TableScanNode tableScanNode = buildTableScanNode(COLUMN_INT);
    CallExpression callExpression = createOperationExpression(OperatorType.ADD, new VariableReferenceExpression(COLUMN_INT.getName(), INTEGER), new ConstantExpression(5, INTEGER));
    List<Symbol> symbols = ImmutableList.of(new Symbol(COLUMN_INT.getName()));
    List<RowExpression> rowExpressions = ImmutableList.of(callExpression);
    ProjectNode projectNode = buildProjectNode(tableScanNode, symbols, rowExpressions);
    AggregationNode aggregationNode = buildCountAggregationNode(projectNode);
    PlanNode output = AGGREGATION_OPTIMIZER.optimize(aggregationNode, OFFLOAD_SESSION, COLUMN_TYPE_MAP, SYMBOL_ALLOCATOR, ID_ALLOCATOR);
    Assignments assignmentsExpected = buildAssignments(symbols, rowExpressions);
    matchProjection(output, assignmentsExpected.getMap());
}

Example 22 with Assignments

use of io.prestosql.spi.plan.Assignments in project hetu-core by openlookeng.

the class TestBaseJdbcPushDownBase method tableScan.

protected TableScanNode tableScan(PlanBuilder planBuilder, JdbcTableHandle connectorTableHandle, JdbcColumnHandle... columnHandles) {
    List<Symbol> symbols = Arrays.stream(columnHandles).map(column -> new Symbol(column.getColumnName().toLowerCase(Locale.ENGLISH))).collect(toImmutableList());
    ImmutableMap.Builder<Symbol, ColumnHandle> assignments = ImmutableMap.builder();
    for (int i = 0; i < symbols.size(); i++) {
        assignments.put(symbols.get(i), columnHandles[i]);
    }
    TableHandle tableHandle = new TableHandle(catalogName, connectorTableHandle, TestingTransactionHandle.create(), Optional.empty());
    return planBuilder.tableScan(tableHandle, symbols, assignments.build());
}

Example 23 with Assignments

use of io.prestosql.spi.plan.Assignments in project hetu-core by openlookeng.

the class TransformCorrelatedSingleRowSubqueryToProject method apply.

@Override
public Result apply(LateralJoinNode parent, Captures captures, Context context) {
    List<ValuesNode> values = searchFrom(parent.getSubquery(), context.getLookup()).recurseOnlyWhen(ProjectNode.class::isInstance).where(ValuesNode.class::isInstance).findAll();
    if (values.size() != 1 || !isSingleRowValuesWithNoColumns(values.get(0))) {
        return Result.empty();
    }
    List<ProjectNode> subqueryProjections = searchFrom(parent.getSubquery(), context.getLookup()).where(node -> node instanceof ProjectNode && !node.getOutputSymbols().equals(parent.getCorrelation())).findAll();
    if (subqueryProjections.size() == 0) {
        return Result.ofPlanNode(parent.getInput());
    }
    if (subqueryProjections.size() == 1) {
        Assignments assignments = Assignments.builder().putAll(AssignmentUtils.identityAsSymbolReferences(parent.getInput().getOutputSymbols())).putAll(subqueryProjections.get(0).getAssignments()).build();
        return Result.ofPlanNode(projectNode(parent.getInput(), assignments, context));
    }
    return Result.empty();
}

Example 24 with Assignments

use of io.prestosql.spi.plan.Assignments in project hetu-core by openlookeng.

the class TransformExistsApplyToLateralNode method rewriteToNonDefaultAggregation.

private Optional<PlanNode> rewriteToNonDefaultAggregation(ApplyNode applyNode, Context context) {
    checkState(applyNode.getSubquery().getOutputSymbols().isEmpty(), "Expected subquery output symbols to be pruned");
    Symbol exists = getOnlyElement(applyNode.getSubqueryAssignments().getSymbols());
    Symbol subqueryTrue = context.getSymbolAllocator().newSymbol("subqueryTrue", BOOLEAN);
    Assignments.Builder assignments = Assignments.builder();
    assignments.putAll(AssignmentUtils.identityAsSymbolReferences(applyNode.getInput().getOutputSymbols()));
    assignments.put(exists, castToRowExpression(new CoalesceExpression(ImmutableList.of(toSymbolReference(subqueryTrue), BooleanLiteral.FALSE_LITERAL))));
    PlanNode subquery = new ProjectNode(context.getIdAllocator().getNextId(), new LimitNode(context.getIdAllocator().getNextId(), applyNode.getSubquery(), 1L, false), Assignments.of(subqueryTrue, castToRowExpression(TRUE_LITERAL)));
    PlanNodeDecorrelator decorrelator = new PlanNodeDecorrelator(context.getIdAllocator(), context.getLookup());
    if (!decorrelator.decorrelateFilters(subquery, applyNode.getCorrelation()).isPresent()) {
        return Optional.empty();
    }
    return Optional.of(new ProjectNode(context.getIdAllocator().getNextId(), new LateralJoinNode(applyNode.getId(), applyNode.getInput(), subquery, applyNode.getCorrelation(), LEFT, TRUE_LITERAL, applyNode.getOriginSubquery()), assignments.build()));
}

Example 25 with Assignments

use of io.prestosql.spi.plan.Assignments in project hetu-core by openlookeng.

the class PushPartialAggregationThroughExchange method pushPartial.

private PlanNode pushPartial(AggregationNode aggregation, ExchangeNode exchange, Context context) {
    List<PlanNode> partials = new ArrayList<>();
    for (int i = 0; i < exchange.getSources().size(); i++) {
        PlanNode source = exchange.getSources().get(i);
        SymbolMapper.Builder mappingsBuilder = SymbolMapper.builder();
        for (int outputIndex = 0; outputIndex < exchange.getOutputSymbols().size(); outputIndex++) {
            Symbol output = exchange.getOutputSymbols().get(outputIndex);
            Symbol input = exchange.getInputs().get(i).get(outputIndex);
            if (!output.equals(input)) {
                mappingsBuilder.put(output, input);
            }
        }
        SymbolMapper symbolMapper = mappingsBuilder.build();
        if (symbolMapper.getTypes() == null) {
            symbolMapper.setTypes(context.getSymbolAllocator().getTypes());
        }
        AggregationNode mappedPartial = symbolMapper.map(aggregation, source, context.getIdAllocator());
        Assignments.Builder assignments = Assignments.builder();
        for (Symbol output : aggregation.getOutputSymbols()) {
            Symbol input = symbolMapper.map(output);
            assignments.put(output, VariableReferenceSymbolConverter.toVariableReference(input, context.getSymbolAllocator().getTypes()));
        }
        partials.add(new ProjectNode(context.getIdAllocator().getNextId(), mappedPartial, assignments.build()));
    }
    for (PlanNode node : partials) {
        verify(aggregation.getOutputSymbols().equals(node.getOutputSymbols()));
    }
    // Since this exchange source is now guaranteed to have the same symbols as the inputs to the the partial
    // aggregation, we don't need to rewrite symbols in the partitioning function
    PartitioningScheme partitioning = new PartitioningScheme(exchange.getPartitioningScheme().getPartitioning(), aggregation.getOutputSymbols(), exchange.getPartitioningScheme().getHashColumn(), exchange.getPartitioningScheme().isReplicateNullsAndAny(), exchange.getPartitioningScheme().getBucketToPartition());
    return new ExchangeNode(context.getIdAllocator().getNextId(), exchange.getType(), exchange.getScope(), partitioning, partials, ImmutableList.copyOf(Collections.nCopies(partials.size(), aggregation.getOutputSymbols())), Optional.empty(), aggregation.getAggregationType());
}