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.List;
 import java.util.Map;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 public class LimitPushDown
         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(idAllocator), plannull);
     }
 
     private static class LimitContext
     {
         private final long count;
 
         public LimitContext(long count)
         {
             this. = count;
         }
 
         public long getCount()
         {
             return ;
         }
     }
 
     private static class Rewriter
             extends PlanRewriter<LimitContext>
     {
         private final PlanNodeIdAllocator idAllocator;
 
         private Rewriter(PlanNodeIdAllocator idAllocator)
         {
             this. = checkNotNull(idAllocator"idAllocator is null");
         }
 
         @Override
         public PlanNode visitPlan(PlanNode nodeRewriteContext<LimitContextcontext)
         {
             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(.getNextId(), rewrittenNodelimit.getCount());
             }
             return rewrittenNode;
         }
 
         @Override
         public PlanNode visitLimit(LimitNode nodeRewriteContext<LimitContextcontext)
         {
             LimitContext limit = context.get();
             if (limit != null && limit.getCount() < node.getCount()) {
                 return context.rewrite(node.getSource(), limit);
            }
            else {
                return context.rewrite(node.getSource(), new LimitContext(node.getCount()));
            }
        }
        @Override
        public PlanNode visitAggregation(AggregationNode nodeRewriteContext<LimitContextcontext)
        {
            LimitContext limit = context.get();
            if (limit != null &&
                    node.getAggregations().isEmpty() &&
                    node.getOutputSymbols().size() == node.getGroupBy().size() &&
                    node.getOutputSymbols().containsAll(node.getGroupBy())) {
                checkArgument(!node.getSampleWeight().isPresent(), "DISTINCT aggregation has sample weight symbol");
                PlanNode rewrittenSource = context.rewrite(node.getSource());
                return new DistinctLimitNode(.getNextId(), rewrittenSourcelimit.getCount(), Optional.empty());
            }
            PlanNode rewrittenNode = context.defaultRewrite(node);
            if (limit != null) {
                // Drop in a LimitNode b/c limits cannot be pushed through aggregations
                rewrittenNode = new LimitNode(.getNextId(), rewrittenNodelimit.getCount());
            }
            return rewrittenNode;
        }
        @Override
        public PlanNode visitMarkDistinct(MarkDistinctNode nodeRewriteContext<LimitContextcontext)
        {
            // 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(nodecontext.get());
        }
        @Override
        public PlanNode visitProject(ProjectNode nodeRewriteContext<LimitContextcontext)
        {
            // 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(nodecontext.get());
        }
        @Override
        public PlanNode visitTopN(TopNNode nodeRewriteContext<LimitContextcontext)
        {
            LimitContext limit = context.get();
            PlanNode rewrittenSource = context.rewrite(node.getSource());
            if (rewrittenSource == node.getSource() && limit == null) {
                return node;
            }
            long count = node.getCount();
            if (limit != null) {
                count = Math.min(countlimit.getCount());
            }
            return new TopNNode(node.getId(), rewrittenSourcecountnode.getOrderBy(), node.getOrderings(), node.isPartial());
        }
        @Override
        public PlanNode visitSort(SortNode nodeRewriteContext<LimitContextcontext)
        {
            LimitContext limit = context.get();
            PlanNode rewrittenSource = context.rewrite(node.getSource());
            if (limit != null) {
                return new TopNNode(node.getId(), rewrittenSourcelimit.getCount(), node.getOrderBy(), node.getOrderings(), false);
            }
            else if (rewrittenSource != node.getSource()) {
                return new SortNode(node.getId(), rewrittenSourcenode.getOrderBy(), node.getOrderings());
            }
            return node;
        }
        @Override
        public PlanNode visitUnion(UnionNode nodeRewriteContext<LimitContextcontext)
        {
            LimitContext limit = context.get();
            List<PlanNodesources = 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(), sourcesnode.getSymbolMapping());
            if (limit != null) {
                output = new LimitNode(.getNextId(), outputlimit.getCount());
            }
            return output;
        }
        @Override
        public PlanNode visitSemiJoin(SemiJoinNode nodeRewriteContext<LimitContextcontext)
        {
            PlanNode source = context.rewrite(node.getSource(), context.get());
            if (source != node.getSource()) {
                return new SemiJoinNode(node.getId(), sourcenode.getFilteringSource(), node.getSourceJoinSymbol(), node.getFilteringSourceJoinSymbol(), node.getSemiJoinOutput(), node.getSourceHashSymbol(), node.getFilteringSourceHashSymbol());
            }
            return node;
        }
    }
New to GrepCode? Check out our FAQ X