Examples with Lookup io.trino.sql.planner.iterative.Lookup used on opensource projects

Example 6 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class RuleAssert method applyRule.

private RuleApplication applyRule() {
    SymbolAllocator symbolAllocator = new SymbolAllocator(types.allTypes());
    Memo memo = new Memo(idAllocator, plan);
    Lookup lookup = Lookup.from(planNode -> Stream.of(memo.resolve(planNode)));
    PlanNode memoRoot = memo.getNode(memo.getRootGroup());
    return inTransaction(session -> applyRule(rule, memoRoot, ruleContext(statsCalculator, costCalculator, symbolAllocator, memo, lookup, session)));
}

Example 7 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class PruneDistinctAggregation method apply.

@Override
public Result apply(AggregationNode node, Captures captures, Context context) {
    Lookup lookup = context.getLookup();
    DistinctAggregationRewriter rewriter = new DistinctAggregationRewriter(lookup);
    List<PlanNode> newSources = node.getSources().stream().flatMap(lookup::resolveGroup).map(source -> source.accept(rewriter, true)).collect(Collectors.toList());
    if (rewriter.isRewritten()) {
        return Result.ofPlanNode(replaceChildren(node, newSources));
    } else {
        return Result.empty();
    }
}

Example 8 with Lookup

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

Example 9 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class SetOperationMerge method merge.

/**
 * Merge all matching source nodes. This method is assumed to be used only for associative set operations: UNION and INTERSECT.
 *
 * @return Merged plan node if applied.
 */
public Optional<SetOperationNode> merge() {
    checkState(node instanceof UnionNode || node instanceof IntersectNode, "unexpected node type: %s", node.getClass().getSimpleName());
    Lookup lookup = context.getLookup();
    List<PlanNode> sources = node.getSources().stream().flatMap(lookup::resolveGroup).collect(Collectors.toList());
    ImmutableListMultimap.Builder<Symbol, Symbol> newMappingsBuilder = ImmutableListMultimap.builder();
    boolean resultIsDistinct = false;
    boolean rewritten = false;
    for (int i = 0; i < sources.size(); i++) {
        PlanNode source = sources.get(i);
        // Determine if set operations can be merged and whether the resulting set operation is quantified DISTINCT or ALL
        Optional<Boolean> mergedQuantifier = mergedQuantifierIsDistinct(node, source);
        if (mergedQuantifier.isPresent()) {
            addMergedMappings((SetOperationNode) source, i, newMappingsBuilder);
            resultIsDistinct |= mergedQuantifier.get();
            rewritten = true;
        } else {
            // Keep mapping as it is
            addOriginalMappings(source, i, newMappingsBuilder);
        }
    }
    if (!rewritten) {
        return Optional.empty();
    }
    if (node instanceof UnionNode) {
        return Optional.of(new UnionNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols()));
    }
    return Optional.of(new IntersectNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols(), resultIsDistinct));
}