Skip to content

Commit dcee505

Browse files
author
zclllyybb
authored
[refactor](be) Enforce COW ownership for assume_mutable (#63001)
Problem Summary: This PR changes the BE COW contract around `assume_mutable`: callers may only use it when they already own the column exclusively, and the helper now behaves as an ownership assertion instead of a silent mutable borrow. Shared `ColumnPtr` / `Block` data must go through `mutate()` or one of the scoped owner-slot mutation helpers before being modified. This lets blocks and columns rely on real COW semantics and avoids both unsafe in-place mutation of shared data and unnecessary cross-operator copies. The main changes are: - Make `assume_mutable` / `assume_mutable_ref` validate exclusive ownership and document audited usage in `docs/dev/be-cow-assume-mutable-audit.md`. - Add scoped COW mutation APIs for common owner-slot patterns, including whole-block scoped mutation, single-column scoped mutation, and rvalue-only stealing mutation for `Block::mutate_columns()`. - Migrate BE paths that previously assumed mutable access to explicit COW ownership transfer, scoped restore, or `MutableBlock` / `MutableColumns` usage in hot paths. - Fix affected storage, scanner, external reader, parquet/orc/json/table-format, variant, aggregation, local exchange, and schema scanner paths so errors restore moved-out columns and nested subcolumns are written back through their owner slots. - Add focused BE UT coverage for the COW contract: shared-column detach, scoped restore on early error, block schema access through scoped guards, LocalExchanger restore-on-error, and table-format partition/missing-column COW behavior.
1 parent 4938d63 commit dcee505

232 files changed

Lines changed: 5171 additions & 2096 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

be/src/core/block/block.cpp

Lines changed: 190 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,51 @@ template void clear_blocks<Block>(moodycamel::ConcurrentQueue<Block>&,
7979
template void clear_blocks<BlockUPtr>(moodycamel::ConcurrentQueue<BlockUPtr>&,
8080
RuntimeProfile::Counter* memory_used_counter);
8181

82+
namespace {
83+
84+
// The no-clone fast path is only safe when the whole column tree is uniquely
85+
// owned. A composite column with shared children still needs COW detachment.
86+
bool is_recursively_exclusive(const IColumn& column) {
87+
if (!column.is_exclusive()) {
88+
return false;
89+
}
90+
91+
bool exclusive = true;
92+
IColumn::ColumnCallback callback = [&](IColumn::WrappedPtr& subcolumn) {
93+
if (!exclusive) {
94+
return;
95+
}
96+
const ColumnPtr& subcolumn_ptr = const_cast<const IColumn::WrappedPtr&>(subcolumn);
97+
DCHECK(subcolumn_ptr);
98+
exclusive = is_recursively_exclusive(*subcolumn_ptr);
99+
};
100+
// `for_each_subcolumn` only exposes a mutable callback type. This callback
101+
// only reads the wrapped pointers and never calls the non-const accessors.
102+
const_cast<IColumn&>(column).for_each_subcolumn(callback);
103+
return exclusive;
104+
}
105+
106+
// Acquire one live Block slot transactionally. Shared columns are detached while
107+
// the original slot is still intact, so a clone failure cannot leave Block with
108+
// a moved-from/null column. Exclusive column trees keep the stealing fast path.
109+
MutableColumnPtr scoped_mutate_column(ColumnPtr& column, const DataTypePtr& type) {
110+
DCHECK(type);
111+
if (!column) {
112+
return type->create_column();
113+
}
114+
115+
MutableColumnPtr mutable_column;
116+
if (is_recursively_exclusive(*column)) {
117+
mutable_column = std::move(*column).mutate();
118+
} else {
119+
mutable_column = IColumn::mutate(column);
120+
}
121+
column = nullptr;
122+
return mutable_column;
123+
}
124+
125+
} // namespace
126+
82127
Block::Block(std::initializer_list<ColumnWithTypeAndName> il) : data {il} {}
83128

84129
Block::Block(ColumnsWithTypeAndName data_) : data {std::move(data_)} {}
@@ -576,12 +621,127 @@ Columns Block::get_columns_and_convert() {
576621
return columns;
577622
}
578623

579-
MutableColumns Block::mutate_columns() {
624+
Block::ScopedMutableColumns::ScopedMutableColumns(Block& block) : _block(&block) {
625+
const size_t num_columns = block.data.size();
626+
_columns.resize(num_columns);
627+
size_t acquired_columns = 0;
628+
try {
629+
for (; acquired_columns < num_columns; ++acquired_columns) {
630+
auto& column_with_type_and_name = block.data[acquired_columns];
631+
_columns[acquired_columns] = scoped_mutate_column(column_with_type_and_name.column,
632+
column_with_type_and_name.type);
633+
}
634+
} catch (...) {
635+
for (size_t i = 0; i < acquired_columns; ++i) {
636+
block.data[i].column = std::move(_columns[i]);
637+
}
638+
_block = nullptr;
639+
throw;
640+
}
641+
}
642+
643+
Block::ScopedMutableColumns::~ScopedMutableColumns() {
644+
restore();
645+
}
646+
647+
Block::ScopedMutableColumns::ScopedMutableColumns(ScopedMutableColumns&& other) noexcept
648+
: _block(std::exchange(other._block, nullptr)), _columns(std::move(other._columns)) {}
649+
650+
Block::ScopedMutableColumns& Block::ScopedMutableColumns::operator=(
651+
ScopedMutableColumns&& other) noexcept {
652+
if (this != &other) {
653+
restore();
654+
_block = std::exchange(other._block, nullptr);
655+
_columns = std::move(other._columns);
656+
}
657+
return *this;
658+
}
659+
660+
const DataTypePtr& Block::ScopedMutableColumns::get_datatype_by_position(size_t position) const {
661+
DCHECK(_block != nullptr);
662+
return _block->get_by_position(position).type;
663+
}
664+
665+
const std::string& Block::ScopedMutableColumns::get_name_by_position(size_t position) const {
666+
DCHECK(_block != nullptr);
667+
return _block->get_by_position(position).name;
668+
}
669+
670+
MutableColumns Block::ScopedMutableColumns::release() {
671+
DCHECK(_block != nullptr);
672+
_block = nullptr;
673+
return std::move(_columns);
674+
}
675+
676+
void Block::ScopedMutableColumns::restore() {
677+
if (_block != nullptr) {
678+
_block->set_columns(std::move(_columns));
679+
_block = nullptr;
680+
}
681+
}
682+
683+
Block::ScopedMutableColumn::ScopedMutableColumn(Block& block, size_t position)
684+
: _block(&block), _position(position) {
685+
DCHECK_LT(_position, _block->data.size());
686+
auto& column_with_type_and_name = _block->data[_position];
687+
DCHECK(column_with_type_and_name.type);
688+
_column =
689+
scoped_mutate_column(column_with_type_and_name.column, column_with_type_and_name.type);
690+
}
691+
692+
Block::ScopedMutableColumn::~ScopedMutableColumn() {
693+
restore();
694+
}
695+
696+
Block::ScopedMutableColumn::ScopedMutableColumn(ScopedMutableColumn&& other) noexcept
697+
: _block(std::exchange(other._block, nullptr)),
698+
_position(other._position),
699+
_column(std::move(other._column)) {}
700+
701+
Block::ScopedMutableColumn& Block::ScopedMutableColumn::operator=(
702+
ScopedMutableColumn&& other) noexcept {
703+
if (this != &other) {
704+
restore();
705+
_block = std::exchange(other._block, nullptr);
706+
_position = other._position;
707+
_column = std::move(other._column);
708+
}
709+
return *this;
710+
}
711+
712+
void Block::ScopedMutableColumn::restore() {
713+
if (_block != nullptr) {
714+
DCHECK_LT(_position, _block->data.size());
715+
_block->data[_position].column = std::move(_column);
716+
_block = nullptr;
717+
}
718+
}
719+
720+
Block::ScopedMutableColumns Block::mutate_columns_scoped() & {
721+
return ScopedMutableColumns(*this);
722+
}
723+
724+
Block::ScopedMutableColumn Block::mutate_column_scoped(size_t position) & {
725+
return ScopedMutableColumn(*this, position);
726+
}
727+
728+
ScopedMutableBlock::ScopedMutableBlock(Block* block) {
729+
DCHECK(block != nullptr);
730+
DataTypes data_types = block->get_data_types();
731+
std::vector<std::string> names = block->get_names();
732+
auto columns_guard = block->mutate_columns_scoped();
733+
_mutable_block.data_types() = std::move(data_types);
734+
_mutable_block.get_names() = std::move(names);
735+
_mutable_block.set_mutable_columns(columns_guard.release());
736+
_block = block;
737+
}
738+
739+
MutableColumns Block::mutate_columns() && {
580740
size_t num_columns = data.size();
581741
MutableColumns columns(num_columns);
582742
for (size_t i = 0; i < num_columns; ++i) {
583743
DCHECK(data[i].type);
584-
columns[i] = data[i].column ? (*std::move(data[i].column)).mutate()
744+
columns[i] = data[i].column ? IColumn::mutate(std::move(data[i].column))
585745
: data[i].type->create_column();
586746
}
587747
return columns;
@@ -644,7 +804,7 @@ void Block::clear() {
644804
data.clear();
645805
}
646806

647-
void Block::clear_column_data(int64_t column_size) noexcept {
807+
void Block::clear_column_data(int64_t column_size) {
648808
SCOPED_SKIP_MEMORY_CHECK();
649809
// data.size() greater than column_size, means here have some
650810
// function exec result in block, need erase it here
@@ -655,9 +815,26 @@ void Block::clear_column_data(int64_t column_size) noexcept {
655815
}
656816
for (auto& d : data) {
657817
if (d.column) {
658-
// Temporarily disable reference count check because a column might be referenced multiple times within a block.
659-
// Queries like this: `select c, c from t1;`
660-
(*std::move(d.column)).assume_mutable()->clear();
818+
if (d.column->is_exclusive()) {
819+
d.column->assume_mutable()->clear();
820+
} else {
821+
d.column = d.column->clone_empty();
822+
}
823+
}
824+
}
825+
}
826+
827+
void Block::clear_column_data(const std::vector<uint32_t>& columns_to_clear) {
828+
SCOPED_SKIP_MEMORY_CHECK();
829+
for (auto col : columns_to_clear) {
830+
DCHECK_LT(col, data.size());
831+
auto& column = data[col].column;
832+
if (column) {
833+
if (column->is_exclusive()) {
834+
column->assume_mutable()->clear();
835+
} else {
836+
column = column->clone_empty();
837+
}
661838
}
662839
}
663840
}
@@ -1085,7 +1262,13 @@ void Block::shrink_char_type_column_suffix_zero(const std::vector<size_t>& char_
10851262
for (auto idx : char_type_idx) {
10861263
if (idx < data.size()) {
10871264
auto& col_and_name = this->get_by_position(idx);
1088-
col_and_name.column->assume_mutable()->shrink_padding_chars();
1265+
if (col_and_name.column->is_exclusive()) {
1266+
col_and_name.column->assume_mutable()->shrink_padding_chars();
1267+
} else {
1268+
auto mutable_col = std::move(*col_and_name.column).mutate();
1269+
mutable_col->shrink_padding_chars();
1270+
col_and_name.column = std::move(mutable_col);
1271+
}
10891272
}
10901273
}
10911274
}

0 commit comments

Comments
 (0)