From de8120d69af8404a420a0a7e56759c747323af8f Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Sun, 17 May 2020 10:17:54 +0300
Subject: [PATCH] Improvement; added a test

---
 src/Processors/ISource.cpp                    |  5 --
 src/Processors/Sources/SourceWithProgress.cpp | 14 ++++-
 .../Transforms/LimitsCheckingTransform.cpp    |  2 +-
 .../01132_max_rows_to_read.reference          | 55 ++++++++++++++++++-
 .../0_stateless/01132_max_rows_to_read.sql    | 15 ++++-
 5 files changed, 80 insertions(+), 11 deletions(-)

diff --git a/src/Processors/ISource.cpp b/src/Processors/ISource.cpp
index 7c620a98a74..90f3962b83e 100644
--- a/src/Processors/ISource.cpp
+++ b/src/Processors/ISource.cpp
@@ -56,11 +56,6 @@ void ISource::work()
         finished = true;
         throw;
     }
-//    {
-//        current_chunk = std::current_exception();
-//        has_input = true;
-//        got_exception = true;
-//    }
 }
 
 }
diff --git a/src/Processors/Sources/SourceWithProgress.cpp b/src/Processors/Sources/SourceWithProgress.cpp
index 80844da16cd..8d7a0a3d946 100644
--- a/src/Processors/Sources/SourceWithProgress.cpp
+++ b/src/Processors/Sources/SourceWithProgress.cpp
@@ -15,7 +15,9 @@ namespace ErrorCodes
 void SourceWithProgress::work()
 {
     if (!limits.speed_limits.checkTimeLimit(total_stopwatch.elapsed(), limits.timeout_overflow_mode))
+    {
         cancel();
+    }
     else
     {
         was_progress_called = false;
@@ -57,7 +59,13 @@ void SourceWithProgress::progress(const Progress & value)
         /// The total amount of data processed or intended for processing in all sources, possibly on remote servers.
 
         ProgressValues progress = process_list_elem->getProgressIn();
-        size_t total_rows_estimate = std::max(progress.read_rows, progress.total_rows_to_read);
+
+        /// If the mode is "throw" and estimate of total rows is known, then throw early if an estimate is too high.
+        /// If the mode is "break", then allow to read before limit even if estimate is very high.
+
+        size_t rows_to_check_limit = progress.read_rows;
+        if (limits.size_limits.overflow_mode == OverflowMode::THROW && progress.total_rows_to_read > progress.read_rows)
+            rows_to_check_limit = progress.total_rows_to_read;
 
         /// Check the restrictions on the
         ///  * amount of data to read
@@ -67,9 +75,11 @@ void SourceWithProgress::progress(const Progress & value)
 
         if (limits.mode == LimitsMode::LIMITS_TOTAL)
         {
-            if (!limits.size_limits.check(total_rows_estimate, progress.read_bytes, "rows to read",
+            if (!limits.size_limits.check(rows_to_check_limit, progress.read_bytes, "rows or bytes to read",
                                           ErrorCodes::TOO_MANY_ROWS, ErrorCodes::TOO_MANY_BYTES))
+            {
                 cancel();
+            }
         }
 
         size_t total_rows = progress.total_rows_to_read;
diff --git a/src/Processors/Transforms/LimitsCheckingTransform.cpp b/src/Processors/Transforms/LimitsCheckingTransform.cpp
index c4739ee9dde..56edd5f0317 100644
--- a/src/Processors/Transforms/LimitsCheckingTransform.cpp
+++ b/src/Processors/Transforms/LimitsCheckingTransform.cpp
@@ -58,7 +58,7 @@ void LimitsCheckingTransform::checkQuota(Chunk & chunk)
     switch (limits.mode)
     {
         case LimitsMode::LIMITS_TOTAL:
-            /// Checked in `progress` method.
+            /// Checked in SourceWithProgress::progress method.
             break;
 
         case LimitsMode::LIMITS_CURRENT:
diff --git a/tests/queries/0_stateless/01132_max_rows_to_read.reference b/tests/queries/0_stateless/01132_max_rows_to_read.reference
index 91653a56ef7..5087d15b87c 100644
--- a/tests/queries/0_stateless/01132_max_rows_to_read.reference
+++ b/tests/queries/0_stateless/01132_max_rows_to_read.reference
@@ -1,6 +1,57 @@
 19
 20
-30
 19
 20
-21
+20
+20
+20
+20
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
diff --git a/tests/queries/0_stateless/01132_max_rows_to_read.sql b/tests/queries/0_stateless/01132_max_rows_to_read.sql
index f8e5e32c0b5..b7923a27d04 100644
--- a/tests/queries/0_stateless/01132_max_rows_to_read.sql
+++ b/tests/queries/0_stateless/01132_max_rows_to_read.sql
@@ -7,11 +7,24 @@ SELECT count() FROM numbers(19);
 SELECT count() FROM numbers(20);
 SELECT count() FROM numbers(21); -- { serverError 158 }
 
+-- check early exception if the estimated number of rows is high
+SELECT * FROM numbers(30); -- { serverError 158 }
+
 SET read_overflow_mode = 'break';
 
 SELECT count() FROM numbers(19);
 SELECT count() FROM numbers(20);
 SELECT count() FROM numbers(21);
-SELECT count() FROM numbers(29); -- one extra block is read and it is Ok.
+SELECT count() FROM numbers(29);
 SELECT count() FROM numbers(30);
 SELECT count() FROM numbers(31);
+
+-- check that partial result is returned even if the estimated number of rows is high
+SELECT * FROM numbers(30);
+
+-- the same for uneven block sizes
+-- NOTE: currently it outputs less amount of data; it will be better to output the latest block also
+SET max_block_size = 11;
+SELECT * FROM numbers(30);
+SET max_block_size = 9;
+SELECT * FROM numbers(30);