Examples with Assignments io.trino.sql.planner.plan.Assignments used on opensource projects

Example 26 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class TestPruneTableScanColumns method testPushColumnPruningProjection.

@Test
public void testPushColumnPruningProjection() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        String mockCatalog = "mock_catalog";
        String testSchema = "test_schema";
        String testTable = "test_table";
        SchemaTableName testSchemaTable = new SchemaTableName(testSchema, testTable);
        ColumnHandle columnHandleA = new MockConnectorColumnHandle("cola", DATE);
        ColumnHandle columnHandleB = new MockConnectorColumnHandle("colb", DOUBLE);
        Map<String, ColumnHandle> assignments = ImmutableMap.of("cola", columnHandleA, "colb", columnHandleB);
        // Create catalog with applyProjection
        MockConnectorFactory factory = MockConnectorFactory.builder().withListSchemaNames(connectorSession -> ImmutableList.of(testSchema)).withListTables((connectorSession, schema) -> testSchema.equals(schema) ? ImmutableList.of(testSchemaTable) : ImmutableList.of()).withGetColumns(schemaTableName -> assignments.entrySet().stream().map(entry -> new ColumnMetadata(entry.getKey(), ((MockConnectorColumnHandle) entry.getValue()).getType())).collect(toImmutableList())).withApplyProjection(this::mockApplyProjection).build();
        ruleTester.getQueryRunner().createCatalog(mockCatalog, factory, ImmutableMap.of());
        ruleTester.assertThat(new PruneTableScanColumns(ruleTester.getMetadata())).on(p -> {
            Symbol symbolA = p.symbol("cola", DATE);
            Symbol symbolB = p.symbol("colb", DOUBLE);
            return p.project(Assignments.of(p.symbol("x"), symbolB.toSymbolReference()), p.tableScan(new TableHandle(new CatalogName(mockCatalog), new MockConnectorTableHandle(testSchemaTable), MockConnectorTransactionHandle.INSTANCE), ImmutableList.of(symbolA, symbolB), ImmutableMap.of(symbolA, columnHandleA, symbolB, columnHandleB)));
        }).withSession(testSessionBuilder().setCatalog(mockCatalog).setSchema(testSchema).build()).matches(strictProject(ImmutableMap.of("expr", PlanMatchPattern.expression("COLB")), tableScan(new MockConnectorTableHandle(testSchemaTable, TupleDomain.all(), Optional.of(ImmutableList.of(columnHandleB)))::equals, TupleDomain.all(), ImmutableMap.of("COLB", columnHandleB::equals))));
    }
}

Example 27 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class TestPruneTableScanColumns method mockApplyProjection.

private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
    MockConnectorTableHandle handle = (MockConnectorTableHandle) tableHandle;
    List<Variable> variables = projections.stream().map(Variable.class::cast).collect(toImmutableList());
    List<ColumnHandle> newColumns = variables.stream().map(variable -> assignments.get(variable.getName())).collect(toImmutableList());
    if (handle.getColumns().isPresent() && newColumns.equals(handle.getColumns().get())) {
        return Optional.empty();
    }
    return Optional.of(new ProjectionApplicationResult<>(new MockConnectorTableHandle(handle.getTableName(), handle.getConstraint(), Optional.of(newColumns)), projections, variables.stream().map(variable -> new Assignment(variable.getName(), assignments.get(variable.getName()), ((MockConnectorColumnHandle) assignments.get(variable.getName())).getType())).collect(toImmutableList()), false));
}

Example 28 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class TestMergeProjectWithValues method testMergeProjectWithValues.

