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

Example 26 with TableScanNode

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

the class TestCostCalculator method testUnion.

@Test
public void testUnion() {
    TableScanNode ts1 = tableScan("ts1", "orderkey");
    TableScanNode ts2 = tableScan("ts2", "orderkey_0");
    ImmutableListMultimap.Builder<Symbol, Symbol> outputMappings = ImmutableListMultimap.builder();
    outputMappings.put(new Symbol("orderkey_1"), new Symbol("orderkey"));
    outputMappings.put(new Symbol("orderkey_1"), new Symbol("orderkey_0"));
    UnionNode union = new UnionNode(new PlanNodeId("union"), ImmutableList.of(ts1, ts2), outputMappings.build(), ImmutableList.of(new Symbol("orderkey_1")));
    Map<String, PlanNodeStatsEstimate> stats = ImmutableMap.of("ts1", statsEstimate(ts1, 4000), "ts2", statsEstimate(ts2, 1000), "union", statsEstimate(ts1, 5000));
    Map<String, PlanCostEstimate> costs = ImmutableMap.of("ts1", cpuCost(1000), "ts2", cpuCost(1000));
    Map<String, Type> types = ImmutableMap.of("orderkey", BIGINT, "orderkey_0", BIGINT, "orderkey_1", BIGINT);
    assertCost(union, costs, stats, types).cpu(2000).memory(0).network(0);
    assertCostEstimatedExchanges(union, costs, stats, types).cpu(2000).memory(0).network(5000 * IS_NULL_OVERHEAD);
}

Example 27 with TableScanNode

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

the class TestSourcePartitionedScheduler method createFragment.

private static PlanFragment createFragment() {
    Symbol symbol = new Symbol("column");
    Symbol buildSymbol = new Symbol("buildColumn");
    // table scan with splitCount splits
    TableScanNode tableScan = TableScanNode.newInstance(TABLE_SCAN_NODE_ID, TEST_TABLE_HANDLE, ImmutableList.of(symbol), ImmutableMap.of(symbol, new TestingColumnHandle("column")), false, Optional.empty());
    FilterNode filterNode = new FilterNode(new PlanNodeId("filter_node_id"), tableScan, createDynamicFilterExpression(TEST_SESSION, createTestMetadataManager(), DYNAMIC_FILTER_ID, VARCHAR, symbol.toSymbolReference()));
    RemoteSourceNode remote = new RemoteSourceNode(new PlanNodeId("remote_id"), new PlanFragmentId("plan_fragment_id"), ImmutableList.of(buildSymbol), Optional.empty(), REPLICATE, RetryPolicy.NONE);
    return new PlanFragment(new PlanFragmentId("plan_id"), new JoinNode(new PlanNodeId("join_id"), INNER, filterNode, remote, ImmutableList.of(), tableScan.getOutputSymbols(), remote.getOutputSymbols(), false, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableMap.of(DYNAMIC_FILTER_ID, buildSymbol), Optional.empty()), ImmutableMap.of(symbol, VARCHAR), SOURCE_DISTRIBUTION, ImmutableList.of(TABLE_SCAN_NODE_ID), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty());
}

Example 28 with TableScanNode

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

the class ApplyTableScanRedirection method apply.

