From 2b738fbab79f29ffb3be074b788b98626e6c6683 Mon Sep 17 00:00:00 2001 From: Ti Chi Robot Date: Wed, 6 Dec 2023 23:55:20 +0800 Subject: [PATCH] Avoid out of bound InsertRangeFrom exception in MergingSortedBlockInputStream (#8445) (#8456) 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 87794f5a9c4..b0e082fc488 100644 --- a/dbms/src/DataStreams/MergingSortedBlockInputStream.cpp +++ b/dbms/src/DataStreams/MergingSortedBlockInputStream.cpp @@ -175,7 +175,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; @@ -216,14 +216,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);