From b4852e4c69077b7704ec58a3d8f1bc04076160ba Mon Sep 17 00:00:00 2001
From: Nikolai Kochetov <nik-kochetov@yandex-team.ru>
Date: Thu, 23 Jul 2020 13:20:38 +0300
Subject: [PATCH] Limit push down.

---
 src/Processors/QueryPlan/QueryPlan.cpp | 43 ++++++++++++++++++++++++++
 src/Processors/QueryPlan/QueryPlan.h   |  2 ++
 2 files changed, 45 insertions(+)

diff --git a/src/Processors/QueryPlan/QueryPlan.cpp b/src/Processors/QueryPlan/QueryPlan.cpp
index 2e5d2a3a5f4..3bc8ca1f994 100644
--- a/src/Processors/QueryPlan/QueryPlan.cpp
+++ b/src/Processors/QueryPlan/QueryPlan.cpp
@@ -4,6 +4,7 @@
 #include <IO/WriteBuffer.h>
 #include <IO/Operators.h>
 #include <stack>
+#include <Processors/QueryPlan/LimitStep.h>
 
 namespace DB
 {
@@ -307,4 +308,46 @@ void QueryPlan::explainPipeline(WriteBuffer & buffer, const ExplainPipelineOptio
     }
 }
 
+static void tryPushDownLimit(QueryPlanStepPtr & parent, QueryPlanStepPtr & child)
+{
+    const auto * limit = typeid_cast<const LimitStep *>(parent.get());
+
+    if (!limit)
+        return;
+
+    /// Now we should decide if pushing down limit possible for this step.
+}
+
+void QueryPlan::optimize()
+{
+    struct Frame
+    {
+        Node * node;
+        size_t next_child = 0;
+    };
+
+    std::stack<Frame> stack;
+    stack.push(Frame{.node = root});
+
+    while (!stack.empty())
+    {
+        auto & frame = stack.top();
+
+        if (frame.next_child == 0)
+        {
+            /// First entrance, try push down.
+            if (frame.node->children.size() == 1)
+                tryPushDownLimit(frame.node->step, frame.node->children.front()->step);
+        }
+
+        if (frame.next_child < frame.node->children.size())
+        {
+            stack.push(Frame{frame.node->children[frame.next_child]});
+            ++frame.next_child;
+        }
+        else
+            stack.pop();
+    }
+}
+
 }
diff --git a/src/Processors/QueryPlan/QueryPlan.h b/src/Processors/QueryPlan/QueryPlan.h
index 6f965f7daec..a933688ead7 100644
--- a/src/Processors/QueryPlan/QueryPlan.h
+++ b/src/Processors/QueryPlan/QueryPlan.h
@@ -33,6 +33,8 @@ public:
     bool isCompleted() const; /// Tree is not empty and root hasOutputStream()
     const DataStream & getCurrentDataStream() const; /// Checks that (isInitialized() && !isCompleted())
 
+    void optimize();
+
     QueryPipelinePtr buildQueryPipeline();
 
     struct ExplainPlanOptions