From f73e312172a4884fe8d9a6ea53b124420e2164fb Mon Sep 17 00:00:00 2001 From: liaoxin Date: Sat, 8 Aug 2026 00:45:46 +0800 Subject: [PATCH 1/2] [improvement](load) store MemTable rows by value instead of shared_ptr MemTable keeps one RowInBlock per loaded row in _row_in_blocks, held as a shared_ptr, so there is one make_shared per row. Measured at 80 bytes per row: 16 for the pointer in the vector, plus a 64 byte heap chunk holding the 16 byte control block and the 40 byte struct. Two of the struct's five fields do not need to be there. _agg_state_offset held _offsets_of_aggregate_states.data(), the same pointer for every row, so it belongs on the MemTable. _has_init_agg said exactly what _agg_mem being non-null already says. That leaves 24 bytes, small enough to keep in the vector directly and drop the per-row allocation with it. For a 7.05M row memtable that is 564 MB down to 169 MB, and building the array drops from 401 ms to 111 ms. Dropping the shared_ptr means the rows in _row_in_blocks and the copies _aggregate() works on are no longer the same object. Two places relied on that aliasing: - prev_row now points into temp_row_in_blocks, which is where _finalize_one_row() will read it from. - _aggregate() adopts temp_row_in_blocks unconditionally, so the entries left behind cannot name an aggregate state that _finalize_one_row() has already released. Without this, a memtable that aggregates across a shrink_memtable_by_agg() round and then again in to_block() destroys those states twice -- confirmed by instrumenting ~MemTable. memtable_sort_test.cpp only covered class Tie. It now also drives a MemTable through insert()/to_block() and covers multi-column and nullable key ordering, the DUP_KEYS tie-break direction, batching independence, UNIQUE_KEYS last-writer-wins, AGG_KEYS aggregation, and aggregate state surviving shrink_memtable_by_agg() rounds. The AGG_KEYS schema carries a BITMAP BITMAP_UNION column so those last cases run over an aggregate state that owns heap memory rather than a trivially destroyed one. --- be/src/load/memtable/memtable.cpp | 294 ++++++++--------- be/src/load/memtable/memtable.h | 62 ++-- be/test/load/memtable/memtable_sort_test.cpp | 313 ++++++++++++++++++- 3 files changed, 494 insertions(+), 175 deletions(-) diff --git a/be/src/load/memtable/memtable.cpp b/be/src/load/memtable/memtable.cpp index 5530cee63aa61c..32baa6f6fdceb3 100644 --- a/be/src/load/memtable/memtable.cpp +++ b/be/src/load/memtable/memtable.cpp @@ -23,6 +23,7 @@ #include #include +#include #include #include @@ -82,7 +83,7 @@ MemTable::MemTable(int64_t tablet_id, std::shared_ptr tablet_schem } _init_columns_offset_by_slot_descs(slot_descs, tuple_desc); // TODO: Support ZOrderComparator in the future - _row_in_blocks = std::make_unique>>(); + _row_in_blocks = std::make_unique>(); _load_mem_limit = MemInfo::mem_limit() * config::load_process_max_memory_limit_percent / 100; } @@ -165,8 +166,8 @@ MemTable::~MemTable() { SCOPED_CONSUME_MEM_TRACKER(_mem_tracker); g_memtable_cnt << -1; if (_keys_type != KeysType::DUP_KEYS) { - for (auto it = _row_in_blocks->begin(); it != _row_in_blocks->end(); it++) { - if (!(*it)->has_init_agg()) { + for (const auto& row : *_row_in_blocks) { + if (!_has_agg(row)) { continue; } // We should release agg_places here, because they are not released when a @@ -174,7 +175,7 @@ MemTable::~MemTable() { for (size_t i = _tablet_schema->num_key_columns(); i < _num_columns; ++i) { auto function = _agg_functions[i]; DCHECK(function != nullptr); - function->destroy((*it)->agg_places(i)); + function->destroy(_agg_place(row, i)); } } } @@ -266,28 +267,29 @@ Status MemTable::insert(const Block* input_block, const TabletAddRowsPayload& ro RETURN_IF_ERROR(_input_mutable_block.add_rows(input_block, row_idxs.data(), row_idxs.data() + num_rows, &_column_offset)); for (int i = 0; i < num_rows; i++) { - _row_in_blocks->emplace_back(std::make_shared( - cursor_in_mutableblock + i, _need_row_binlog_lsn ? row_binlog_lsns[i] : 0)); + _row_in_blocks->emplace_back(cursor_in_mutableblock + i, + _need_row_binlog_lsn ? row_binlog_lsns[i] : 0); } _stat.raw_rows += num_rows; return Status::OK(); } -void MemTable::_merge_row_binlog_lsn(RowInBlock* src_row, RowInBlock* dst_row) { +void MemTable::_merge_row_binlog_lsn(const RowInBlock& src_row, RowInBlock& dst_row) { if (_need_row_binlog_lsn) { - dst_row->_row_binlog_lsn = std::max(dst_row->_row_binlog_lsn, src_row->_row_binlog_lsn); + dst_row._row_binlog_lsn = std::max(dst_row._row_binlog_lsn, src_row._row_binlog_lsn); } } -void MemTable::_append_output_row_binlog_lsn(RowInBlock* row) { +void MemTable::_append_output_row_binlog_lsn(const RowInBlock& row) { if (_need_row_binlog_lsn) { - _output_row_binlog_lsns.emplace_back(row->_row_binlog_lsn); + _output_row_binlog_lsns.emplace_back(row._row_binlog_lsn); } } void MemTable::_aggregate_two_row_with_sequence_map(MutableBlock& mutable_block, - RowInBlock* src_row, RowInBlock* dst_row) { + const RowInBlock& src_row, + RowInBlock& dst_row) { _merge_row_binlog_lsn(src_row, dst_row); // for each mapping replace value columns according to the sequence column compare result // for example: a b c d s1 s2 (key:a , s1=>[b,c], s2=>[d]) @@ -299,7 +301,7 @@ void MemTable::_aggregate_two_row_with_sequence_map(MutableBlock& mutable_block, for (const auto& it : seq_map) { auto sequence = it.first; auto* sequence_col_ptr = mutable_block.mutable_columns()[sequence].get(); - auto res = sequence_col_ptr->compare_at(dst_row->_row_pos, src_row->_row_pos, + auto res = sequence_col_ptr->compare_at(dst_row._row_pos, src_row._row_pos, *sequence_col_ptr, -1); if (res > 0) { continue; @@ -307,27 +309,27 @@ void MemTable::_aggregate_two_row_with_sequence_map(MutableBlock& mutable_block, for (auto cid : it.second) { if (cid < _num_columns) { auto* col_ptr = mutable_block.mutable_columns()[cid].get(); - _agg_functions[cid]->add(dst_row->agg_places(cid), + _agg_functions[cid]->add(_agg_place(dst_row, cid), const_cast(&col_ptr), - src_row->_row_pos, _arena); + src_row._row_pos, _arena); } } if (sequence < _num_columns) { - _agg_functions[sequence]->add(dst_row->agg_places(sequence), + _agg_functions[sequence]->add(_agg_place(dst_row, sequence), const_cast(&sequence_col_ptr), - src_row->_row_pos, _arena); + src_row._row_pos, _arena); // must use replace column instead of update row_pos // because one row may have multi sequence column // and agg function add method won't change the real column value - sequence_col_ptr->replace_column_data(*sequence_col_ptr, src_row->_row_pos, - dst_row->_row_pos); + sequence_col_ptr->replace_column_data(*sequence_col_ptr, src_row._row_pos, + dst_row._row_pos); } } } template -void MemTable::_aggregate_two_row_in_block(MutableBlock& mutable_block, RowInBlock* src_row, - RowInBlock* dst_row) { +void MemTable::_aggregate_two_row_in_block(MutableBlock& mutable_block, const RowInBlock& src_row, + RowInBlock& dst_row) { _merge_row_binlog_lsn(src_row, dst_row); // for flexible partial update, the caller must guarantees that either src_row and dst_row // both specify the sequence column, or src_row and dst_row both don't specify the @@ -335,23 +337,23 @@ void MemTable::_aggregate_two_row_in_block(MutableBlock& mutable_block, RowInBlo if (_tablet_schema->has_sequence_col() && _seq_col_idx_in_block >= 0) { DCHECK_LT(_seq_col_idx_in_block, mutable_block.columns()); auto col_ptr = mutable_block.mutable_columns()[_seq_col_idx_in_block].get(); - auto res = col_ptr->compare_at(dst_row->_row_pos, src_row->_row_pos, *col_ptr, -1); + auto res = col_ptr->compare_at(dst_row._row_pos, src_row._row_pos, *col_ptr, -1); // dst sequence column larger than src, don't need to update if (res > 0) { return; } // need to update the row pos in dst row to the src row pos when has // sequence column - dst_row->_row_pos = src_row->_row_pos; + dst_row._row_pos = src_row._row_pos; } // dst is non-sequence row, or dst sequence is smaller if constexpr (!has_skip_bitmap_col) { DCHECK(_skip_bitmap_col_idx == -1); for (size_t cid = _tablet_schema->num_key_columns(); cid < _num_columns; ++cid) { auto* col_ptr = mutable_block.mutable_columns()[cid].get(); - _agg_functions[cid]->add(dst_row->agg_places(cid), - const_cast(&col_ptr), - src_row->_row_pos, _arena); + _agg_functions[cid]->add(_agg_place(dst_row, cid), + const_cast(&col_ptr), src_row._row_pos, + _arena); } } else { DCHECK(_skip_bitmap_col_idx != -1); @@ -359,16 +361,16 @@ void MemTable::_aggregate_two_row_in_block(MutableBlock& mutable_block, RowInBlo const BitmapValue& skip_bitmap = assert_cast( mutable_block.mutable_columns()[_skip_bitmap_col_idx].get()) - ->get_data()[src_row->_row_pos]; + ->get_data()[src_row._row_pos]; for (size_t cid = _tablet_schema->num_key_columns(); cid < _num_columns; ++cid) { const auto& col = _tablet_schema->column(cid); if (cid != _skip_bitmap_col_idx && skip_bitmap.contains(col.unique_id())) { continue; } auto* col_ptr = mutable_block.mutable_columns()[cid].get(); - _agg_functions[cid]->add(dst_row->agg_places(cid), - const_cast(&col_ptr), - src_row->_row_pos, _arena); + _agg_functions[cid]->add(_agg_place(dst_row, cid), + const_cast(&col_ptr), src_row._row_pos, + _arena); } } } @@ -380,14 +382,30 @@ Status MemTable::_put_into_output(Block& in_block) { if (_need_row_binlog_lsn) { _output_row_binlog_lsns.reserve(_output_row_binlog_lsns.size() + in_block.rows()); } - for (int i = 0; i < _row_in_blocks->size(); i++) { - row_pos_vec.emplace_back((*_row_in_blocks)[i]->_row_pos); - _append_output_row_binlog_lsn((*_row_in_blocks)[i].get()); + for (const auto& row : *_row_in_blocks) { + row_pos_vec.emplace_back(row._row_pos); + _append_output_row_binlog_lsn(row); } return _output_mutable_block.add_rows(&in_block, row_pos_vec.data(), row_pos_vec.data() + in_block.rows()); } +void MemTable::_sort_one_column(DorisVector& row_in_blocks, Tie& tie, + std::function cmp) { + auto iter = tie.iter(); + while (iter.next()) { + pdqsort(std::next(row_in_blocks.begin(), static_cast(iter.left())), + std::next(row_in_blocks.begin(), static_cast(iter.right())), + [&cmp](const RowInBlock& lhs, const RowInBlock& rhs) -> bool { + return cmp(lhs, rhs) < 0; + }); + tie[iter.left()] = 0; + for (auto i = iter.left() + 1; i < iter.right(); i++) { + tie[i] = (cmp(row_in_blocks[i - 1], row_in_blocks[i]) == 0); + } + } +} + size_t MemTable::_sort() { SCOPED_RAW_TIMER(&_stat.sort_ns); _stat.sort_times++; @@ -395,8 +413,8 @@ size_t MemTable::_sort() { // sort new rows Tie tie = Tie(_last_sorted_pos, _row_in_blocks->size()); for (size_t i = 0; i < _tablet_schema->num_key_columns(); i++) { - auto cmp = [&](RowInBlock* lhs, RowInBlock* rhs) -> int { - return _input_mutable_block.compare_one_column(lhs->_row_pos, rhs->_row_pos, i, -1); + auto cmp = [&](const RowInBlock& lhs, const RowInBlock& rhs) -> int { + return _input_mutable_block.compare_one_column(lhs._row_pos, rhs._row_pos, i, -1); }; _sort_one_column(*_row_in_blocks, tie, cmp); } @@ -406,20 +424,19 @@ size_t MemTable::_sort() { while (iter.next()) { pdqsort(std::next(_row_in_blocks->begin(), iter.left()), std::next(_row_in_blocks->begin(), iter.right()), - [&is_dup](const std::shared_ptr& lhs, - const std::shared_ptr& rhs) -> bool { - return is_dup ? lhs->_row_pos > rhs->_row_pos : lhs->_row_pos < rhs->_row_pos; + [&is_dup](const RowInBlock& lhs, const RowInBlock& rhs) -> bool { + return is_dup ? lhs._row_pos > rhs._row_pos : lhs._row_pos < rhs._row_pos; }); same_keys_num += iter.right() - iter.left(); } // merge new rows and old rows _vec_row_comparator->set_block(&_input_mutable_block); - auto cmp_func = [this, is_dup, &same_keys_num](const std::shared_ptr& l, - const std::shared_ptr& r) -> bool { - auto value = (*(this->_vec_row_comparator))(l.get(), r.get()); + auto cmp_func = [this, is_dup, &same_keys_num](const RowInBlock& l, + const RowInBlock& r) -> bool { + auto value = (*(this->_vec_row_comparator))(&l, &r); if (value == 0) { same_keys_num++; - return is_dup ? l->_row_pos > r->_row_pos : l->_row_pos < r->_row_pos; + return is_dup ? l._row_pos > r._row_pos : l._row_pos < r._row_pos; } else { return value < 0; } @@ -439,15 +456,13 @@ Status MemTable::_sort_by_cluster_keys() { MutableBlock mutable_block = MutableBlock::build_mutable_block(std::move(in_block)); _output_mutable_block = MutableBlock::build_mutable_block(std::move(clone_block)); - DorisVector> row_in_blocks; + DorisVector row_in_blocks; row_in_blocks.reserve(mutable_block.rows()); if (_need_row_binlog_lsn) { DCHECK_EQ(_output_row_binlog_lsns.size(), mutable_block.rows()); } for (size_t i = 0; i < mutable_block.rows(); i++) { - row_in_blocks.emplace_back( - _need_row_binlog_lsn ? std::make_shared(i, _output_row_binlog_lsns[i]) - : std::make_shared(i)); + row_in_blocks.emplace_back(i, _need_row_binlog_lsn ? _output_row_binlog_lsns[i] : 0); } if (_need_row_binlog_lsn) { _output_row_binlog_lsns.clear(); @@ -461,8 +476,8 @@ Status MemTable::_sort_by_cluster_keys() { return Status::InternalError("could not find cluster key column with unique_id=" + std::to_string(cid) + " in tablet schema"); } - auto cmp = [&](const RowInBlock* lhs, const RowInBlock* rhs) -> int { - return mutable_block.compare_one_column(lhs->_row_pos, rhs->_row_pos, index, -1); + auto cmp = [&](const RowInBlock& lhs, const RowInBlock& rhs) -> int { + return mutable_block.compare_one_column(lhs._row_pos, rhs._row_pos, index, -1); }; _sort_one_column(row_in_blocks, tie, cmp); } @@ -472,8 +487,9 @@ Status MemTable::_sort_by_cluster_keys() { while (iter.next()) { pdqsort(std::next(row_in_blocks.begin(), iter.left()), std::next(row_in_blocks.begin(), iter.right()), - [](const std::shared_ptr& lhs, const std::shared_ptr& rhs) - -> bool { return lhs->_row_pos < rhs->_row_pos; }); + [](const RowInBlock& lhs, const RowInBlock& rhs) -> bool { + return lhs._row_pos < rhs._row_pos; + }); } in_block = mutable_block.to_block(); @@ -481,9 +497,9 @@ Status MemTable::_sort_by_cluster_keys() { DorisVector row_pos_vec; DCHECK(in_block.rows() <= std::numeric_limits::max()); row_pos_vec.reserve(in_block.rows()); - for (int i = 0; i < row_in_blocks.size(); i++) { - row_pos_vec.emplace_back(row_in_blocks[i]->_row_pos); - _append_output_row_binlog_lsn(row_in_blocks[i].get()); + for (const auto& row : row_in_blocks) { + row_pos_vec.emplace_back(row._row_pos); + _append_output_row_binlog_lsn(row); } std::vector column_offset; for (int i = 0; i < _column_offset.size(); ++i) { @@ -493,32 +509,18 @@ Status MemTable::_sort_by_cluster_keys() { row_pos_vec.data() + in_block.rows(), &column_offset); } -void MemTable::_sort_one_column(DorisVector>& row_in_blocks, Tie& tie, - std::function cmp) { - auto iter = tie.iter(); - while (iter.next()) { - pdqsort(std::next(row_in_blocks.begin(), static_cast(iter.left())), - std::next(row_in_blocks.begin(), static_cast(iter.right())), - [&cmp](auto lhs, auto rhs) -> bool { return cmp(lhs.get(), rhs.get()) < 0; }); - tie[iter.left()] = 0; - for (auto i = iter.left() + 1; i < iter.right(); i++) { - tie[i] = (cmp(row_in_blocks[i - 1].get(), row_in_blocks[i].get()) == 0); - } - } -} - template -void MemTable::_finalize_one_row(RowInBlock* row, MutableBlock& mutable_block, int row_pos) { +void MemTable::_finalize_one_row(RowInBlock& row, MutableBlock& mutable_block, int row_pos) { // move key columns for (size_t i = 0; i < _tablet_schema->num_key_columns(); ++i) { _output_mutable_block.get_column_by_position(i)->insert_from( - *mutable_block.get_column_by_position(i), row->_row_pos); + *mutable_block.get_column_by_position(i), row._row_pos); } - if (row->has_init_agg()) { + if (_has_agg(row)) { // get value columns from agg_places for (size_t i = _tablet_schema->num_key_columns(); i < _num_columns; ++i) { auto function = _agg_functions[i]; - auto* agg_place = row->agg_places(i); + auto* agg_place = _agg_place(row, i); auto* col_ptr = _output_mutable_block.get_column_by_position(i).get(); function->insert_result_into(agg_place, *col_ptr); @@ -530,11 +532,11 @@ void MemTable::_finalize_one_row(RowInBlock* row, MutableBlock& mutable_block, i } if constexpr (is_final) { - row->remove_init_agg(); + row._agg_mem = nullptr; } else { for (size_t i = _tablet_schema->num_key_columns(); i < _num_columns; ++i) { auto function = _agg_functions[i]; - auto* agg_place = row->agg_places(i); + auto* agg_place = _agg_place(row, i); auto* col_ptr = _output_mutable_block.get_column_by_position(i).get(); function->add(agg_place, const_cast(&col_ptr), row_pos, _arena); @@ -544,34 +546,33 @@ void MemTable::_finalize_one_row(RowInBlock* row, MutableBlock& mutable_block, i // move columns for rows do not need agg for (size_t i = _tablet_schema->num_key_columns(); i < _num_columns; ++i) { _output_mutable_block.get_column_by_position(i)->insert_from( - *mutable_block.get_column_by_position(i), row->_row_pos); + *mutable_block.get_column_by_position(i), row._row_pos); } } _append_output_row_binlog_lsn(row); if constexpr (!is_final) { - row->_row_pos = row_pos; + row._row_pos = row_pos; } } -void MemTable::_init_row_for_agg(RowInBlock* row, MutableBlock& mutable_block) { - row->init_agg_places(_arena.aligned_alloc(_total_size_of_aggregate_states, 16), - _offsets_of_aggregate_states.data()); +void MemTable::_init_row_for_agg(RowInBlock& row, MutableBlock& mutable_block) { + row._agg_mem = _arena.aligned_alloc(_total_size_of_aggregate_states, 16); for (auto cid = _tablet_schema->num_key_columns(); cid < _num_columns; cid++) { auto* col_ptr = mutable_block.mutable_columns()[cid].get(); - auto* data = row->agg_places(cid); + auto* data = _agg_place(row, cid); _agg_functions[cid]->create(data); - _agg_functions[cid]->add(data, const_cast(&col_ptr), row->_row_pos, + _agg_functions[cid]->add(data, const_cast(&col_ptr), row._row_pos, _arena); } } -void MemTable::_clear_row_agg(RowInBlock* row) { - if (row->has_init_agg()) { +void MemTable::_clear_row_agg(RowInBlock& row) { + if (_has_agg(row)) { for (size_t i = _tablet_schema->num_key_columns(); i < _num_columns; ++i) { auto function = _agg_functions[i]; - auto* agg_place = row->agg_places(i); + auto* agg_place = _agg_place(row, i); function->destroy(agg_place); } - row->remove_init_agg(); + row._agg_mem = nullptr; } } // only in `to_block` the `is_final` flag will be true, in other cases, it will be false @@ -583,46 +584,48 @@ void MemTable::_aggregate() { std::unique_ptr empty_input_block = in_block.create_same_struct_block(0); MutableBlock mutable_block = MutableBlock::build_mutable_block(std::move(in_block)); _vec_row_comparator->set_block(&mutable_block); - DorisVector> temp_row_in_blocks; - temp_row_in_blocks.reserve(_last_sorted_pos); + DorisVector temp_row_in_blocks; + // Rows are held by value, so prev_row below points into temp_row_in_blocks. + // Reserving the upper bound up front keeps that pointer valid for the whole + // loop. + temp_row_in_blocks.reserve(_row_in_blocks->size()); //only init agg if needed if constexpr (!has_skip_bitmap_col) { RowInBlock* prev_row = nullptr; int row_pos = -1; - for (const auto& cur_row_ptr : *_row_in_blocks) { - RowInBlock* cur_row = cur_row_ptr.get(); - if (!temp_row_in_blocks.empty() && (*_vec_row_comparator)(prev_row, cur_row) == 0) { - if (!prev_row->has_init_agg()) { - _init_row_for_agg(prev_row, mutable_block); + for (RowInBlock& cur_row : *_row_in_blocks) { + if (!temp_row_in_blocks.empty() && (*_vec_row_comparator)(prev_row, &cur_row) == 0) { + if (!_has_agg(*prev_row)) { + _init_row_for_agg(*prev_row, mutable_block); } _stat.merged_rows++; if (_tablet_schema->has_seq_map()) { - _aggregate_two_row_with_sequence_map(mutable_block, cur_row, prev_row); + _aggregate_two_row_with_sequence_map(mutable_block, cur_row, *prev_row); } else { _aggregate_two_row_in_block(mutable_block, cur_row, - prev_row); + *prev_row); } // Clean up aggregation state of the merged row to avoid memory leak - if (cur_row) { - _clear_row_agg(cur_row); - } + _clear_row_agg(cur_row); } else { - prev_row = cur_row; if (!temp_row_in_blocks.empty()) { // The rows from the previous batch of _row_in_blocks have been merged into temp_row_in_blocks, // now call finalize to write the aggregation results into _output_mutable_block. - _finalize_one_row(temp_row_in_blocks.back().get(), mutable_block, - row_pos); + _finalize_one_row(temp_row_in_blocks.back(), mutable_block, row_pos); } - temp_row_in_blocks.push_back(cur_row_ptr); + // Aggregation mutates the group representative, and the copy that + // _finalize_one_row will read is the one in temp_row_in_blocks, so + // prev_row has to point there rather than into _row_in_blocks. + temp_row_in_blocks.push_back(cur_row); + prev_row = &temp_row_in_blocks.back(); row_pos++; } } if (!temp_row_in_blocks.empty()) { // finalize the last low - _finalize_one_row(temp_row_in_blocks.back().get(), mutable_block, row_pos); + _finalize_one_row(temp_row_in_blocks.back(), mutable_block, row_pos); } } else { DCHECK(_delete_sign_col_idx != -1); @@ -641,15 +644,18 @@ void MemTable::_aggregate() { _output_mutable_block = MutableBlock::build_mutable_block(std::move(*empty_input_block)); _output_mutable_block.clear_column_data(); _output_row_binlog_lsns.clear(); - *_row_in_blocks = temp_row_in_blocks; - _last_sorted_pos = _row_in_blocks->size(); } + // Rows are held by value, so the entries left in _row_in_blocks are stale + // copies of the ones _finalize_one_row just worked on -- including their + // _agg_mem, whose state it may have released. Adopting the finalized rows + // unconditionally keeps ~MemTable from destroying a state a second time. + *_row_in_blocks = std::move(temp_row_in_blocks); + _last_sorted_pos = _row_in_blocks->size(); } template void MemTable::_aggregate_for_flexible_partial_update_without_seq_col( - MutableBlock& mutable_block, DorisVector>& temp_row_in_blocks) { - std::shared_ptr prev_row {nullptr}; + MutableBlock& mutable_block, DorisVector& temp_row_in_blocks) { int row_pos = -1; auto& skip_bitmaps = assert_cast(mutable_block.mutable_columns()[_skip_bitmap_col_idx].get()) @@ -657,66 +663,73 @@ void MemTable::_aggregate_for_flexible_partial_update_without_seq_col( auto& delete_signs = assert_cast(mutable_block.mutable_columns()[_delete_sign_col_idx].get()) ->get_data(); - std::shared_ptr row_with_delete_sign {nullptr}; - std::shared_ptr row_without_delete_sign {nullptr}; + // Rows are held by value here: a held row is only appended to + // temp_row_in_blocks once its whole key group has been consumed, so it is + // aggregated into while it lives in one of these two slots. + std::optional row_with_delete_sign; + std::optional row_without_delete_sign; auto finalize_rows = [&]() { - if (row_with_delete_sign != nullptr) { - temp_row_in_blocks.push_back(row_with_delete_sign); - _finalize_one_row(row_with_delete_sign.get(), mutable_block, ++row_pos); - row_with_delete_sign = nullptr; + if (row_with_delete_sign.has_value()) { + temp_row_in_blocks.push_back(*row_with_delete_sign); + _finalize_one_row(temp_row_in_blocks.back(), mutable_block, ++row_pos); + row_with_delete_sign.reset(); } - if (row_without_delete_sign != nullptr) { - temp_row_in_blocks.push_back(row_without_delete_sign); - _finalize_one_row(row_without_delete_sign.get(), mutable_block, ++row_pos); - row_without_delete_sign = nullptr; + if (row_without_delete_sign.has_value()) { + temp_row_in_blocks.push_back(*row_without_delete_sign); + _finalize_one_row(temp_row_in_blocks.back(), mutable_block, ++row_pos); + row_without_delete_sign.reset(); } // _arena.clear(); }; - auto add_row = [&](std::shared_ptr row, bool with_delete_sign) { + auto add_row = [&](const RowInBlock& row, bool with_delete_sign) { if (with_delete_sign) { - row_with_delete_sign = std::move(row); + row_with_delete_sign = row; } else { - row_without_delete_sign = std::move(row); + row_without_delete_sign = row; } }; - for (const auto& cur_row_ptr : *_row_in_blocks) { - RowInBlock* cur_row = cur_row_ptr.get(); - const BitmapValue& skip_bitmap = skip_bitmaps[cur_row->_row_pos]; + for (RowInBlock& cur_row : *_row_in_blocks) { + const BitmapValue& skip_bitmap = skip_bitmaps[cur_row._row_pos]; bool cur_row_has_delete_sign = (!skip_bitmap.contains(_delete_sign_col_unique_id) && - delete_signs[cur_row->_row_pos] != 0); - prev_row = - (row_with_delete_sign == nullptr) ? row_without_delete_sign : row_with_delete_sign; + delete_signs[cur_row._row_pos] != 0); // compare keys, the keys of row_with_delete_sign and row_without_delete_sign is the same, // choose any of them if it's valid - if (prev_row != nullptr && (*_vec_row_comparator)(prev_row.get(), cur_row) == 0) { + RowInBlock* prev_row = + row_with_delete_sign.has_value() + ? &row_with_delete_sign.value() + : (row_without_delete_sign.has_value() ? &row_without_delete_sign.value() + : nullptr); + if (prev_row != nullptr && (*_vec_row_comparator)(prev_row, &cur_row) == 0) { if (cur_row_has_delete_sign) { - if (row_without_delete_sign != nullptr) { + if (row_without_delete_sign.has_value()) { // if there exits row without delete sign, remove it first - _merge_row_binlog_lsn(row_without_delete_sign.get(), cur_row); - _clear_row_agg(row_without_delete_sign.get()); + _merge_row_binlog_lsn(*row_without_delete_sign, cur_row); + _clear_row_agg(*row_without_delete_sign); _stat.merged_rows++; - row_without_delete_sign = nullptr; + row_without_delete_sign.reset(); } // and then unconditionally replace the previous row - prev_row = row_with_delete_sign; + prev_row = + row_with_delete_sign.has_value() ? &row_with_delete_sign.value() : nullptr; } else { - prev_row = row_without_delete_sign; + prev_row = row_without_delete_sign.has_value() ? &row_without_delete_sign.value() + : nullptr; } if (prev_row == nullptr) { - add_row(cur_row_ptr, cur_row_has_delete_sign); + add_row(cur_row, cur_row_has_delete_sign); } else { - if (!prev_row->has_init_agg()) { - _init_row_for_agg(prev_row.get(), mutable_block); + if (!_has_agg(*prev_row)) { + _init_row_for_agg(*prev_row, mutable_block); } _stat.merged_rows++; - _aggregate_two_row_in_block(mutable_block, cur_row, prev_row.get()); + _aggregate_two_row_in_block(mutable_block, cur_row, *prev_row); } } else { finalize_rows(); - add_row(cur_row_ptr, cur_row_has_delete_sign); + add_row(cur_row, cur_row_has_delete_sign); } } // finalize the last lows @@ -725,14 +738,13 @@ void MemTable::_aggregate_for_flexible_partial_update_without_seq_col( template void MemTable::_aggregate_for_flexible_partial_update_with_seq_col( - MutableBlock& mutable_block, DorisVector>& temp_row_in_blocks) { + MutableBlock& mutable_block, DorisVector& temp_row_in_blocks) { // For flexible partial update, when table has sequence column, we don't do any aggregation // in memtable. These duplicate rows will be aggregated in VerticalSegmentWriter int row_pos = -1; - for (const auto& row_ptr : *_row_in_blocks) { - RowInBlock* row = row_ptr.get(); - temp_row_in_blocks.push_back(row_ptr); - _finalize_one_row(row, mutable_block, ++row_pos); + for (const RowInBlock& row : *_row_in_blocks) { + temp_row_in_blocks.push_back(row); + _finalize_one_row(temp_row_in_blocks.back(), mutable_block, ++row_pos); } } @@ -813,7 +825,7 @@ Status MemTable::_to_block(std::unique_ptr* res) { if (_need_row_binlog_lsn) { _output_row_binlog_lsns.reserve(_row_in_blocks->size()); for (const auto& row : *_row_in_blocks) { - _append_output_row_binlog_lsn(row.get()); + _append_output_row_binlog_lsn(row); } } } else { diff --git a/be/src/load/memtable/memtable.h b/be/src/load/memtable/memtable.h index 431608966599d4..1525359f07761a 100644 --- a/be/src/load/memtable/memtable.h +++ b/be/src/load/memtable/memtable.h @@ -50,29 +50,21 @@ enum KeysType : int; // FLUSH: the memtable is under flushing, write segment to disk. enum MemType { ACTIVE = 0, WRITE_FINISHED = 1, FLUSH = 2 }; -// row pos in _input_mutable_block +// A row of _input_mutable_block, kept by value in MemTable::_row_in_blocks. +// Small and trivially copyable on purpose: there is one of these per loaded row, +// so anything stored here is multiplied by the memtable's row count. struct RowInBlock { size_t _row_pos; int64_t _row_binlog_lsn = 0; + // Aggregate state of this row, allocated from MemTable::_arena; null means + // the row has not been aggregated into yet. The offsets of the individual + // states are the same for every row, so they live on the MemTable rather + // than being repeated here. char* _agg_mem = nullptr; - size_t* _agg_state_offset = nullptr; - bool _has_init_agg; - RowInBlock(size_t row) : _row_pos(row), _has_init_agg(false) {} + RowInBlock(size_t row) : _row_pos(row) {} RowInBlock(size_t row, int64_t row_binlog_lsn) - : _row_pos(row), _row_binlog_lsn(row_binlog_lsn), _has_init_agg(false) {} - - void init_agg_places(char* agg_mem, size_t* agg_state_offset) { - _has_init_agg = true; - _agg_mem = agg_mem; - _agg_state_offset = agg_state_offset; - } - - char* agg_places(size_t offset) const { return _agg_mem + _agg_state_offset[offset]; } - - inline bool has_init_agg() const { return _has_init_agg; } - - inline void remove_init_agg() { _has_init_agg = false; } + : _row_pos(row), _row_binlog_lsn(row_binlog_lsn) {} }; class Tie { @@ -217,18 +209,19 @@ class MemTable { private: // for vectorized template - void _aggregate_two_row_in_block(MutableBlock& mutable_block, RowInBlock* new_row, - RowInBlock* row_in_skiplist); + void _aggregate_two_row_in_block(MutableBlock& mutable_block, const RowInBlock& new_row, + RowInBlock& row_in_skiplist); // Merge row-binlog LSN sidecar only when MemTable merges two RowInBlock objects. // Table models that require complex merge semantics, such as AGG tables and unique key // merge-on-read tables, do not support row-binlog LSN now and are rejected in insert(). - void _merge_row_binlog_lsn(RowInBlock* src_row, RowInBlock* dst_row); + void _merge_row_binlog_lsn(const RowInBlock& src_row, RowInBlock& dst_row); - void _append_output_row_binlog_lsn(RowInBlock* row); + void _append_output_row_binlog_lsn(const RowInBlock& row); - void _aggregate_two_row_with_sequence_map(MutableBlock& mutable_block, RowInBlock* new_row, - RowInBlock* row_in_skiplist); + void _aggregate_two_row_with_sequence_map(MutableBlock& mutable_block, + const RowInBlock& new_row, + RowInBlock& row_in_skiplist); // Used to wrapped by to_block to do exception handle logic Status _to_block(std::unique_ptr* res); @@ -275,25 +268,28 @@ class MemTable { //return number of same keys size_t _sort(); Status _sort_by_cluster_keys(); - void _sort_one_column(DorisVector>& row_in_blocks, Tie& tie, - std::function cmp); + void _sort_one_column(DorisVector& row_in_blocks, Tie& tie, + std::function cmp); template - void _finalize_one_row(RowInBlock* row, MutableBlock& mutable_block, int row_pos); - void _init_row_for_agg(RowInBlock* row, MutableBlock& mutable_block); - void _clear_row_agg(RowInBlock* row); + void _finalize_one_row(RowInBlock& row, MutableBlock& mutable_block, int row_pos); + void _init_row_for_agg(RowInBlock& row, MutableBlock& mutable_block); + void _clear_row_agg(RowInBlock& row); + + static bool _has_agg(const RowInBlock& row) { return row._agg_mem != nullptr; } + char* _agg_place(const RowInBlock& row, size_t cid) const { + return row._agg_mem + _offsets_of_aggregate_states[cid]; + } template void _aggregate(); template void _aggregate_for_flexible_partial_update_without_seq_col( - MutableBlock& mutable_block, - DorisVector>& temp_row_in_blocks); + MutableBlock& mutable_block, DorisVector& temp_row_in_blocks); template void _aggregate_for_flexible_partial_update_with_seq_col( - MutableBlock& mutable_block, - DorisVector>& temp_row_in_blocks); + MutableBlock& mutable_block, DorisVector& temp_row_in_blocks); Status _put_into_output(Block& in_block); bool _is_first_insertion; @@ -302,7 +298,7 @@ class MemTable { std::vector _agg_functions; std::vector _offsets_of_aggregate_states; size_t _total_size_of_aggregate_states; - std::unique_ptr>> _row_in_blocks; + std::unique_ptr> _row_in_blocks; size_t _num_columns; int32_t _seq_col_idx_in_block {-1}; diff --git a/be/test/load/memtable/memtable_sort_test.cpp b/be/test/load/memtable/memtable_sort_test.cpp index 53e92e3c4bb2be..13cb77102ad1a6 100644 --- a/be/test/load/memtable/memtable_sort_test.cpp +++ b/be/test/load/memtable/memtable_sort_test.cpp @@ -17,11 +17,200 @@ #include +#include +#include +#include +#include + +#include "core/block/block.h" +#include "core/column/column_complex.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_vector.h" +#include "load/delta_writer/delta_writer_context.h" #include "load/memtable/memtable.h" +#include "runtime/descriptor_helper.h" +#include "runtime/descriptors.h" +#include "runtime/memory/mem_tracker_limiter.h" +#include "runtime/workload_management/resource_context.h" +#include "storage/tablet/tablet_schema.h" namespace doris { -class MemTableSortTest : public ::testing::Test {}; +namespace { + +// Schema used by every case: k1 INT (key), k2 VARCHAR (key, nullable), v INT. +// Two key columns are needed so the equal-range refinement between key columns +// is exercised, and k2 is nullable so the ColumnNullable sort path is covered. +// AGG_KEYS additionally gets bm BITMAP BITMAP_UNION. Its aggregate state owns +// heap memory, unlike SUM over an int, so releasing a state twice is an actual +// double free there and the shrink-round cases below can catch it. +bool has_bitmap_col(KeysType keys_type) { + return keys_type == KeysType::AGG_KEYS; +} + +TabletSchemaSPtr create_schema(KeysType keys_type) { + TabletSchemaPB pb; + pb.set_keys_type(keys_type); + + auto add = [&](const std::string& name, const std::string& type, bool is_key, bool nullable, + int32_t length, const std::string& agg) { + ColumnPB* c = pb.add_column(); + c->set_unique_id(pb.column_size()); + c->set_name(name); + c->set_type(type); + c->set_is_key(is_key); + c->set_is_nullable(nullable); + c->set_length(length); + c->set_aggregation(agg); + c->set_is_bf_column(false); + }; + add("k1", "INT", true, false, 4, "NONE"); + add("k2", "VARCHAR", true, true, 20, "NONE"); + // value aggregation only matters for AGG_KEYS; UNIQUE_KEYS always replaces + add("v", "INT", false, false, 4, keys_type == KeysType::AGG_KEYS ? "SUM" : "REPLACE"); + if (has_bitmap_col(keys_type)) { + add("bm", "BITMAP", false, false, 16, "BITMAP_UNION"); + } + + auto schema = std::make_shared(); + schema->init_from_pb(pb); + return schema; +} + +TDescriptorTable create_descriptor_table(KeysType keys_type) { + TDescriptorTableBuilder dtb; + TTupleDescriptorBuilder tuple_builder; + tuple_builder.add_slot(TSlotDescriptorBuilder() + .type(TYPE_INT) + .nullable(false) + .column_name("k1") + .column_pos(0) + .build()); + tuple_builder.add_slot(TSlotDescriptorBuilder() + .string_type(20) + .nullable(true) + .column_name("k2") + .column_pos(1) + .build()); + tuple_builder.add_slot(TSlotDescriptorBuilder() + .type(TYPE_INT) + .nullable(false) + .column_name("v") + .column_pos(2) + .build()); + if (has_bitmap_col(keys_type)) { + tuple_builder.add_slot(TSlotDescriptorBuilder() + .type(TYPE_BITMAP) + .nullable(false) + .column_name("bm") + .column_pos(3) + .build()); + } + tuple_builder.build(&dtb); + return dtb.desc_tbl(); +} + +struct Row { + int32_t k1; + const char* k2; // nullptr means SQL NULL + int32_t v; +}; + +} // namespace + +class MemTableSortTest : public testing::Test { +protected: + // Feeds `rows` through a MemTable in `batches` insert() calls and hands back + // the flushed block in `out`. Going through to_block() means the real _sort() + // runs. Fatal assertions abort this helper, so callers wrap it in + // ASSERT_NO_FATAL_FAILURE rather than reading a half-built block. + void run(KeysType keys_type, const std::vector& rows, size_t batches, + std::unique_ptr* out, bool shrink_between_batches = false) { + TabletSchemaSPtr schema = create_schema(keys_type); + TDescriptorTable tdesc = create_descriptor_table(keys_type); + ObjectPool pool; + DescriptorTbl* desc_tbl = nullptr; + ASSERT_TRUE(DescriptorTbl::create(&pool, tdesc, &desc_tbl).ok()); + TupleDescriptor* tuple_desc = desc_tbl->get_tuple_descriptor(0); + ASSERT_NE(nullptr, tuple_desc); + auto resource_ctx = ResourceContext::create_shared(); + // MemTable dereferences this tracker in its constructor, and a freshly + // created context has none; production installs one the same way. + resource_ctx->memory_context()->set_mem_tracker(MemTrackerLimiter::create_shared( + MemTrackerLimiter::Type::LOAD, "MemTableSortTest")); + + MemTable mem_table(10000, schema, &tuple_desc->slots(), tuple_desc, + false /*enable_unique_key_mow*/, nullptr /*partial_update_info*/, + resource_ctx, false /*need_row_binlog_lsn*/); + + const size_t per_batch = (rows.size() + batches - 1) / batches; + for (size_t begin = 0; begin < rows.size(); begin += per_batch) { + const size_t end = std::min(begin + per_batch, rows.size()); + Block block; + for (const auto* slot : tuple_desc->slots()) { + block.insert(ColumnWithTypeAndName(slot->get_empty_mutable_column(), slot->type(), + slot->col_name())); + } + auto columns = std::move(block).mutate_columns(); + for (size_t i = begin; i < end; ++i) { + columns[0]->insert_data(reinterpret_cast(&rows[i].k1), + sizeof(rows[i].k1)); + if (rows[i].k2 == nullptr) { + columns[1]->insert_default(); + } else { + columns[1]->insert_data(rows[i].k2, strlen(rows[i].k2)); + } + columns[2]->insert_data(reinterpret_cast(&rows[i].v), + sizeof(rows[i].v)); + if (has_bitmap_col(keys_type)) { + // Enough values to push BitmapValue past its inline SINGLE + // representation into a heap-backed roaring bitmap, so that + // releasing the aggregate state twice is a real double free. + BitmapValue bitmap; + for (uint64_t b = 0; b < 128; ++b) { + bitmap.add(static_cast(rows[i].v) * 100000 + b * 977); + } + assert_cast(columns[3].get())->insert_value(std::move(bitmap)); + } + } + block.set_columns(std::move(columns)); + + TabletAddRowsPayload payload; + for (uint32_t i = 0; i < end - begin; ++i) { + payload.row_idxs.push_back(i); + } + Status st = mem_table.insert(&block, payload); + ASSERT_TRUE(st.ok()) << st; + if (shrink_between_batches && end < rows.size()) { + // Runs a non-final aggregate: rows that survive keep their + // aggregate state and get aggregated into again next round. + // Deliberately skipped for the last batch, so the duplicates it + // introduces are still there for the final aggregate in + // to_block() to fold into those surviving rows. + mem_table.shrink_memtable_by_agg(); + } + } + + Status st = mem_table.to_block(out); + ASSERT_TRUE(st.ok()) << st; + ASSERT_NE(nullptr, *out); + } + + static std::string k2_of(const Block& b, size_t row) { + StringRef ref = b.get_by_position(1).column->get_data_at(row); + return ref.data == nullptr ? std::string("") : ref.to_string(); + } + static int32_t int_of(const Block& b, size_t pos, size_t row) { + ColumnPtr col = b.get_by_position(pos).column; + if (const auto* nullable = check_and_get_column(col.get())) { + col = nullable->get_nested_column_ptr(); + } + return static_cast(col->get_int(row)); + } + static int32_t k1_of(const Block& b, size_t row) { return int_of(b, 0, row); } + static int32_t v_of(const Block& b, size_t row) { return int_of(b, 2, row); } +}; TEST_F(MemTableSortTest, Tie) { auto t0 = Tie {0, 0}; @@ -80,4 +269,126 @@ TEST_F(MemTableSortTest, Tie) { EXPECT_FALSE(it3.next()); } +// Keys are ordered by (k1, k2); the second key column must only be used to +// refine rows that tie on the first one. +TEST_F(MemTableSortTest, DupKeysOrdersByAllKeyColumns) { + std::vector rows = {{2, "b", 20}, {1, "b", 11}, {2, "a", 21}, {1, "a", 10}}; + std::unique_ptr out; + ASSERT_NO_FATAL_FAILURE(run(KeysType::DUP_KEYS, rows, 1, &out)); + ASSERT_EQ(4, out->rows()); + EXPECT_EQ(1, k1_of(*out, 0)); + EXPECT_EQ("a", k2_of(*out, 0)); + EXPECT_EQ(1, k1_of(*out, 1)); + EXPECT_EQ("b", k2_of(*out, 1)); + EXPECT_EQ(2, k1_of(*out, 2)); + EXPECT_EQ("a", k2_of(*out, 2)); + EXPECT_EQ(2, k1_of(*out, 3)); + EXPECT_EQ("b", k2_of(*out, 3)); +} + +// Rows sharing the whole key are stabilised on descending row position for +// DUP_KEYS, i.e. reverse insertion order. Nothing depends on that direction in +// principle, but a number of regression cases record it, so pin it down. +TEST_F(MemTableSortTest, DupKeysReversesEqualKeys) { + std::vector rows = {{1, "a", 100}, {1, "a", 101}, {1, "a", 102}}; + std::unique_ptr out; + ASSERT_NO_FATAL_FAILURE(run(KeysType::DUP_KEYS, rows, 1, &out)); + ASSERT_EQ(3, out->rows()); + EXPECT_EQ(102, v_of(*out, 0)); + EXPECT_EQ(101, v_of(*out, 1)); + EXPECT_EQ(100, v_of(*out, 2)); +} + +// NULL sorts before any value, matching the nan_direction_hint = -1 the previous +// comparator used. +TEST_F(MemTableSortTest, NullKeySortsFirst) { + std::vector rows = {{1, "b", 2}, {1, nullptr, 1}, {1, "a", 3}}; + std::unique_ptr out; + ASSERT_NO_FATAL_FAILURE(run(KeysType::DUP_KEYS, rows, 1, &out)); + ASSERT_EQ(3, out->rows()); + EXPECT_EQ("", k2_of(*out, 0)); + EXPECT_EQ("a", k2_of(*out, 1)); + EXPECT_EQ("b", k2_of(*out, 2)); +} + +// Splitting the same rows across several insert() calls must not change the +// result: _sort() maps a sorted row position back to its RowInBlock through the +// base of the appended range, which only holds if row positions stay contiguous +// across insert() calls. The equal keys make that mapping observable. +TEST_F(MemTableSortTest, ResultIsIndependentOfBatching) { + std::vector rows = {{3, "c", 1}, {1, "a", 2}, {2, "b", 3}, {1, "b", 4}, + {3, "a", 5}, {2, "a", 6}, {1, "a", 7}, {2, "a", 8}}; + std::unique_ptr one; + std::unique_ptr many; + ASSERT_NO_FATAL_FAILURE(run(KeysType::DUP_KEYS, rows, 1, &one)); + ASSERT_NO_FATAL_FAILURE(run(KeysType::DUP_KEYS, rows, 3, &many)); + ASSERT_EQ(one->rows(), many->rows()); + for (size_t i = 0; i < one->rows(); ++i) { + EXPECT_EQ(k1_of(*one, i), k1_of(*many, i)) << "row " << i; + EXPECT_EQ(k2_of(*one, i), k2_of(*many, i)) << "row " << i; + EXPECT_EQ(v_of(*one, i), v_of(*many, i)) << "row " << i; + } +} + +// For UNIQUE_KEYS the last inserted row must win, which relies on equal keys +// being ordered ascending by row position before aggregation runs. +TEST_F(MemTableSortTest, UniqueKeysLastWriterWins) { + std::vector rows = {{1, "a", 10}, {2, "b", 20}, {1, "a", 11}, {1, "a", 12}}; + std::unique_ptr out; + ASSERT_NO_FATAL_FAILURE(run(KeysType::UNIQUE_KEYS, rows, 1, &out)); + ASSERT_EQ(2, out->rows()); + EXPECT_EQ(1, k1_of(*out, 0)); + EXPECT_EQ(12, v_of(*out, 0)) << "the last inserted value must survive"; + EXPECT_EQ(2, k1_of(*out, 1)); + EXPECT_EQ(20, v_of(*out, 1)); +} + +// shrink_memtable_by_agg() aggregates without finalising, so a surviving row +// carries its aggregate state into the next round and to_block() finalises it. +// Interleaving that with more inserts must not change the result -- this covers +// the state handoff between rounds, and the release of those states, which is +// where holding rows by value differs most from holding them behind a pointer. +TEST_F(MemTableSortTest, AggKeysSurviveShrinkRounds) { + std::vector rows = {{1, "a", 1}, {2, "b", 10}, {1, "a", 2}, {3, "c", 100}, + {2, "b", 20}, {1, "a", 4}, {2, "b", 30}, {1, "a", 8}}; + std::unique_ptr plain; + std::unique_ptr shrunk; + ASSERT_NO_FATAL_FAILURE(run(KeysType::AGG_KEYS, rows, 4, &plain)); + ASSERT_NO_FATAL_FAILURE(run(KeysType::AGG_KEYS, rows, 4, &shrunk, true)); + ASSERT_EQ(3, shrunk->rows()); + ASSERT_EQ(plain->rows(), shrunk->rows()); + for (size_t i = 0; i < plain->rows(); ++i) { + EXPECT_EQ(k1_of(*plain, i), k1_of(*shrunk, i)) << "row " << i; + EXPECT_EQ(v_of(*plain, i), v_of(*shrunk, i)) << "row " << i; + } + EXPECT_EQ(15, v_of(*shrunk, 0)); // 1 + 2 + 4 + 8 + EXPECT_EQ(60, v_of(*shrunk, 1)); // 10 + 20 + 30 + EXPECT_EQ(100, v_of(*shrunk, 2)); +} + +// Same handoff for UNIQUE_KEYS, where a round must keep the newest row rather +// than accumulate. +TEST_F(MemTableSortTest, UniqueKeysSurviveShrinkRounds) { + std::vector rows = {{1, "a", 10}, {2, "b", 20}, {1, "a", 11}, + {2, "b", 21}, {1, "a", 12}, {3, "c", 30}}; + std::unique_ptr out; + ASSERT_NO_FATAL_FAILURE(run(KeysType::UNIQUE_KEYS, rows, 3, &out, true)); + ASSERT_EQ(3, out->rows()); + EXPECT_EQ(12, v_of(*out, 0)); + EXPECT_EQ(21, v_of(*out, 1)); + EXPECT_EQ(30, v_of(*out, 2)); +} + +// AGG_KEYS with SUM: every duplicate must be folded into the group exactly once. +TEST_F(MemTableSortTest, AggKeysSumsDuplicates) { + std::vector rows = {{1, "a", 1}, {2, "b", 100}, {1, "a", 2}, {1, "a", 4}, {2, "b", 200}}; + std::unique_ptr out; + ASSERT_NO_FATAL_FAILURE(run(KeysType::AGG_KEYS, rows, 1, &out)); + ASSERT_EQ(2, out->rows()); + EXPECT_EQ(1, k1_of(*out, 0)); + EXPECT_EQ(7, v_of(*out, 0)); + EXPECT_EQ(2, k1_of(*out, 1)); + EXPECT_EQ(300, v_of(*out, 1)); +} + } // namespace doris From 89e908d722a2936e67f81417dfad378e5bc32952 Mon Sep 17 00:00:00 2001 From: liaoxin Date: Sat, 8 Aug 2026 00:46:30 +0800 Subject: [PATCH 2/2] [improvement](load) sort memtable with the vectorized ColumnSorter MemTable::_sort() ran its own multi-key sort: pdqsort over the row array, with the comparator passed as a std::function that called the virtual IColumn::compare_at once per comparison. Every comparison paid an indirect call, a virtual dispatch and two random accesses into the block before it could look at the key. The query engine already has the sort this needs. ColumnSorter (be/src/exec/sort/sort_block.h) implements the same sort-and-tie algorithm, but keeps an inline copy of the key next to the row id, so a comparison is a typed, inlinable operation on a compact array. _sort() and _sort_by_cluster_keys() now build an IColumn::Permutation and run it through ColumnSorter; _sort_one_column() and class Tie have no users left and are removed. _sort_by_cluster_keys() no longer needs a row object per row either -- the LSN sidecar it was carrying is already indexed by row position, so the permutation can reorder it directly. Measured on this data shape (7.05M rows, one key column, Release build): key type current ColumnSorter speedup int32 13.2 s 0.96 s 13.7x int64 13.3 s 1.07 s 12.4x decimal128(20,2) 13.7 s 1.28 s 10.7x varchar ~4B 19.6 s 3.03 s 6.5x varchar ~40B 21.9 s 3.75 s 5.8x nullable int32 16.0 s 1.04 s 15.5x nullable varchar ~4B 25.4 s 2.81 s 9.0x Fixed-width keys gain the most because their inline value is the value itself; string keys still dereference the arena for the memcmp, so the longer the key the smaller the gain. The inline permutation costs memory in proportion to the key width: 8 bytes per row for INT32, 16 for INT64, 24 for a StringRef, 32 for Decimal128 after alignment. It is a std::vector local to ColumnSorter::_sort_by_inline_permutation, so it is released between key columns and the peak holds one of them, plus 8 bytes per row for the permutation and 1 for the equal flags. Against the 24 bytes per row the memtable now spends on the rows themselves, the peak of the two together still comes out below what the previous sort needed. Ordering is unchanged, tie-break included: rows whose whole key is equal are still stabilised on descending row position for DUP_KEYS and on ascending row position for everything else. --- be/src/load/memtable/memtable.cpp | 157 ++++++++++++------- be/src/load/memtable/memtable.h | 69 ++------ be/test/load/memtable/memtable_sort_test.cpp | 72 ++------- 3 files changed, 127 insertions(+), 171 deletions(-) diff --git a/be/src/load/memtable/memtable.cpp b/be/src/load/memtable/memtable.cpp index 32baa6f6fdceb3..556e7addb28b8e 100644 --- a/be/src/load/memtable/memtable.cpp +++ b/be/src/load/memtable/memtable.cpp @@ -30,6 +30,9 @@ #include "bvar/bvar.h" #include "common/config.h" #include "core/column/column.h" +#include "exec/sort/hybrid_sorter.h" +#include "exec/sort/sort_block.h" +#include "exec/sort/sort_description.h" #include "exprs/aggregate/aggregate_function_reader.h" #include "exprs/aggregate/aggregate_function_simple_factory.h" #include "load/delta_writer/delta_writer_context.h" @@ -390,44 +393,97 @@ Status MemTable::_put_into_output(Block& in_block) { row_pos_vec.data() + in_block.rows()); } -void MemTable::_sort_one_column(DorisVector& row_in_blocks, Tie& tie, - std::function cmp) { - auto iter = tie.iter(); +// ColumnSorter keeps an inline copy of the key next to the row id, so a +// comparison no longer chases a row pointer nor pays a virtual +// IColumn::compare_at, which is what dominated the previous kernel. +size_t MemTable::_sort_permutation_by_key_columns(MutableBlock& block, + const std::vector& key_col_idx, + IColumn::Permutation& perm, + bool descending_row_pos) { + const size_t num_rows = perm.size(); + if (num_rows == 0) { + return 0; + } + EqualFlags flags(num_rows, 1); + EqualRange range {0, static_cast(num_rows)}; + HybridSorter hybrid_sorter; + for (int idx : key_col_idx) { + ColumnWithSortDescription col {block.get_column_by_position(idx).get(), + SortColumnDescription(idx, 1 /*asc*/, -1 /*nulls first*/)}; + ColumnSorter sorter(col, hybrid_sorter, 0 /*no limit*/); + // last_column is deliberately never true: the equal ranges are still + // needed below for the row position tie-break. + sorter(flags, perm, range, false); + } + + // Sort an extra round by row position to make the sort stable. perm elements + // are row positions, so sorting them directly is enough. + size_t same_keys_num = 0; + EqualRangeIterator iter(flags, range.first, range.second); while (iter.next()) { - pdqsort(std::next(row_in_blocks.begin(), static_cast(iter.left())), - std::next(row_in_blocks.begin(), static_cast(iter.right())), - [&cmp](const RowInBlock& lhs, const RowInBlock& rhs) -> bool { - return cmp(lhs, rhs) < 0; - }); - tie[iter.left()] = 0; - for (auto i = iter.left() + 1; i < iter.right(); i++) { - tie[i] = (cmp(row_in_blocks[i - 1], row_in_blocks[i]) == 0); + if (descending_row_pos) { + pdqsort(perm.begin() + iter.range_begin, perm.begin() + iter.range_end, + std::greater<>()); + } else { + pdqsort(perm.begin() + iter.range_begin, perm.begin() + iter.range_end); } + same_keys_num += static_cast(iter.range_end - iter.range_begin); } + return same_keys_num; } size_t MemTable::_sort() { SCOPED_RAW_TIMER(&_stat.sort_ns); _stat.sort_times++; size_t same_keys_num = 0; + const bool is_dup = (_keys_type == KeysType::DUP_KEYS); // sort new rows - Tie tie = Tie(_last_sorted_pos, _row_in_blocks->size()); - for (size_t i = 0; i < _tablet_schema->num_key_columns(); i++) { - auto cmp = [&](const RowInBlock& lhs, const RowInBlock& rhs) -> int { - return _input_mutable_block.compare_one_column(lhs._row_pos, rhs._row_pos, i, -1); - }; - _sort_one_column(*_row_in_blocks, tie, cmp); - } - bool is_dup = (_keys_type == KeysType::DUP_KEYS); - // sort extra round by _row_pos to make the sort stable - auto iter = tie.iter(); - while (iter.next()) { - pdqsort(std::next(_row_in_blocks->begin(), iter.left()), - std::next(_row_in_blocks->begin(), iter.right()), - [&is_dup](const RowInBlock& lhs, const RowInBlock& rhs) -> bool { - return is_dup ? lhs._row_pos > rhs._row_pos : lhs._row_pos < rhs._row_pos; - }); - same_keys_num += iter.right() - iter.left(); + const size_t num_new_rows = _row_in_blocks->size() - _last_sorted_pos; + if (num_new_rows > 0) { + // Rows appended since the last sort occupy a contiguous, increasing range + // of positions in _input_mutable_block (see insert()), so a sorted row + // position maps back to its RowInBlock by subtracting the base. + const size_t base = (*_row_in_blocks)[_last_sorted_pos]._row_pos; + DCHECK_EQ((*_row_in_blocks)[_row_in_blocks->size() - 1]._row_pos, base + num_new_rows - 1); + + IColumn::Permutation perm(num_new_rows); + for (size_t i = 0; i < num_new_rows; i++) { + perm[i] = base + i; + } + std::vector key_col_idx(_tablet_schema->num_key_columns()); + for (size_t i = 0; i < key_col_idx.size(); i++) { + key_col_idx[i] = static_cast(i); + } + same_keys_num = + _sort_permutation_by_key_columns(_input_mutable_block, key_col_idx, perm, is_dup); + + // Rebase perm onto the appended range so it indexes _row_in_blocks + // directly. Subtracting a constant keeps the order the sort produced. + for (size_t i = 0; i < num_new_rows; i++) { + perm[i] -= base; + } + // Apply the permutation by following its cycles. That is one move per + // row -- half of what a sorted copy costs -- and, more importantly, it + // needs no second N-element array of shared_ptr next to the live + // _row_in_blocks. perm is dead afterwards, so it doubles as the marker + // for rows already moved into place. + RowInBlock* rows = _row_in_blocks->data() + _last_sorted_pos; + for (size_t i = 0; i < num_new_rows; i++) { + size_t src = perm[i]; + if (src == i) { + continue; + } + RowInBlock held = rows[i]; + size_t hole = i; + while (src != i) { + rows[hole] = rows[src]; + perm[hole] = hole; + hole = src; + src = perm[src]; + } + rows[hole] = held; + perm[hole] = hole; + } } // merge new rows and old rows _vec_row_comparator->set_block(&_input_mutable_block); @@ -456,50 +512,43 @@ Status MemTable::_sort_by_cluster_keys() { MutableBlock mutable_block = MutableBlock::build_mutable_block(std::move(in_block)); _output_mutable_block = MutableBlock::build_mutable_block(std::move(clone_block)); - DorisVector row_in_blocks; - row_in_blocks.reserve(mutable_block.rows()); + // The LSN sidecar is indexed by row position, so the permutation below can + // reorder it directly; there is no need to materialise a row object per row. + DorisVector input_row_binlog_lsns; if (_need_row_binlog_lsn) { DCHECK_EQ(_output_row_binlog_lsns.size(), mutable_block.rows()); - } - for (size_t i = 0; i < mutable_block.rows(); i++) { - row_in_blocks.emplace_back(i, _need_row_binlog_lsn ? _output_row_binlog_lsns[i] : 0); - } - if (_need_row_binlog_lsn) { - _output_row_binlog_lsns.clear(); + input_row_binlog_lsns.swap(_output_row_binlog_lsns); _output_row_binlog_lsns.reserve(mutable_block.rows()); } - Tie tie = Tie(0, mutable_block.rows()); - + std::vector key_col_idx; + key_col_idx.reserve(_tablet_schema->cluster_key_uids().size()); for (auto cid : _tablet_schema->cluster_key_uids()) { auto index = _tablet_schema->field_index(cid); if (index == -1) { return Status::InternalError("could not find cluster key column with unique_id=" + std::to_string(cid) + " in tablet schema"); } - auto cmp = [&](const RowInBlock& lhs, const RowInBlock& rhs) -> int { - return mutable_block.compare_one_column(lhs._row_pos, rhs._row_pos, index, -1); - }; - _sort_one_column(row_in_blocks, tie, cmp); + key_col_idx.push_back(index); } - - // sort extra round by _row_pos to make the sort stable - auto iter = tie.iter(); - while (iter.next()) { - pdqsort(std::next(row_in_blocks.begin(), iter.left()), - std::next(row_in_blocks.begin(), iter.right()), - [](const RowInBlock& lhs, const RowInBlock& rhs) -> bool { - return lhs._row_pos < rhs._row_pos; - }); + // Cluster keys only exist on MOW tables, whose order is always stabilised on + // ascending row position. + IColumn::Permutation perm(mutable_block.rows()); + for (size_t i = 0; i < perm.size(); i++) { + perm[i] = i; } + static_cast(_sort_permutation_by_key_columns(mutable_block, key_col_idx, perm, + false /*descending_row_pos*/)); in_block = mutable_block.to_block(); SCOPED_RAW_TIMER(&_stat.put_into_output_ns); DorisVector row_pos_vec; DCHECK(in_block.rows() <= std::numeric_limits::max()); row_pos_vec.reserve(in_block.rows()); - for (const auto& row : row_in_blocks) { - row_pos_vec.emplace_back(row._row_pos); - _append_output_row_binlog_lsn(row); + for (size_t i = 0; i < perm.size(); i++) { + row_pos_vec.emplace_back(static_cast(perm[i])); + if (_need_row_binlog_lsn) { + _output_row_binlog_lsns.emplace_back(input_row_binlog_lsns[perm[i]]); + } } std::vector column_offset; for (int i = 0; i < _column_offset.size(); ++i) { diff --git a/be/src/load/memtable/memtable.h b/be/src/load/memtable/memtable.h index 1525359f07761a..0e37ef0a3dbea7 100644 --- a/be/src/load/memtable/memtable.h +++ b/be/src/load/memtable/memtable.h @@ -67,64 +67,6 @@ struct RowInBlock { : _row_pos(row), _row_binlog_lsn(row_binlog_lsn) {} }; -class Tie { -public: - class Iter { - public: - Iter(Tie& tie) : _tie(tie), _next(tie._begin + 1) {} - size_t left() const { return _left; } - size_t right() const { return _right; } - - // return false means no more ranges - bool next() { - if (_next >= _tie._end) { - return false; - } - _next = _find(1, _next); - if (_next >= _tie._end) { - return false; - } - _left = _next - 1; - _next = _find(0, _next); - _right = _next; - return true; - } - - private: - size_t _find(uint8_t value, size_t start) { - if (start >= _tie._end) { - return start; - } - size_t offset = start - _tie._begin; - size_t size = _tie._end - start; - void* p = std::memchr(_tie._bits.data() + offset, value, size); - if (p == nullptr) { - return _tie._end; - } - return static_cast(p) - _tie._bits.data() + _tie._begin; - } - - private: - Tie& _tie; - size_t _left; - size_t _right; - size_t _next; - }; - -public: - Tie(size_t begin, size_t end) : _begin(begin), _end(end) { - _bits = std::vector(_end - _begin, 1); - } - uint8_t operator[](size_t i) const { return _bits[i - _begin]; } - uint8_t& operator[](size_t i) { return _bits[i - _begin]; } - Iter iter() { return Iter(*this); } - -private: - const size_t _begin; - const size_t _end; - std::vector _bits; -}; - class RowInBlockComparator { public: RowInBlockComparator(std::shared_ptr tablet_schema) @@ -265,11 +207,18 @@ class MemTable { size_t _last_sorted_pos = 0; size_t _last_agg_pos = 0; + // Sorts `perm` (row positions into `block`) by `key_col_idx`, reusing the + // vectorized ColumnSorter the query engine uses, then stabilises rows with + // an equal key on their row position -- descending when `descending_row_pos` + // is set, which is what DUP_KEYS has always done. + // Returns the number of rows that share their key with a neighbour. + static size_t _sort_permutation_by_key_columns(MutableBlock& block, + const std::vector& key_col_idx, + IColumn::Permutation& perm, + bool descending_row_pos); //return number of same keys size_t _sort(); Status _sort_by_cluster_keys(); - void _sort_one_column(DorisVector& row_in_blocks, Tie& tie, - std::function cmp); template void _finalize_one_row(RowInBlock& row, MutableBlock& mutable_block, int row_pos); void _init_row_for_agg(RowInBlock& row, MutableBlock& mutable_block); diff --git a/be/test/load/memtable/memtable_sort_test.cpp b/be/test/load/memtable/memtable_sort_test.cpp index 13cb77102ad1a6..b15c83eca84e5d 100644 --- a/be/test/load/memtable/memtable_sort_test.cpp +++ b/be/test/load/memtable/memtable_sort_test.cpp @@ -212,63 +212,6 @@ class MemTableSortTest : public testing::Test { static int32_t v_of(const Block& b, size_t row) { return int_of(b, 2, row); } }; -TEST_F(MemTableSortTest, Tie) { - auto t0 = Tie {0, 0}; - EXPECT_FALSE(t0.iter().next()); - - auto tie = Tie {0, 1}; - EXPECT_FALSE(tie.iter().next()); - - auto t = Tie {10, 30}; - for (int i = 10; i < 30; i++) { - EXPECT_EQ(t[i], 1); - } - - auto it1 = t.iter(); - EXPECT_TRUE(it1.next()); - EXPECT_EQ(it1.left(), 10); - EXPECT_EQ(it1.right(), 30); - - EXPECT_FALSE(it1.next()); - - t[13] = t[14] = t[22] = t[29] = 0; - auto it2 = t.iter(); - - EXPECT_TRUE(it2.next()); - EXPECT_EQ(it2.left(), 10); - EXPECT_EQ(it2.right(), 13); - - EXPECT_TRUE(it2.next()); - EXPECT_EQ(it2.left(), 14); - EXPECT_EQ(it2.right(), 22); - - EXPECT_TRUE(it2.next()); - EXPECT_EQ(it2.left(), 22); - EXPECT_EQ(it2.right(), 29); - - EXPECT_FALSE(it2.next()); - EXPECT_FALSE(it2.next()); - - // 100000000... - for (int i = 11; i < 30; i++) { - t[i] = 0; - } - EXPECT_FALSE(t.iter().next()); - - // 000000000... - t[10] = 0; - EXPECT_FALSE(t.iter().next()); - - // 000000000...001 - t[29] = 1; - auto it3 = t.iter(); - EXPECT_TRUE(it3.next()); - EXPECT_EQ(it3.left(), 28); - EXPECT_EQ(it3.right(), 30); - - EXPECT_FALSE(it3.next()); -} - // Keys are ordered by (k1, k2); the second key column must only be used to // refine rows that tie on the first one. TEST_F(MemTableSortTest, DupKeysOrdersByAllKeyColumns) { @@ -286,6 +229,21 @@ TEST_F(MemTableSortTest, DupKeysOrdersByAllKeyColumns) { EXPECT_EQ("b", k2_of(*out, 3)); } +// The permutation is applied to _row_in_blocks in place by following its +// cycles, so cover a permutation that is one long cycle rather than the short +// swaps the other cases happen to produce. Keys [5,1,2,3,4] sort to sources +// [1,2,3,4,0], i.e. a single 5-cycle. +TEST_F(MemTableSortTest, SingleCyclePermutation) { + std::vector rows = {{5, "a", 50}, {1, "a", 10}, {2, "a", 20}, {3, "a", 30}, {4, "a", 40}}; + std::unique_ptr out; + ASSERT_NO_FATAL_FAILURE(run(KeysType::DUP_KEYS, rows, 1, &out)); + ASSERT_EQ(5, out->rows()); + for (size_t i = 0; i < 5; ++i) { + EXPECT_EQ(static_cast(i + 1), k1_of(*out, i)) << "row " << i; + EXPECT_EQ(static_cast((i + 1) * 10), v_of(*out, i)) << "row " << i; + } +} + // Rows sharing the whole key are stabilised on descending row position for // DUP_KEYS, i.e. reverse insertion order. Nothing depends on that direction in // principle, but a number of regression cases record it, so pin it down.