@Test
public void testMergeProjectWithValues() {
    tester().assertThat(new MergeProjectWithValues(tester().getMetadata())).on(p -> {
        Symbol a = p.symbol("a");
        Symbol b = p.symbol("b");
        Symbol c = p.symbol("c");
        Symbol d = p.symbol("d");
        Symbol e = p.symbol("e");
        Symbol f = p.symbol("f");
        Assignments.Builder assignments = Assignments.builder();
        // identity assignment
        assignments.putIdentity(a);
        // renaming assignment
        assignments.put(d, b.toSymbolReference());
        // expression involving input symbol
        assignments.put(e, new IsNullPredicate(a.toSymbolReference()));
        // constant expression
        assignments.put(f, new LongLiteral("1"));
        return p.project(assignments.build(), p.valuesOfExpressions(ImmutableList.of(a, b, c), ImmutableList.of(new Row(ImmutableList.of(new CharLiteral("x"), new BooleanLiteral("true"), new LongLiteral("1"))), new Row(ImmutableList.of(new CharLiteral("y"), new BooleanLiteral("false"), new LongLiteral("2"))), new Row(ImmutableList.of(new CharLiteral("z"), new BooleanLiteral("true"), new LongLiteral("3"))))));
    }).matches(values(ImmutableList.of("a", "d", "e", "f"), ImmutableList.of(ImmutableList.of(new CharLiteral("x"), new BooleanLiteral("true"), new IsNullPredicate(new CharLiteral("x")), new LongLiteral("1")), ImmutableList.of(new CharLiteral("y"), new BooleanLiteral("false"), new IsNullPredicate(new CharLiteral("y")), new LongLiteral("1")), ImmutableList.of(new CharLiteral("z"), new BooleanLiteral("true"), new IsNullPredicate(new CharLiteral("z")), new LongLiteral("1")))));
    // ValuesNode has no rows
    tester().assertThat(new MergeProjectWithValues(tester().getMetadata())).on(p -> {
        Symbol a = p.symbol("a");
        Symbol b = p.symbol("b");
        Symbol c = p.symbol("c");
        Symbol d = p.symbol("d");
        Symbol e = p.symbol("e");
        Symbol f = p.symbol("f");
        Assignments.Builder assignments = Assignments.builder();
        // identity assignment
        assignments.putIdentity(a);
        // renaming assignment
        assignments.put(d, b.toSymbolReference());
        // expression involving input symbol
        assignments.put(e, new IsNullPredicate(a.toSymbolReference()));
        // constant expression
        assignments.put(f, new LongLiteral("1"));
        return p.project(assignments.build(), p.values(ImmutableList.of(a, b, c), ImmutableList.of()));
    }).matches(values(ImmutableList.of("a", "d", "e", "f"), ImmutableList.of()));
}

Example 29 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class ExtractDereferencesFromFilterAboveScan method apply.

@Override
public Result apply(FilterNode node, Captures captures, Context context) {
    Set<SubscriptExpression> dereferences = extractRowSubscripts(ImmutableList.of(node.getPredicate()), true, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    if (dereferences.isEmpty()) {
        return Result.empty();
    }
    Assignments assignments = Assignments.of(dereferences, context.getSession(), context.getSymbolAllocator(), typeAnalyzer);
    Map<Expression, SymbolReference> mappings = HashBiMap.create(assignments.getMap()).inverse().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
    PlanNode source = node.getSource();
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new FilterNode(context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), source, Assignments.builder().putIdentities(source.getOutputSymbols()).putAll(assignments).build()), replaceExpression(node.getPredicate(), mappings)), Assignments.identity(node.getOutputSymbols())));
}

Example 30 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class DecorrelateInnerUnnestWithGlobalAggregation method withGroupingAndMask.

private static AggregationNode withGroupingAndMask(AggregationNode aggregationNode, List<Symbol> groupingSymbols, Symbol mask, PlanNode source, SymbolAllocator symbolAllocator, PlanNodeIdAllocator idAllocator) {
    // Every original aggregation agg() will be rewritten to agg() mask(mask_symbol). If the aggregation
    // already has a mask, it will be replaced with conjunction of the existing mask and mask_symbol.
    ImmutableMap.Builder<Symbol, Symbol> masks = ImmutableMap.builder();
    Assignments.Builder assignmentsBuilder = Assignments.builder();
    for (Map.Entry<Symbol, AggregationNode.Aggregation> entry : aggregationNode.getAggregations().entrySet()) {
        AggregationNode.Aggregation aggregation = entry.getValue();
        if (aggregation.getMask().isPresent()) {
            Symbol newMask = symbolAllocator.newSymbol("mask", BOOLEAN);
            Expression expression = and(aggregation.getMask().get().toSymbolReference(), mask.toSymbolReference());
            assignmentsBuilder.put(newMask, expression);
            masks.put(entry.getKey(), newMask);
        } else {
            masks.put(entry.getKey(), mask);
        }
    }
    Assignments maskAssignments = assignmentsBuilder.build();
    if (!maskAssignments.isEmpty()) {
        source = new ProjectNode(idAllocator.getNextId(), source, Assignments.builder().putIdentities(source.getOutputSymbols()).putAll(maskAssignments).build());
    }
    return new AggregationNode(aggregationNode.getId(), source, rewriteWithMasks(aggregationNode.getAggregations(), masks.buildOrThrow()), singleGroupingSet(groupingSymbols), ImmutableList.of(), SINGLE, Optional.empty(), Optional.empty());
}