Examples with SymbolAllocator io.trino.sql.planner.SymbolAllocator used on opensource projects

Example 1 with SymbolAllocator

use of io.trino.sql.planner.SymbolAllocator in project trino by trinodb.

the class PushPredicateIntoTableScan method pushFilterIntoTableScan.

public static Optional<PlanNode> pushFilterIntoTableScan(FilterNode filterNode, TableScanNode node, boolean pruneWithPredicateExpression, Session session, SymbolAllocator symbolAllocator, PlannerContext plannerContext, TypeAnalyzer typeAnalyzer, StatsProvider statsProvider, DomainTranslator domainTranslator) {
    if (!isAllowPushdownIntoConnectors(session)) {
        return Optional.empty();
    }
    SplitExpression splitExpression = splitExpression(plannerContext, filterNode.getPredicate());
    DomainTranslator.ExtractionResult decomposedPredicate = DomainTranslator.getExtractionResult(plannerContext, session, splitExpression.getDeterministicPredicate(), symbolAllocator.getTypes());
    TupleDomain<ColumnHandle> newDomain = decomposedPredicate.getTupleDomain().transformKeys(node.getAssignments()::get).intersect(node.getEnforcedConstraint());
    Map<NodeRef<Expression>, Type> remainingExpressionTypes = typeAnalyzer.getTypes(session, symbolAllocator.getTypes(), decomposedPredicate.getRemainingExpression());
    Optional<ConnectorExpression> connectorExpression = new ConnectorExpressionTranslator.SqlToConnectorExpressionTranslator(session, remainingExpressionTypes, plannerContext).process(decomposedPredicate.getRemainingExpression());
    Map<String, ColumnHandle> connectorExpressionAssignments = connectorExpression.map(ignored -> node.getAssignments().entrySet().stream().collect(toImmutableMap(entry -> entry.getKey().getName(), Map.Entry::getValue))).orElse(ImmutableMap.of());
    Map<ColumnHandle, Symbol> assignments = ImmutableBiMap.copyOf(node.getAssignments()).inverse();
    Constraint constraint;
    // use evaluator only when there is some predicate which could not be translated into tuple domain
    if (pruneWithPredicateExpression && !TRUE_LITERAL.equals(decomposedPredicate.getRemainingExpression())) {
        LayoutConstraintEvaluator evaluator = new LayoutConstraintEvaluator(plannerContext, typeAnalyzer, session, symbolAllocator.getTypes(), node.getAssignments(), combineConjuncts(plannerContext.getMetadata(), splitExpression.getDeterministicPredicate(), // which would be expensive to evaluate in the call to isCandidate below.
        domainTranslator.toPredicate(session, newDomain.simplify().transformKeys(assignments::get))));
        constraint = new Constraint(newDomain, connectorExpression.orElse(TRUE), connectorExpressionAssignments, evaluator::isCandidate, evaluator.getArguments());
    } else {
        // Currently, invoking the expression interpreter is very expensive.
        // TODO invoke the interpreter unconditionally when the interpreter becomes cheap enough.
        constraint = new Constraint(newDomain, connectorExpression.orElse(TRUE), connectorExpressionAssignments);
    }
    // check if new domain is wider than domain already provided by table scan
    if (constraint.predicate().isEmpty() && // TODO do we need to track enforced ConnectorExpression in TableScanNode?
    TRUE.equals(connectorExpression.orElse(TRUE)) && newDomain.contains(node.getEnforcedConstraint())) {
        Expression resultingPredicate = createResultingPredicate(plannerContext, session, symbolAllocator, typeAnalyzer, splitExpression.getDynamicFilter(), TRUE_LITERAL, splitExpression.getNonDeterministicPredicate(), decomposedPredicate.getRemainingExpression());
        if (!TRUE_LITERAL.equals(resultingPredicate)) {
            return Optional.of(new FilterNode(filterNode.getId(), node, resultingPredicate));
        }
        return Optional.of(node);
    }
    if (newDomain.isNone()) {
        // to turn the subtree into a Values node
        return Optional.of(new ValuesNode(node.getId(), node.getOutputSymbols(), ImmutableList.of()));
    }
    Optional<ConstraintApplicationResult<TableHandle>> result = plannerContext.getMetadata().applyFilter(session, node.getTable(), constraint);
    if (result.isEmpty()) {
        return Optional.empty();
    }
    TableHandle newTable = result.get().getHandle();
    TableProperties newTableProperties = plannerContext.getMetadata().getTableProperties(session, newTable);
    Optional<TablePartitioning> newTablePartitioning = newTableProperties.getTablePartitioning();
    if (newTableProperties.getPredicate().isNone()) {
        return Optional.of(new ValuesNode(node.getId(), node.getOutputSymbols(), ImmutableList.of()));
    }
    TupleDomain<ColumnHandle> remainingFilter = result.get().getRemainingFilter();
    Optional<ConnectorExpression> remainingConnectorExpression = result.get().getRemainingExpression();
    boolean precalculateStatistics = result.get().isPrecalculateStatistics();
    verifyTablePartitioning(session, plannerContext.getMetadata(), node, newTablePartitioning);
    TableScanNode tableScan = new TableScanNode(node.getId(), newTable, node.getOutputSymbols(), node.getAssignments(), computeEnforced(newDomain, remainingFilter), // TODO (https://github.com/trinodb/trino/issues/8144) distinguish between predicate pushed down and remaining
    deriveTableStatisticsForPushdown(statsProvider, session, precalculateStatistics, filterNode), node.isUpdateTarget(), node.getUseConnectorNodePartitioning());
    Expression remainingDecomposedPredicate;
    if (remainingConnectorExpression.isEmpty() || remainingConnectorExpression.equals(connectorExpression)) {
        remainingDecomposedPredicate = decomposedPredicate.getRemainingExpression();
    } else {
        Map<String, Symbol> variableMappings = assignments.values().stream().collect(toImmutableMap(Symbol::getName, Function.identity()));
        Expression translatedExpression = ConnectorExpressionTranslator.translate(session, remainingConnectorExpression.get(), plannerContext, variableMappings, new LiteralEncoder(plannerContext));
        if (connectorExpression.isEmpty()) {
            remainingDecomposedPredicate = ExpressionUtils.combineConjuncts(plannerContext.getMetadata(), translatedExpression, decomposedPredicate.getRemainingExpression());
        } else {
            remainingDecomposedPredicate = translatedExpression;
        }
    }
    Expression resultingPredicate = createResultingPredicate(plannerContext, session, symbolAllocator, typeAnalyzer, splitExpression.getDynamicFilter(), domainTranslator.toPredicate(session, remainingFilter.transformKeys(assignments::get)), splitExpression.getNonDeterministicPredicate(), remainingDecomposedPredicate);
    if (!TRUE_LITERAL.equals(resultingPredicate)) {
        return Optional.of(new FilterNode(filterNode.getId(), tableScan, resultingPredicate));
    }
    return Optional.of(tableScan);
}

