From 6a40749a18d814b543609c2fd298a2d04bc4491f Mon Sep 17 00:00:00 2001 From: Ti Chi Robot Date: Fri, 8 Dec 2023 08:46:21 +0800 Subject: [PATCH] Avoid out of bound InsertRangeFrom exception in MergingSortedBlockInputStream (#8445) (#8454) close pingcap/tiflash#8438 --- .../MergingSortedBlockInputStream.cpp | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/dbms/src/DataStreams/MergingSortedBlockInputStream.cpp b/dbms/src/DataStreams/MergingSortedBlockInputStream.cpp index 02545636902..39b88b65891 100644 --- 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,21 @@ 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) { + RUNTIME_CHECK_MSG( + limit >= total_merged_rows, + "Unexpect limit and total_merged_rows {} {}", + limit, + total_merged_rows); 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);