@Override
public Result apply(TableScanNode scanNode, Captures captures, Context context) {
    Optional<TableScanRedirectApplicationResult> tableScanRedirectApplicationResult = plannerContext.getMetadata().applyTableScanRedirect(context.getSession(), scanNode.getTable());
    if (tableScanRedirectApplicationResult.isEmpty()) {
        return Result.empty();
    }
    CatalogSchemaTableName destinationTable = tableScanRedirectApplicationResult.get().getDestinationTable();
    QualifiedObjectName destinationObjectName = convertFromSchemaTableName(destinationTable.getCatalogName()).apply(destinationTable.getSchemaTableName());
    Optional<QualifiedObjectName> redirectedObjectName = plannerContext.getMetadata().getRedirectionAwareTableHandle(context.getSession(), destinationObjectName).getRedirectedTableName();
    redirectedObjectName.ifPresent(name -> {
        throw new TrinoException(NOT_SUPPORTED, format("Further redirection of destination table '%s' to '%s' is not supported", destinationObjectName, name));
    });
    TableMetadata tableMetadata = plannerContext.getMetadata().getTableMetadata(context.getSession(), scanNode.getTable());
    CatalogSchemaTableName sourceTable = new CatalogSchemaTableName(tableMetadata.getCatalogName().getCatalogName(), tableMetadata.getTable());
    if (destinationTable.equals(sourceTable)) {
        return Result.empty();
    }
    Optional<TableHandle> destinationTableHandle = plannerContext.getMetadata().getTableHandle(context.getSession(), convertFromSchemaTableName(destinationTable.getCatalogName()).apply(destinationTable.getSchemaTableName()));
    if (destinationTableHandle.isEmpty()) {
        throw new TrinoException(TABLE_NOT_FOUND, format("Destination table %s from table scan redirection not found", destinationTable));
    }
    Map<ColumnHandle, String> columnMapping = tableScanRedirectApplicationResult.get().getDestinationColumns();
    Map<String, ColumnHandle> destinationColumnHandles = plannerContext.getMetadata().getColumnHandles(context.getSession(), destinationTableHandle.get());
    ImmutableMap.Builder<Symbol, Cast> casts = ImmutableMap.builder();
    ImmutableMap.Builder<Symbol, ColumnHandle> newAssignmentsBuilder = ImmutableMap.builder();
    for (Map.Entry<Symbol, ColumnHandle> assignment : scanNode.getAssignments().entrySet()) {
        String destinationColumn = columnMapping.get(assignment.getValue());
        if (destinationColumn == null) {
            throw new TrinoException(COLUMN_NOT_FOUND, format("Did not find mapping for source column %s in table scan redirection", assignment.getValue()));
        }
        ColumnHandle destinationColumnHandle = destinationColumnHandles.get(destinationColumn);
        if (destinationColumnHandle == null) {
            throw new TrinoException(COLUMN_NOT_FOUND, format("Did not find handle for column %s in destination table %s", destinationColumn, destinationTable));
        }
        // insert ts if redirected types don't match source types
        Type sourceType = context.getSymbolAllocator().getTypes().get(assignment.getKey());
        Type redirectedType = plannerContext.getMetadata().getColumnMetadata(context.getSession(), destinationTableHandle.get(), destinationColumnHandle).getType();
        if (!sourceType.equals(redirectedType)) {
            Symbol redirectedSymbol = context.getSymbolAllocator().newSymbol(destinationColumn, redirectedType);
            Cast cast = getCast(context.getSession(), destinationTable, destinationColumn, redirectedType, redirectedSymbol, sourceTable, assignment.getValue(), sourceType);
            casts.put(assignment.getKey(), cast);
            newAssignmentsBuilder.put(redirectedSymbol, destinationColumnHandle);
        } else {
            newAssignmentsBuilder.put(assignment.getKey(), destinationColumnHandle);
        }
    }
    TupleDomain<String> requiredFilter = tableScanRedirectApplicationResult.get().getFilter();
    if (requiredFilter.isAll()) {
        ImmutableMap<Symbol, ColumnHandle> newAssignments = newAssignmentsBuilder.buildOrThrow();
        return Result.ofPlanNode(applyProjection(context.getIdAllocator(), ImmutableSet.copyOf(scanNode.getOutputSymbols()), casts.buildOrThrow(), new TableScanNode(scanNode.getId(), destinationTableHandle.get(), ImmutableList.copyOf(newAssignments.keySet()), newAssignments, TupleDomain.all(), // Use table statistics from destination table
        Optional.empty(), scanNode.isUpdateTarget(), Optional.empty())));
    }
    Map<ColumnHandle, Symbol> inverseAssignments = ImmutableBiMap.copyOf(scanNode.getAssignments()).inverse();
    Map<String, ColumnHandle> inverseColumnsMapping = ImmutableBiMap.copyOf(columnMapping).inverse();
    TupleDomain<Symbol> transformedConstraint = requiredFilter.transformKeys(destinationColumn -> {
        ColumnHandle sourceColumnHandle = inverseColumnsMapping.get(destinationColumn);
        if (sourceColumnHandle == null) {
            throw new TrinoException(COLUMN_NOT_FOUND, format("Did not find mapping for destination column %s in table scan redirection", destinationColumn));
        }
        Symbol symbol = inverseAssignments.get(sourceColumnHandle);
        if (symbol != null) {
            // domain symbol should already be mapped in redirected table scan
            return symbol;
        }
        // Column pruning after predicate is pushed into table scan can remove assignments for filter columns from the scan node
        Type domainType = requiredFilter.getDomains().get().get(destinationColumn).getType();
        symbol = context.getSymbolAllocator().newSymbol(destinationColumn, domainType);
        ColumnHandle destinationColumnHandle = destinationColumnHandles.get(destinationColumn);
        if (destinationColumnHandle == null) {
            throw new TrinoException(COLUMN_NOT_FOUND, format("Did not find handle for column %s in destination table %s", destinationColumn, destinationTable));
        }
        // insert casts if redirected types don't match domain types
        Type redirectedType = plannerContext.getMetadata().getColumnMetadata(context.getSession(), destinationTableHandle.get(), destinationColumnHandle).getType();
        if (!domainType.equals(redirectedType)) {
            Symbol redirectedSymbol = context.getSymbolAllocator().newSymbol(destinationColumn, redirectedType);
            Cast cast = getCast(context.getSession(), destinationTable, destinationColumn, redirectedType, redirectedSymbol, sourceTable, sourceColumnHandle, domainType);
            casts.put(symbol, cast);
            newAssignmentsBuilder.put(redirectedSymbol, destinationColumnHandle);
        } else {
            newAssignmentsBuilder.put(symbol, destinationColumnHandle);
        }
        return symbol;
    });
    Map<Symbol, ColumnHandle> newAssignments = newAssignmentsBuilder.buildOrThrow();
    TableScanNode newScanNode = new TableScanNode(scanNode.getId(), destinationTableHandle.get(), ImmutableList.copyOf(newAssignments.keySet()), newAssignments, TupleDomain.all(), // Use table statistics from destination table
    Optional.empty(), scanNode.isUpdateTarget(), Optional.empty());
    DomainTranslator domainTranslator = new DomainTranslator(plannerContext);
    FilterNode filterNode = new FilterNode(context.getIdAllocator().getNextId(), applyProjection(context.getIdAllocator(), newAssignments.keySet(), casts.buildOrThrow(), newScanNode), domainTranslator.toPredicate(context.getSession(), transformedConstraint));
    return Result.ofPlanNode(applyProjection(context.getIdAllocator(), ImmutableSet.copyOf(scanNode.getOutputSymbols()), ImmutableMap.of(), filterNode));
}