Example 2 with SymbolAllocator

use of io.trino.sql.planner.SymbolAllocator in project trino by trinodb.

the class TestSimplifyExpressions method assertSimplifiesNumericTypes.

private static void assertSimplifiesNumericTypes(String expression, String expected) {
    ParsingOptions parsingOptions = new ParsingOptions();
    Expression actualExpression = rewriteIdentifiersToSymbolReferences(SQL_PARSER.createExpression(expression, parsingOptions));
    Expression expectedExpression = rewriteIdentifiersToSymbolReferences(SQL_PARSER.createExpression(expected, parsingOptions));
    Expression rewritten = rewrite(actualExpression, TEST_SESSION, new SymbolAllocator(numericAndBooleanSymbolTypeMapFor(actualExpression)), PLANNER_CONTEXT, createTestingTypeAnalyzer(PLANNER_CONTEXT));
    assertEquals(normalize(rewritten), normalize(expectedExpression));
}

Example 3 with SymbolAllocator

use of io.trino.sql.planner.SymbolAllocator in project trino by trinodb.

the class TestDynamicFilterService method testMultipleColumnMapping.

@Test
public void testMultipleColumnMapping() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", INTEGER);
    Expression df1 = symbol1.toSymbolReference();
    Expression df2 = symbol2.toSymbolReference();
    QueryId queryId = new QueryId("query");
    StageId stageId1 = new StageId(queryId, 1);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId1), ImmutableSet.of(filterId1), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId1, 0, 1);
    TestingColumnHandle column1 = new TestingColumnHandle("probeColumnA");
    TestingColumnHandle column2 = new TestingColumnHandle("probeColumnB");
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId1, df2)), ImmutableMap.of(symbol1, column1, symbol2, column2), symbolAllocator.getTypes());
    assertEquals(dynamicFilter.getColumnsCovered(), Set.of(column1, column2), "columns covered");
    Domain domain = singleValue(INTEGER, 1L);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 0, 0), ImmutableMap.of(filterId1, domain));
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(column1, domain, column2, domain)));
}

Example 4 with SymbolAllocator

use of io.trino.sql.planner.SymbolAllocator in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilter.

