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

Example 16 with TypeProvider

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

the class TestExpressionEquivalence method assertNotEquivalent.

private static void assertNotEquivalent(@Language("SQL") String left, @Language("SQL") String right) {
    ParsingOptions parsingOptions = new ParsingOptions(AS_DOUBLE);
    Expression leftExpression = planExpression(PLANNER_CONTEXT, TEST_SESSION, TYPE_PROVIDER, SQL_PARSER.createExpression(left, parsingOptions));
    Expression rightExpression = planExpression(PLANNER_CONTEXT, TEST_SESSION, TYPE_PROVIDER, SQL_PARSER.createExpression(right, parsingOptions));
    Set<Symbol> symbols = extractUnique(ImmutableList.of(leftExpression, rightExpression));
    TypeProvider types = TypeProvider.copyOf(symbols.stream().collect(toMap(identity(), TestExpressionEquivalence::generateType)));
    assertFalse(areExpressionEquivalent(leftExpression, rightExpression, types), format("Expected (%s) and (%s) to not be equivalent", left, right));
    assertFalse(areExpressionEquivalent(rightExpression, leftExpression, types), format("Expected (%s) and (%s) to not be equivalent", right, left));
}

Example 17 with TypeProvider

use of io.trino.sql.planner.TypeProvider 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)));
}

Example 18 with TypeProvider

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

the class PruneTableScanColumns method pushDownProjectOff.

@Override
protected Optional<PlanNode> pushDownProjectOff(Context context, TableScanNode node, Set<Symbol> referencedOutputs) {
    Session session = context.getSession();
    TypeProvider types = context.getSymbolAllocator().getTypes();
    return pruneColumns(metadata, types, session, node, referencedOutputs);
}

Example 19 with TypeProvider

use of io.trino.sql.planner.TypeProvider 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()));
}

Example 20 with TypeProvider

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

the class TopNStatsRule method doCalculate.

@Override
protected Optional<PlanNodeStatsEstimate> doCalculate(TopNNode node, StatsProvider statsProvider, Lookup lookup, Session session, TypeProvider types) {
    PlanNodeStatsEstimate sourceStats = statsProvider.getStats(node.getSource());
    double rowCount = sourceStats.getOutputRowCount();
    if (node.getStep() != TopNNode.Step.SINGLE) {
        return Optional.empty();
    }
    if (rowCount <= node.getCount()) {
        return Optional.of(sourceStats);
    }
    long limitCount = node.getCount();
    PlanNodeStatsEstimate resultStats = PlanNodeStatsEstimate.buildFrom(sourceStats).setOutputRowCount(limitCount).build();
    if (limitCount == 0) {
        return Optional.of(resultStats);
    }
    // augment null fraction estimation for first ORDER BY symbol
    // Assuming not empty list
    Symbol firstOrderSymbol = node.getOrderingScheme().getOrderBy().get(0);
    SortOrder sortOrder = node.getOrderingScheme().getOrdering(firstOrderSymbol);
    resultStats = resultStats.mapSymbolColumnStatistics(firstOrderSymbol, symbolStats -> {
        SymbolStatsEstimate.Builder newStats = SymbolStatsEstimate.buildFrom(symbolStats);
        double nullCount = rowCount * symbolStats.getNullsFraction();
        if (sortOrder.isNullsFirst()) {
            if (nullCount > limitCount) {
                newStats.setNullsFraction(1.0);
            } else {
                newStats.setNullsFraction(nullCount / limitCount);
            }
        } else {
            double nonNullCount = (rowCount - nullCount);
            if (nonNullCount > limitCount) {
                newStats.setNullsFraction(0.0);
            } else {
                newStats.setNullsFraction((limitCount - nonNullCount) / limitCount);
            }
        }
        return newStats.build();
    });
    // TopN actually limits (or when there was no row count estimated for source)
    return Optional.of(resultStats);
}