private static List<TableScanNode> findTableScanNodes(PlanNode node) { return searchFrom(node) .where(TableScanNode.class::isInstance) .findAll(); } }
private static Optional<AggregationNode> findAggregation(PlanNode rootNode, Lookup lookup) { return searchFrom(rootNode, lookup) .where(AggregationNode.class::isInstance) .recurseOnlyWhen(isInstanceOfAny(ProjectNode.class, EnforceSingleRowNode.class)) .findFirst(); } }
/** * Use it in optimizer {@link com.facebook.presto.sql.planner.iterative.Rule} only if you truly do not have a better option * <p> * TODO: replace it with a support for plan (physical) properties in rules pattern matching */ public static PlanNodeSearcher searchFrom(PlanNode node, Lookup lookup) { return new PlanNodeSearcher(node, lookup); }
@Override public void validate(PlanNode planNode, Session session, Metadata metadata, SqlParser sqlParser, TypeProvider types, WarningCollector warningCollector) { Map<PlanNodeId, PlanNode> planNodeIds = new HashMap<>(); searchFrom(planNode) .findAll() .forEach(node -> planNodeIds.merge(node.getId(), node, this::reportDuplicateId)); }
private static void assertAtMostOneAggregationBetweenRemoteExchanges(Plan plan) { List<PlanNode> fragments = searchFrom(plan.getRoot()) .where(TestUnion::isRemoteExchange) .findAll() .stream() .flatMap(exchangeNode -> exchangeNode.getSources().stream()) .collect(toList()); for (PlanNode fragment : fragments) { List<PlanNode> aggregations = searchFrom(fragment) .where(AggregationNode.class::isInstance) .recurseOnlyWhen(TestUnion::isNotRemoteExchange) .findAll(); assertFalse(aggregations.size() > 1, "More than a single AggregationNode between remote exchanges"); } }
private static int countOfMatchingNodes(Plan plan, Predicate<PlanNode> predicate) { return searchFrom(plan.getRoot()).where(predicate).count(); }
if (!searchFrom(subquery, context.getLookup()) .where(EnforceSingleRowNode.class::isInstance) .recurseOnlyWhen(ProjectNode.class::isInstance) .matches()) { return Result.empty(); PlanNode rewrittenSubquery = searchFrom(subquery, context.getLookup()) .where(EnforceSingleRowNode.class::isInstance) .recurseOnlyWhen(ProjectNode.class::isInstance) .removeFirst();
private void assertPlanContainsNoFilter(String sql) { assertFalse( searchFrom(plan(sql, LogicalPlanner.Stage.OPTIMIZED).getRoot()) .where(isInstanceOfAny(FilterNode.class)) .matches(), "Unexpected node for query: " + sql); } }
@Test public void testUnionUnderTopN() { Plan plan = plan( "SELECT * FROM (" + " SELECT regionkey FROM nation " + " UNION ALL " + " SELECT nationkey FROM nation" + ") t(a) " + "ORDER BY a LIMIT 1", LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, false); List<PlanNode> remotes = searchFrom(plan.getRoot()) .where(TestUnion::isRemoteExchange) .findAll(); assertEquals(remotes.size(), 1, "There should be exactly one RemoteExchange"); assertEquals(((ExchangeNode) Iterables.getOnlyElement(remotes)).getType(), GATHER); int numberOfpartialTopN = searchFrom(plan.getRoot()) .where(planNode -> planNode instanceof TopNNode && ((TopNNode) planNode).getStep().equals(TopNNode.Step.PARTIAL)) .count(); assertEquals(numberOfpartialTopN, 2, "There should be exactly two partial TopN nodes"); assertPlanIsFullyDistributed(plan); }
private Constraint<ColumnHandle> getConstraint(Plan plan) { Optional<TableScanNode> scanNode = searchFrom(plan.getRoot()) .where(TableScanNode.class::isInstance) .findSingle(); if (!scanNode.isPresent()) { return Constraint.alwaysFalse(); } return new Constraint<>(scanNode.get().getCurrentConstraint()); }
@Override public PlanNode optimize(PlanNode plan, Session session, TypeProvider types, SymbolAllocator symbolAllocator, PlanNodeIdAllocator idAllocator, WarningCollector warningCollector) { searchFrom(plan).where(ApplyNode.class::isInstance) .findFirst() .ifPresent(node -> { ApplyNode applyNode = (ApplyNode) node; throw error(applyNode.getCorrelation(), applyNode.getOriginSubquery()); }); searchFrom(plan).where(LateralJoinNode.class::isInstance) .findFirst() .ifPresent(node -> { LateralJoinNode lateralJoinNode = (LateralJoinNode) node; throw error(lateralJoinNode.getCorrelation(), lateralJoinNode.getOriginSubquery()); }); return plan; }
public StatsCalculatorAssertion withSourceStats(PlanNodeId planNodeId, PlanNodeStatsEstimate sourceStats) { PlanNode sourceNode = PlanNodeSearcher.searchFrom(planNode).where(node -> node.getId().equals(planNodeId)).findOnlyElement(); sourcesStats.put(sourceNode, sourceStats); return this; }
public <T extends PlanNode> T findOnlyElement() { return getOnlyElement(findAll()); }
public static PlanNodeSearcher searchFrom(PlanNode node) { return searchFrom(node, noLookup()); }
public <T extends PlanNode> List<T> findAll() { ImmutableList.Builder<T> nodes = ImmutableList.builder(); findAllRecursive(node, nodes); return nodes.build(); }
public boolean matches() { return findFirst().isPresent(); }
private void assertPlanIsFullyDistributed(Plan plan) { int numberOfGathers = searchFrom(plan.getRoot()) .where(TestUnion::isRemoteGatheringExchange) .findAll() .size(); if (numberOfGathers == 0) { // there are no "gather" nodes, so the plan is expected to be fully distributed return; } assertTrue( searchFrom(plan.getRoot()) .recurseOnlyWhen(TestUnion::isNotRemoteGatheringExchange) .findAll() .stream() .noneMatch(this::shouldBeDistributed), "There is a node that should be distributed between output and first REMOTE GATHER ExchangeNode"); assertEquals(numberOfGathers, 1, "Only a single REMOTE GATHER was expected"); }
private static int countFinalAggregationNodes(Plan plan) { return searchFrom(plan.getRoot()) .where(node -> node instanceof AggregationNode && ((AggregationNode) node).getStep() == FINAL) .count(); }
private void assertPlanContainsNoApplyOrAnyJoin(String sql) { assertFalse( searchFrom(plan(sql, LogicalPlanner.Stage.OPTIMIZED).getRoot()) .where(isInstanceOfAny(ApplyNode.class, JoinNode.class, IndexJoinNode.class, SemiJoinNode.class, LateralJoinNode.class)) .matches(), "Unexpected node for query: " + sql); }
private void validateShowStatsSubquery(ShowStats node, Query query, QuerySpecification querySpecification, Plan plan) { // The following properties of SELECT subquery are required: // - only one relation in FROM // - only predicates that can be pushed down can be in the where clause // - no group by // - no having // - no set quantifier Optional<FilterNode> filterNode = searchFrom(plan.getRoot()) .where(FilterNode.class::isInstance) .findSingle(); check(!filterNode.isPresent(), node, "Only predicates that can be pushed down are supported in the SHOW STATS WHERE clause"); check(querySpecification.getFrom().isPresent(), node, "There must be exactly one table in query passed to SHOW STATS SELECT clause"); check(querySpecification.getFrom().get() instanceof Table, node, "There must be exactly one table in query passed to SHOW STATS SELECT clause"); check(!query.getWith().isPresent(), node, "WITH is not supported by SHOW STATS SELECT clause"); check(!querySpecification.getOrderBy().isPresent(), node, "ORDER BY is not supported in SHOW STATS SELECT clause"); check(!querySpecification.getLimit().isPresent(), node, "LIMIT is not supported by SHOW STATS SELECT clause"); check(!querySpecification.getHaving().isPresent(), node, "HAVING is not supported in SHOW STATS SELECT clause"); check(!querySpecification.getGroupBy().isPresent(), node, "GROUP BY is not supported in SHOW STATS SELECT clause"); check(!querySpecification.getSelect().isDistinct(), node, "DISTINCT is not supported by SHOW STATS SELECT clause"); List<SelectItem> selectItems = querySpecification.getSelect().getSelectItems(); check(selectItems.size() == 1 && selectItems.get(0) instanceof AllColumns, node, "Only SELECT * is supported in SHOW STATS SELECT clause"); }