private Function<Symbol, Optional<Symbol>> outputToInputTranslator(UnionNode node, int sourceIndex) { return symbol -> Optional.of(node.getSymbolMapping().get(symbol).get(sourceIndex)); }
private static Map<Symbol, Symbol> getInputSymbolMapping(UnionNode node, int source) { return node.getSymbolMapping() .keySet() .stream() .collect(toImmutableMap(key -> key, key -> node.getSymbolMapping().get(key).get(source))); } }
@Override protected final Optional<PlanNodeStatsEstimate> doCalculate(UnionNode node, StatsProvider statsProvider, Lookup lookup, Session session, TypeProvider types) { checkArgument(!node.getSources().isEmpty(), "Empty Union is not supported"); Optional<PlanNodeStatsEstimate> estimate = Optional.empty(); for (int i = 0; i < node.getSources().size(); i++) { PlanNode source = node.getSources().get(i); PlanNodeStatsEstimate sourceStats = statsProvider.getStats(source); PlanNodeStatsEstimate sourceStatsWithMappedSymbols = mapToOutputSymbols(sourceStats, node.getSymbolMapping(), i); if (estimate.isPresent()) { estimate = Optional.of(addStatsAndCollapseDistinctValues(estimate.get(), sourceStatsWithMappedSymbols)); } else { estimate = Optional.of(sourceStatsWithMappedSymbols); } } return estimate; }
@Override public PlanNode replaceChildren(List<PlanNode> newChildren) { return new UnionNode(getId(), newChildren, getSymbolMapping(), getOutputSymbols()); } }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<Expression> context) { boolean modified = false; ImmutableList.Builder<PlanNode> builder = ImmutableList.builder(); for (int i = 0; i < node.getSources().size(); i++) { Expression sourcePredicate = inlineSymbols(node.sourceSymbolMap(i), context.get()); PlanNode source = node.getSources().get(i); PlanNode rewrittenSource = context.rewrite(source, sourcePredicate); if (rewrittenSource != source) { modified = true; } builder.add(rewrittenSource); } if (modified) { return new UnionNode(node.getId(), builder.build(), node.getSymbolMapping(), node.getOutputSymbols()); } return node; }
@Override public Result apply(TopNNode topNNode, Captures captures, Context context) { UnionNode unionNode = captures.get(CHILD); ImmutableList.Builder<PlanNode> sources = ImmutableList.builder(); for (PlanNode source : unionNode.getSources()) { SymbolMapper.Builder symbolMapper = SymbolMapper.builder(); Set<Symbol> sourceOutputSymbols = ImmutableSet.copyOf(source.getOutputSymbols()); for (Symbol unionOutput : unionNode.getOutputSymbols()) { Set<Symbol> inputSymbols = ImmutableSet.copyOf(unionNode.getSymbolMapping().get(unionOutput)); Symbol unionInput = getLast(intersection(inputSymbols, sourceOutputSymbols)); symbolMapper.put(unionOutput, unionInput); } sources.add(symbolMapper.build().map(topNNode, source, context.getIdAllocator().getNextId())); } return Result.ofPlanNode(new UnionNode( unionNode.getId(), sources.build(), unionNode.getSymbolMapping(), unionNode.getOutputSymbols())); } }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<Void> context) { return new UnionNode(node.getId(), rewriteSources(node, context).build(), canonicalizeSetOperationSymbolMap(node.getSymbolMapping()), canonicalizeAndDistinct(node.getOutputSymbols())); }
newSymbolMapping.putAll(node.getSymbolMapping()); ImmutableList.Builder<PlanNode> newSources = ImmutableList.builder(); for (int sourceId = 0; sourceId < node.getSources().size(); sourceId++) { outputToInputMap.put(outputSymbol, node.getSymbolMapping().get(outputSymbol).get(sourceId));
@Override public Void visitUnion(UnionNode node, Void context) { visitPlan(node, context); ListMultimap<Symbol, Symbol> symbolMapping = node.getSymbolMapping(); for (Symbol keySymbol : symbolMapping.keySet()) { List<Symbol> valueSymbols = symbolMapping.get(keySymbol); Type expectedType = types.get(keySymbol); for (Symbol valueSymbol : valueSymbols) { verifyTypeSignature(keySymbol, expectedType.getTypeSignature(), types.get(valueSymbol).getTypeSignature()); } } return null; }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<LimitContext> context) { LimitContext limit = context.get(); LimitContext childLimit = null; if (limit != null) { childLimit = new LimitContext(limit.getCount(), true); } List<PlanNode> sources = new ArrayList<>(); for (int i = 0; i < node.getSources().size(); i++) { sources.add(context.rewrite(node.getSources().get(i), childLimit)); } PlanNode output = new UnionNode(node.getId(), sources, node.getSymbolMapping(), node.getOutputSymbols()); if (limit != null) { output = new LimitNode(idAllocator.getNextId(), output, limit.getCount(), limit.isPartial()); } return output; }
@Override public PlanNode visitUnion(UnionNode node, List<PlanNode> newChildren) { return new UnionNode(node.getId(), newChildren, node.getSymbolMapping(), node.getOutputSymbols()); }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<Boolean> context) { ImmutableList.Builder<PlanNode> flattenedSources = ImmutableList.builder(); ImmutableListMultimap.Builder<Symbol, Symbol> flattenedSymbolMap = ImmutableListMultimap.builder(); for (int i = 0; i < node.getSources().size(); i++) { PlanNode subplan = node.getSources().get(i); PlanNode rewrittenSource = context.rewrite(subplan, context.get()); if (rewrittenSource instanceof UnionNode) { // Absorb source's subplans if it is also a UnionNode UnionNode rewrittenUnion = (UnionNode) rewrittenSource; flattenedSources.addAll(rewrittenUnion.getSources()); for (Map.Entry<Symbol, Collection<Symbol>> entry : node.getSymbolMapping().asMap().entrySet()) { Symbol inputSymbol = Iterables.get(entry.getValue(), i); flattenedSymbolMap.putAll(entry.getKey(), rewrittenUnion.getSymbolMapping().get(inputSymbol)); } } else { flattenedSources.add(rewrittenSource); for (Map.Entry<Symbol, Collection<Symbol>> entry : node.getSymbolMapping().asMap().entrySet()) { flattenedSymbolMap.put(entry.getKey(), Iterables.get(entry.getValue(), i)); } } } return new UnionNode(node.getId(), flattenedSources.build(), flattenedSymbolMap.build(), ImmutableList.copyOf(flattenedSymbolMap.build().keySet())); }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<Void> context) { ImmutableList.Builder<PlanNode> rewrittenSources = ImmutableList.builder(); for (PlanNode source : node.getSources()) { rewrittenSources.add(context.rewrite(source)); } return new UnionNode(node.getId(), rewrittenSources.build(), canonicalizeUnionSymbolMap(node.getSymbolMapping()), canonicalize(node.getOutputSymbols())); }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<Set<Symbol>> context) { // Find out which output symbols we need to keep ImmutableListMultimap.Builder<Symbol, Symbol> rewrittenSymbolMappingBuilder = ImmutableListMultimap.builder(); for (Symbol symbol : node.getOutputSymbols()) { if (context.get().contains(symbol)) { rewrittenSymbolMappingBuilder.putAll(symbol, node.getSymbolMapping().get(symbol)); } } ListMultimap<Symbol, Symbol> rewrittenSymbolMapping = rewrittenSymbolMappingBuilder.build(); // Find the corresponding input symbol to the remaining output symbols and prune the subplans ImmutableList.Builder<PlanNode> rewrittenSubPlans = ImmutableList.builder(); for (int i = 0; i < node.getSources().size(); i++) { ImmutableSet.Builder<Symbol> expectedInputSymbols = ImmutableSet.builder(); for (Collection<Symbol> symbols : rewrittenSymbolMapping.asMap().values()) { expectedInputSymbols.add(Iterables.get(symbols, i)); } rewrittenSubPlans.add(context.rewrite(node.getSources().get(i), expectedInputSymbols.build())); } return new UnionNode(node.getId(), rewrittenSubPlans.build(), rewrittenSymbolMapping, ImmutableList.copyOf(rewrittenSymbolMapping.keySet())); }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<LimitContext> context) { LimitContext limit = context.get(); List<PlanNode> sources = new ArrayList<>(); for (int i = 0; i < node.getSources().size(); i++) { sources.add(context.rewrite(node.getSources().get(i), limit)); } PlanNode output = new UnionNode(node.getId(), sources, node.getSymbolMapping(), node.getOutputSymbols()); if (limit != null) { output = new LimitNode(idAllocator.getNextId(), output, limit.getCount()); } return output; }
@Override public PlanNode visitUnion(UnionNode node, RewriteContext<Expression> context) { boolean modified = false; ImmutableList.Builder<PlanNode> builder = ImmutableList.builder(); for (int i = 0; i < node.getSources().size(); i++) { Expression sourcePredicate = ExpressionTreeRewriter.rewriteWith(new ExpressionSymbolInliner(node.sourceSymbolMap(i)), context.get()); PlanNode source = node.getSources().get(i); PlanNode rewrittenSource = context.rewrite(source, sourcePredicate); if (rewrittenSource != source) { modified = true; } builder.add(rewrittenSource); } if (modified) { return new UnionNode(node.getId(), builder.build(), node.getSymbolMapping(), node.getOutputSymbols()); } return node; }
ImmutableList.Builder<Symbol> hashColumnsBuilder = ImmutableList.builder(); for (Symbol column : hashingColumns) { hashColumnsBuilder.add(node.getSymbolMapping().get(column).get(sourceIndex));