Avoid out of bound InsertRangeFrom exception in MergingSortedBlockInp…

…utStream (#8445) (#8453) close #8438
pingcap · Dec 7, 2023 · 8b29776 · 8b29776
1 parent d48dc34
commit 8b29776
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/dbms/src/DataStreams/MergingSortedBlockInputStream.cpp b/dbms/src/DataStreams/MergingSortedBlockInputStream.cpp
@@ -169,7 +169,7 @@ void MergingSortedBlockInputStream::merge(MutableColumns & merged_columns, std::
       */
     auto count_row_and_check_limit = [&, this]() {
         ++total_merged_rows;
-        if (limit && total_merged_rows == limit)
+        if (limit && total_merged_rows >= limit)
         {
             cancel(false);
             finished = true;
@@ -210,14 +210,18 @@ void MergingSortedBlockInputStream::merge(MutableColumns & merged_columns, std::
                     merged_columns[i] = (*std::move(source_blocks[source_num]->getByPosition(i).column)).mutate();
 
                 size_t merged_rows = merged_columns.at(0)->size();
-
-                if (limit && total_merged_rows + merged_rows > limit)
+                if (limit && total_merged_rows + merged_rows >= limit)
                 {
+                    if (total_merged_rows > limit)
+                        throw Exception("Logical error in MergingSortedBlockInputStream", ErrorCodes::LOGICAL_ERROR);
                     merged_rows = limit - total_merged_rows;
-                    for (size_t i = 0; i < num_columns; ++i)
+                    if likely (total_merged_rows + merged_rows > limit)
                     {
-                        auto & column = merged_columns[i];
-                        column = (*column->cut(0, merged_rows)).mutate();
+                        for (size_t i = 0; i < num_columns; ++i)
+                        {
+                            auto & column = merged_columns[i];
+                            column = (*column->cut(0, merged_rows)).mutate();
+                        }
                     }
 
                     cancel(false);