@Override public PlanNode visitMarkDistinct(MarkDistinctNode node, RewriteContext<Expression> context) { Set<Symbol> pushDownableSymbols = ImmutableSet.copyOf(node.getDistinctSymbols()); Map<Boolean, List<Expression>> conjuncts = extractConjuncts(context.get()).stream() .collect(Collectors.partitioningBy(conjunct -> SymbolsExtractor.extractUnique(conjunct).stream().allMatch(pushDownableSymbols::contains))); PlanNode rewrittenNode = context.defaultRewrite(node, combineConjuncts(conjuncts.get(true))); if (!conjuncts.get(false).isEmpty()) { rewrittenNode = new FilterNode(idAllocator.getNextId(), rewrittenNode, combineConjuncts(conjuncts.get(false))); } return rewrittenNode; }
@Override public PlanNode visitTableScan(TableScanNode node, RewriteContext<FragmentProperties> context) { PartitioningHandle partitioning = node.getLayout() .map(layout -> metadata.getLayout(session, layout)) .flatMap(TableLayout::getTablePartitioning) .map(TablePartitioning::getPartitioningHandle) .orElse(SOURCE_DISTRIBUTION); context.get().addSourceDistribution(node.getId(), partitioning, metadata, session); return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitAggregation(AggregationNode node, RewriteContext<Context> context) { // Lookup symbols can only be passed through if they are part of the group by columns Set<Symbol> groupByLookupSymbols = context.get().getLookupSymbols().stream() .filter(node.getGroupingKeys()::contains) .collect(toImmutableSet()); if (groupByLookupSymbols.isEmpty()) { return node; } return context.defaultRewrite(node, new Context(groupByLookupSymbols, context.get().getSuccess())); }
@Override public PlanNode visitAssignUniqueId(AssignUniqueId node, RewriteContext<Expression> context) { Set<Symbol> predicateSymbols = SymbolsExtractor.extractUnique(context.get()); checkState(!predicateSymbols.contains(node.getIdColumn()), "UniqueId in predicate is not yet supported"); return context.defaultRewrite(node, context.get()); } }
/** * Invoke the rewrite logic recursively on children of the given node and swap it * out with an identical copy with the rewritten children */ public PlanNode defaultRewrite(PlanNode node) { return defaultRewrite(node, null); }
@Override public PlanNode visitOutput(OutputNode node, RewriteContext<FragmentProperties> context) { if (isForceSingleNodeOutput(session)) { context.get().setSingleNodeDistribution(); } return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitPlan(PlanNode node, RewriteContext<Boolean> context) { return context.defaultRewrite(node, false); }
@Override public PlanNode visitFilter(FilterNode node, RewriteContext<Context> context) { if (node.getSource() instanceof TableScanNode) { return planTableScan((TableScanNode) node.getSource(), node.getPredicate(), context.get()); } return context.defaultRewrite(node, new Context(context.get().getLookupSymbols(), context.get().getSuccess())); }
@Override public PlanNode visitSample(SampleNode node, RewriteContext<Expression> context) { return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitPlan(PlanNode node, RewriteContext<Expression> context) { PlanNode rewrittenNode = context.defaultRewrite(node, TRUE_LITERAL); if (!context.get().equals(TRUE_LITERAL)) { // Drop in a FilterNode b/c we cannot push our predicate down any further rewrittenNode = new FilterNode(idAllocator.getNextId(), rewrittenNode, context.get()); } return rewrittenNode; }
@Override protected PlanNode visitPlan(PlanNode node, RewriteContext<C> context) { return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitTableWriter(TableWriterNode node, RewriteContext<FragmentProperties> context) { if (node.getPartitioningScheme().isPresent()) { context.get().setDistribution(node.getPartitioningScheme().get().getPartitioning().getHandle(), metadata, session); } return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitFilter(FilterNode node, RewriteContext<Void> context) { FilterNode rewrittenNode = (FilterNode) context.defaultRewrite(node); return new FilterNode(idAllocator.getNextId(), rewrittenNode.getSource(), replaceExpression(rewrittenNode.getPredicate(), mapping)); }
@Override public PlanNode visitPlan(PlanNode node, RewriteContext<LimitContext> context) { PlanNode rewrittenNode = context.defaultRewrite(node); LimitContext limit = context.get(); if (limit != null) { // Drop in a LimitNode b/c we cannot push our limit down any further rewrittenNode = new LimitNode(idAllocator.getNextId(), rewrittenNode, limit.getCount(), limit.isPartial()); } return rewrittenNode; }
@Override public PlanNode visitMetadataDelete(MetadataDeleteNode node, RewriteContext<FragmentProperties> context) { context.get().setCoordinatorOnlyDistribution(); return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitLimit(LimitNode node, RewriteContext<Void> context) { return context.defaultRewrite(node); }
@Override public PlanNode visitMarkDistinct(MarkDistinctNode node, RewriteContext<LimitContext> context) { // the fallback logic (in visitPlan) for node types we don't know about introduces a limit node, // so we need this here to push the limit through this trivial node type return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitSort(SortNode node, RewriteContext<Expression> context) { return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitAssignUniqueId(AssignUniqueId node, RewriteContext<Set<Symbol>> context) { if (!context.get().contains(node.getIdColumn())) { return context.rewrite(node.getSource(), context.get()); } return context.defaultRewrite(node, context.get()); }
@Override public PlanNode visitExplainAnalyze(ExplainAnalyzeNode node, RewriteContext<Set<Symbol>> context) { return context.defaultRewrite(node, ImmutableSet.copyOf(node.getSource().getOutputSymbols())); }