Skip to content

Commit

Permalink
Avoid out of bound InsertRangeFrom exception in MergingSortedBlockInp…
Browse files Browse the repository at this point in the history
…utStream (#8445) (#8456)

close #8438
  • Loading branch information
ti-chi-bot committed Dec 6, 2023
1 parent e920e3d commit 2b738fb
Showing 1 changed file with 13 additions and 6 deletions.
19 changes: 13 additions & 6 deletions dbms/src/DataStreams/MergingSortedBlockInputStream.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit 2b738fb

Please sign in to comment.