Example 29 with TableScanNode

use of io.trino.sql.planner.plan.TableScanNode 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 30 with TableScanNode

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

the class PruneTableScanColumns method pruneColumns.

public static Optional<PlanNode> pruneColumns(Metadata metadata, TypeProvider types, Session session, TableScanNode node, Set<Symbol> referencedOutputs) {
    List<Symbol> newOutputs = filteredCopy(node.getOutputSymbols(), referencedOutputs::contains);
    if (newOutputs.size() == node.getOutputSymbols().size()) {
        return Optional.empty();
    }
    List<ConnectorExpression> projections = newOutputs.stream().map(symbol -> new Variable(symbol.getName(), types.get(symbol))).collect(toImmutableList());
    TableHandle handle = node.getTable();
    Optional<ProjectionApplicationResult<TableHandle>> result = metadata.applyProjection(session, handle, projections, newOutputs.stream().collect(toImmutableMap(Symbol::getName, node.getAssignments()::get)));
    Map<Symbol, ColumnHandle> newAssignments;
    // Bail out if the connector does anything other than limit the list of columns (e.g., if it synthesizes arbitrary expressions)
    if (result.isPresent() && result.get().getProjections().stream().allMatch(Variable.class::isInstance)) {
        handle = result.get().getHandle();
        Map<String, ColumnHandle> assignments = result.get().getAssignments().stream().collect(toImmutableMap(Assignment::getVariable, Assignment::getColumn));
        ImmutableMap.Builder<Symbol, ColumnHandle> builder = ImmutableMap.builder();
        for (int i = 0; i < newOutputs.size(); i++) {
            Variable variable = (Variable) result.get().getProjections().get(i);
            builder.put(newOutputs.get(i), assignments.get(variable.getName()));
        }
        newAssignments = builder.buildOrThrow();
    } else {
        newAssignments = newOutputs.stream().collect(toImmutableMap(Function.identity(), node.getAssignments()::get));
    }
    Set<ColumnHandle> visibleColumns = ImmutableSet.copyOf(newAssignments.values());
    TupleDomain<ColumnHandle> enforcedConstraint = node.getEnforcedConstraint().filter((columnHandle, domain) -> visibleColumns.contains(columnHandle));
    Optional<PlanNodeStatsEstimate> newStatistics = node.getStatistics().map(statistics -> new PlanNodeStatsEstimate(statistics.getOutputRowCount(), statistics.getSymbolStatistics().entrySet().stream().filter(entry -> newAssignments.containsKey(entry.getKey())).collect(toImmutableMap(Entry::getKey, Entry::getValue))));
    return Optional.of(new TableScanNode(node.getId(), handle, newOutputs, newAssignments, enforcedConstraint, newStatistics, node.isUpdateTarget(), node.getUseConnectorNodePartitioning()));
}