@@ -79,6 +79,51 @@ template void clear_blocks<Block>(moodycamel::ConcurrentQueue<Block>&,
7979template 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+
82127Block::Block (std::initializer_list<ColumnWithTypeAndName> il) : data {il} {}
83128
84129Block::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