@Test
public void testDynamicFilter() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    DynamicFilterId filterId3 = new DynamicFilterId("df3");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", INTEGER);
    Symbol symbol3 = symbolAllocator.newSymbol("DF_SYMBOL3", INTEGER);
    Expression df1 = symbol1.toSymbolReference();
    Expression df2 = symbol2.toSymbolReference();
    Expression df3 = symbol3.toSymbolReference();
    QueryId queryId = new QueryId("query");
    StageId stageId1 = new StageId(queryId, 1);
    StageId stageId2 = new StageId(queryId, 2);
    StageId stageId3 = new StageId(queryId, 3);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId1, filterId2, filterId3), ImmutableSet.of(filterId1, filterId2, filterId3), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId1, 0, 2);
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId2, 0, 2);
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId3, 0, 2);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2), new DynamicFilters.Descriptor(filterId3, df3)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA"), symbol2, new TestingColumnHandle("probeColumnA"), symbol3, new TestingColumnHandle("probeColumnB")), symbolAllocator.getTypes());
    assertEquals(dynamicFilter.getColumnsCovered(), Set.of(new TestingColumnHandle("probeColumnA"), new TestingColumnHandle("probeColumnB")), "columns covered");
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    // assert initial dynamic filtering stats
    DynamicFiltersStats stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getTotalDynamicFilters(), 3);
    assertEquals(stats.getDynamicFiltersCompleted(), 0);
    assertEquals(stats.getLazyDynamicFilters(), 3);
    assertEquals(stats.getReplicatedDynamicFilters(), 0);
    // dynamic filter should be blocked waiting for tuple domain to be provided
    CompletableFuture<?> blockedFuture = dynamicFilter.isBlocked();
    assertFalse(blockedFuture.isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 0, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 1L)));
    // tuple domain from two tasks are needed for dynamic filter to be narrowed down
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    assertFalse(blockedFuture.isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 1, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 2L)));
    // dynamic filter (id1) has been collected as tuple domains from two tasks have been provided
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), multipleValues(INTEGER, ImmutableList.of(1L, 2L)))));
    assertTrue(blockedFuture.isDone());
    assertFalse(blockedFuture.isCompletedExceptionally());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 1);
    // there are still more dynamic filters to be collected
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    blockedFuture = dynamicFilter.isBlocked();
    assertFalse(blockedFuture.isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId2, 0, 0), ImmutableMap.of(filterId2, singleValue(INTEGER, 2L)));
    // tuple domain from two tasks (stage 2) are needed for dynamic filter to be narrowed down
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), multipleValues(INTEGER, ImmutableList.of(1L, 2L)))));
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    assertFalse(blockedFuture.isDone());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 1);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId2, 1, 0), ImmutableMap.of(filterId2, singleValue(INTEGER, 3L)));
    // dynamic filter (id2) has been collected as tuple domains from two tasks have been provided
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), singleValue(INTEGER, 2L))));
    assertTrue(blockedFuture.isDone());
    assertFalse(blockedFuture.isCompletedExceptionally());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 2);
    // there are still more dynamic filters to be collected for columns A and B
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    blockedFuture = dynamicFilter.isBlocked();
    assertFalse(blockedFuture.isDone());
    // make sure dynamic filter on just column A is now completed
    DynamicFilter dynamicFilterColumnA = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA"), symbol2, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    assertEquals(dynamicFilterColumnA.getColumnsCovered(), Set.of(new TestingColumnHandle("probeColumnA")), "columns covered");
    assertTrue(dynamicFilterColumnA.isComplete());
    assertFalse(dynamicFilterColumnA.isAwaitable());
    assertTrue(dynamicFilterColumnA.isBlocked().isDone());
    assertEquals(dynamicFilterColumnA.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), singleValue(INTEGER, 2L))));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId3, 0, 0), ImmutableMap.of(filterId3, none(INTEGER)));
    // tuple domain from two tasks (stage 3) are needed for dynamic filter to be narrowed down
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), singleValue(INTEGER, 2L))));
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    assertFalse(blockedFuture.isDone());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 2);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId3, 1, 0), ImmutableMap.of(filterId3, none(INTEGER)));
    // "none" dynamic filter (id3) has been collected for column B as tuple domains from two tasks have been provided
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.none());
    assertTrue(blockedFuture.isDone());
    assertFalse(blockedFuture.isCompletedExceptionally());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 3);
    assertEquals(stats.getLazyDynamicFilters(), 3);
    assertEquals(stats.getReplicatedDynamicFilters(), 0);
    assertEquals(ImmutableSet.copyOf(stats.getDynamicFilterDomainStats()), ImmutableSet.of(new DynamicFilterDomainStats(filterId1, getSimplifiedDomainString(1L, 2L, 2, INTEGER)), new DynamicFilterDomainStats(filterId2, getSimplifiedDomainString(2L, 3L, 2, INTEGER)), new DynamicFilterDomainStats(filterId3, Domain.none(INTEGER).toString(session.toConnectorSession()))));
    // all dynamic filters have been collected, no need for more requests
    assertTrue(dynamicFilter.isComplete());
    assertFalse(dynamicFilter.isAwaitable());
    assertTrue(dynamicFilter.isBlocked().isDone());
}

Example 5 with SymbolAllocator

use of io.trino.sql.planner.SymbolAllocator in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilterCoercion.

@Test
public void testDynamicFilterCoercion() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Expression df1 = new Cast(symbol1.toSymbolReference(), toSqlType(BIGINT));
    QueryId queryId = new QueryId("query");
    StageId stageId1 = new StageId(queryId, 1);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId1), ImmutableSet.of(filterId1), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId1, 0, 1);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    assertEquals(dynamicFilter.getColumnsCovered(), Set.of(new TestingColumnHandle("probeColumnA")), "columns covered");
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 0, 0), ImmutableMap.of(filterId1, multipleValues(BIGINT, ImmutableList.of(1L, 2L, 3L))));
    assertTrue(dynamicFilter.isComplete());
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), multipleValues(INTEGER, ImmutableList.of(1L, 2L, 3L)))));
}