Start line:  
End line:  

Snippet Preview

Snippet HTML Code

Stack Overflow Questions
  /*
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package com.facebook.presto.sql.planner.optimizations;
 
 
 import java.util.Map;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
 public class SetFlatteningOptimizer
         extends PlanOptimizer
 {
     @Override
     public PlanNode optimize(PlanNode planSession sessionMap<SymbolTypetypesSymbolAllocator symbolAllocatorPlanNodeIdAllocator idAllocator)
     {
         checkNotNull(plan"plan is null");
         checkNotNull(session"session is null");
         checkNotNull(types"types is null");
         checkNotNull(symbolAllocator"symbolAllocator is null");
         checkNotNull(idAllocator"idAllocator is null");
 
         return PlanRewriter.rewriteWith(new Rewriter(), planfalse);
     }
 
     // TODO: remove expectation that UNION DISTINCT => distinct aggregation directly above union node
     private static class Rewriter
             extends PlanRewriter<Boolean>
     {
         @Override
         public PlanNode visitPlan(PlanNode nodeRewriteContext<Booleancontext)
         {
             return context.defaultRewrite(nodefalse);
         }
 
         @Override
         public PlanNode visitUnion(UnionNode nodeRewriteContext<Booleancontext)
         {
             ImmutableList.Builder<PlanNodeflattenedSources = ImmutableList.builder();
             ImmutableListMultimap.Builder<SymbolSymbolflattenedSymbolMap = ImmutableListMultimap.builder();
             for (int i = 0; i < node.getSources().size(); i++) {
                 PlanNode subplan = node.getSources().get(i);
                 PlanNode rewrittenSource = context.rewrite(subplancontext.get());
 
                 if (rewrittenSource instanceof UnionNode) {
                     // Absorb source's subplans if it is also a UnionNode
                     UnionNode rewrittenUnion = (UnionNoderewrittenSource;
                     flattenedSources.addAll(rewrittenUnion.getSources());
                     for (Map.Entry<SymbolCollection<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<SymbolCollection<Symbol>> entry : node.getSymbolMapping().asMap().entrySet()) {
                         flattenedSymbolMap.put(entry.getKey(), Iterables.get(entry.getValue(), i));
                     }
                 }
             }
             return new UnionNode(node.getId(), flattenedSources.build(), flattenedSymbolMap.build());
         }
 
         @Override
         public PlanNode visitAggregation(AggregationNode nodeRewriteContext<Booleancontext)
         {
             boolean distinct = isDistinctOperator(node);
 
             PlanNode rewrittenNode = context.rewrite(node.getSource(), distinct);
 
             if (context.get() && distinct) {
                 // Assumes underlying node has same output symbols as this distinct node
                 return rewrittenNode;
             }
 
             return new AggregationNode(node.getId(), rewrittenNodenode.getGroupBy(), node.getAggregations(), node.getFunctions(), node.getMasks(), node.getSampleWeight(), node.getConfidence(), node.getHashSymbol());
        }
        private static boolean isDistinctOperator(AggregationNode node)
        {
            return node.getAggregations().isEmpty();
        }
    }
New to GrepCode? Check out our FAQ X