From 5aa3d98ea2f4fe886d47a4ce846406e22850be95 Mon Sep 17 00:00:00 2001
From: Nikolai Kochetov <nik-kochetov@yandex-team.ru>
Date: Fri, 12 Apr 2019 17:43:28 +0300
Subject: [PATCH] Fix limit transfrom prepare.

---
 dbms/src/Processors/LimitTransform.cpp | 35 +++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/dbms/src/Processors/LimitTransform.cpp b/dbms/src/Processors/LimitTransform.cpp
index aeec65ac1ef..0244733f12e 100644
--- a/dbms/src/Processors/LimitTransform.cpp
+++ b/dbms/src/Processors/LimitTransform.cpp
@@ -18,30 +18,35 @@ LimitTransform::LimitTransform(
 
 LimitTransform::Status LimitTransform::prepare()
 {
+    /// Check if we are done with pushing.
+    bool pushing_is_finished = rows_read >= offset + limit;
+
     /// Check can output.
 
     if (output.isFinished())
     {
-        input.close();
-        return Status::Finished;
+        pushing_is_finished = true;
+        if (!always_read_till_end)
+        {
+            input.close();
+            return Status::Finished;
+        }
     }
 
-    if (!output.canPush())
+    if (!pushing_is_finished && !output.canPush())
     {
         input.setNotNeeded();
         return Status::PortFull;
     }
 
     /// Push block if can.
-    if (has_block && block_processed)
+    if (!pushing_is_finished && has_block && block_processed)
     {
         output.push(std::move(current_chunk));
         has_block = false;
         block_processed = false;
     }
 
-    /// Check if we are done with pushing.
-    bool pushing_is_finished = rows_read >= offset + limit;
     if (pushing_is_finished)
     {
         if (!always_read_till_end)
@@ -77,6 +82,12 @@ LimitTransform::Status LimitTransform::prepare()
         current_chunk.clear();
         has_block = false;
 
+        if (input.isFinished())
+        {
+            output.finish();
+            return Status::Finished;
+        }
+
         /// Now, we pulled from input, and it must be empty.
         return Status::NeedData;
     }
@@ -90,6 +101,12 @@ LimitTransform::Status LimitTransform::prepare()
         current_chunk.clear();
         has_block = false;
 
+        if (input.isFinished())
+        {
+            output.finish();
+            return Status::Finished;
+        }
+
         /// Now, we pulled from input, and it must be empty.
         return Status::NeedData;
     }
@@ -103,6 +120,12 @@ LimitTransform::Status LimitTransform::prepare()
         output.push(std::move(current_chunk));
         has_block = false;
 
+        if (input.isFinished())
+        {
+            output.finish();
+            return Status::Finished;
+        }
+
         return Status::NeedData;
     }