[Enhancement] split chunk of HashTable

be/src/column/column_helper.cpp

@@ -25,6 +25,7 @@
 #include "column/vectorized_fwd.h"
 #include "gutil/casts.h"
 #include "simd/simd.h"
+#include "storage/chunk_helper.h"
 #include "types/logical_type_infra.h"
 #include "util/date_func.h"
 #include "util/percentile_value.h"
@@ -469,7 +470,7 @@ size_t ChunkSliceTemplate<Ptr>::skip(size_t skip_rows) {
 
 // Cutoff required rows from this chunk
 template <class Ptr>
-Ptr ChunkSliceTemplate<Ptr>::cutoff(size_t required_rows) {
+ChunkUniquePtr ChunkSliceTemplate<Ptr>::cutoff(size_t required_rows) {
     DCHECK(!empty());
     size_t cut_rows = std::min(rows(), required_rows);
     auto res = chunk->clone_empty(cut_rows);
@@ -482,7 +483,35 @@ Ptr ChunkSliceTemplate<Ptr>::cutoff(size_t required_rows) {
     return res;
 }
 
+// Specialized for SegmentedChunkPtr
+template <>
+ChunkUniquePtr ChunkSliceTemplate<SegmentedChunkPtr>::cutoff(size_t required_rows) {
+    DCHECK(!empty());
+    // cutoff a chunk from current segment, if it doesn't meet the requirement just let it be
+    ChunkPtr segment = chunk->segments()[segment_id];
+    size_t segment_offset = offset % chunk->segment_size();
+    size_t cut_rows = std::min(segment->num_rows() - segment_offset, required_rows);
+
+    auto res = segment->clone_empty(cut_rows);
+    res->append(*segment, segment_offset, cut_rows);
+    offset += cut_rows;
+
+    // move to next segment and release previous one
+    size_t new_segment_id = offset / chunk->segment_size();
+    if (new_segment_id != segment_id) {
+        chunk->segments()[segment_id].reset();
+        segment_id = new_segment_id;
+    }
+
+    if (empty()) {
+        chunk.reset();
+        offset = 0;
+    }
+    return res;
+}
+
 template struct ChunkSliceTemplate<ChunkPtr>;
 template struct ChunkSliceTemplate<ChunkUniquePtr>;
+template struct ChunkSliceTemplate<SegmentedChunkPtr>;
 
 } // namespace starrocks

be/src/column/column_helper.h

@@ -513,12 +513,13 @@ class ColumnHelper {
 template <class Ptr = ChunkUniquePtr>
 struct ChunkSliceTemplate {
     Ptr chunk;
+    size_t segment_id = 0;
     size_t offset = 0;
 
     bool empty() const;
     size_t rows() const;
     size_t skip(size_t skip_rows);
-    Ptr cutoff(size_t required_rows);
+    ChunkUniquePtr cutoff(size_t required_rows);
     void reset(Ptr input);
 };
 
@@ -548,5 +549,6 @@ APPLY_FOR_ALL_STRING_TYPE(GET_CONTAINER)
 
 using ChunkSlice = ChunkSliceTemplate<ChunkUniquePtr>;
 using ChunkSharedSlice = ChunkSliceTemplate<ChunkPtr>;
+using SegmentedChunkSlice = ChunkSliceTemplate<SegmentedChunkPtr>;
 
 } // namespace starrocks

be/src/column/vectorized_fwd.h

@@ -112,6 +112,12 @@ using ChunkPtr = std::shared_ptr<Chunk>;
 using ChunkUniquePtr = std::unique_ptr<Chunk>;
 using Chunks = std::vector<ChunkPtr>;
 
+class SegmentedColumn;
+class SegmentedChunk;
+using SegmentedColumnPtr = std::shared_ptr<SegmentedColumn>;
+using SegmentedColumns = std::vector<SegmentedColumnPtr>;
+using SegmentedChunkPtr = std::shared_ptr<SegmentedChunk>;
+
 using SchemaPtr = std::shared_ptr<Schema>;
 
 using Fields = std::vector<std::shared_ptr<Field>>;

be/src/exec/join_hash_map.cpp

@@ -335,7 +335,7 @@ void JoinHashTable::create(const HashTableParam& param) {
         _probe_state->probe_counter = param.probe_counter;
     }
 
-    _table_items->build_chunk = std::make_shared<Chunk>();
+    _table_items->build_chunk = std::make_shared<SegmentedChunk>(param.build_chunk_segment_size);
     _table_items->with_other_conjunct = param.with_other_conjunct;
     _table_items->join_type = param.join_type;
     _table_items->mor_reader_mode = param.mor_reader_mode;
@@ -545,15 +545,6 @@ Status JoinHashTable::build(RuntimeState* state) {
     RETURN_IF_ERROR(_table_items->build_chunk->upgrade_if_overflow());
     _table_items->has_large_column = _table_items->build_chunk->has_large_column();
 
-    // If the join key is column ref of build chunk, fetch from build chunk directly
-    size_t join_key_count = _table_items->join_keys.size();
-    for (size_t i = 0; i < join_key_count; i++) {
-        if (_table_items->join_keys[i].col_ref != nullptr) {
-            SlotId slot_id = _table_items->join_keys[i].col_ref->slot_id();
-            _table_items->key_columns[i] = _table_items->build_chunk->get_column_by_slot_id(slot_id);
-        }
-    }
-
     RETURN_IF_ERROR(_upgrade_key_columns_if_overflow());
 
     _hash_map_type = _choose_join_hash_map();
@@ -626,31 +617,22 @@ Status JoinHashTable::probe_remain(RuntimeState* state, ChunkPtr* chunk, bool* e
 }
 
 void JoinHashTable::append_chunk(const ChunkPtr& chunk, const Columns& key_columns) {
-    Columns& columns = _table_items->build_chunk->columns();
-
+    // TODO: simplify the SlotId mapping, if the slot of input chunk is same as build_chunk, we don't need to remap them
+    std::vector<SlotId> slots;
     for (size_t i = 0; i < _table_items->build_column_count; i++) {
         SlotDescriptor* slot = _table_items->build_slots[i].slot;
-        ColumnPtr& column = chunk->get_column_by_slot_id(slot->id());
-
-        if (!columns[i]->is_nullable() && column->is_nullable()) {
-            // upgrade to nullable column
-            columns[i] = NullableColumn::create(columns[i], NullColumn::create(columns[i]->size(), 0));
-        }
-        columns[i]->append(*column);
+        slots.push_back(slot->id());
     }
+    _table_items->build_chunk->append_chunk(chunk, slots);
 
     for (size_t i = 0; i < _table_items->key_columns.size(); i++) {
-        // If the join key is slot ref, will get from build chunk directly,
-        // otherwise will append from key_column of input
-        if (_table_items->join_keys[i].col_ref == nullptr) {
-            // upgrade to nullable column
-            if (!_table_items->key_columns[i]->is_nullable() && key_columns[i]->is_nullable()) {
-                size_t row_count = _table_items->key_columns[i]->size();
-                _table_items->key_columns[i] =
-                        NullableColumn::create(_table_items->key_columns[i], NullColumn::create(row_count, 0));
-            }
-            _table_items->key_columns[i]->append(*key_columns[i]);
+        // upgrade to nullable column
+        if (!_table_items->key_columns[i]->is_nullable() && key_columns[i]->is_nullable()) {
+            size_t row_count = _table_items->key_columns[i]->size();
+            _table_items->key_columns[i] =
+                    NullableColumn::create(_table_items->key_columns[i], NullColumn::create(row_count, 0));
         }
+        _table_items->key_columns[i]->append(*key_columns[i]);
     }
 
     _table_items->row_count += chunk->num_rows();
@@ -659,16 +641,16 @@ void JoinHashTable::append_chunk(const ChunkPtr& chunk, const Columns& key_colum
 void JoinHashTable::merge_ht(const JoinHashTable& ht) {
     _table_items->row_count += ht._table_items->row_count;
 
-    Columns& columns = _table_items->build_chunk->columns();
-    Columns& other_columns = ht._table_items->build_chunk->columns();
+    auto& columns = _table_items->build_chunk->columns();
+    auto& other_columns = ht._table_items->build_chunk->columns();
 
     for (size_t i = 0; i < _table_items->build_column_count; i++) {
         if (!columns[i]->is_nullable() && other_columns[i]->is_nullable()) {
             // upgrade to nullable column
-            columns[i] = NullableColumn::create(columns[i], NullColumn::create(columns[i]->size(), 0));
+            columns[i]->upgrade_to_nullable();
         }
-        columns[i]->append(*other_columns[i], 1, other_columns[i]->size() - 1);
     }
+    _table_items->build_chunk->append(ht._table_items->build_chunk, 1);
 }
 
 ChunkPtr JoinHashTable::convert_to_spill_schema(const ChunkPtr& chunk) const {

be/src/exec/join_hash_map.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "storage/chunk_helper.h"
 #include "util/runtime_profile.h"
 #define JOIN_HASH_MAP_H
 
@@ -99,8 +100,7 @@ struct HashTableSlotDescriptor {
 };
 
 struct JoinHashTableItems {
-    //TODO: memory continues problem?
-    ChunkPtr build_chunk = nullptr;
+    SegmentedChunkPtr build_chunk = nullptr;
     Columns key_columns;
     std::vector<HashTableSlotDescriptor> build_slots;
     std::vector<HashTableSlotDescriptor> probe_slots;
@@ -294,6 +294,9 @@ struct HashTableParam {
     RuntimeProfile::Counter* output_probe_column_timer = nullptr;
     RuntimeProfile::Counter* probe_counter = nullptr;
     bool mor_reader_mode = false;
+
+    // TODO: optimize this according to chunk width
+    size_t build_chunk_segment_size = 1 << 16;
 };
 
 template <class T>
@@ -688,9 +691,10 @@ class JoinHashMap {
 
     void _copy_probe_nullable_column(ColumnPtr* src_column, ChunkPtr* chunk, const SlotDescriptor* slot);
 
-    void _copy_build_column(const ColumnPtr& src_column, ChunkPtr* chunk, const SlotDescriptor* slot, bool to_nullable);
+    void _copy_build_column(const SegmentedColumnPtr& src_column, ChunkPtr* chunk, const SlotDescriptor* slot,
+                            bool to_nullable);
 
-    void _copy_build_nullable_column(const ColumnPtr& src_column, ChunkPtr* chunk, const SlotDescriptor* slot);
+    void _copy_build_nullable_column(const SegmentedColumnPtr& src_column, ChunkPtr* chunk, const SlotDescriptor* slot);
 
     void _probe_index_output(ChunkPtr* chunk);
     void _build_index_output(ChunkPtr* chunk);
@@ -838,7 +842,7 @@ class JoinHashTable {
     // convert input column to spill schema order
     ChunkPtr convert_to_spill_schema(const ChunkPtr& chunk) const;
 
-    const ChunkPtr& get_build_chunk() const { return _table_items->build_chunk; }
+    const SegmentedChunkPtr& get_build_chunk() const { return _table_items->build_chunk; }
     Columns& get_key_columns() { return _table_items->key_columns; }
     const Columns& get_key_columns() const { return _table_items->key_columns; }
     uint32_t get_row_count() const { return _table_items->row_count; }

be/src/exec/join_hash_map.tpp

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "column/vectorized_fwd.h"
 #include "simd/simd.h"
 #include "util/runtime_profile.h"
 
@@ -413,6 +414,7 @@ void JoinHashMap<LT, BuildFunc, ProbeFunc>::probe_prepare(RuntimeState* state) {
 
 template <LogicalType LT, class BuildFunc, class ProbeFunc>
 void JoinHashMap<LT, BuildFunc, ProbeFunc>::build(RuntimeState* state) {
+    _table_items->build_chunk->append_finished();
     BuildFunc().construct_hash_table(state, _table_items, _probe_state);
 }
 
@@ -619,7 +621,7 @@ void JoinHashMap<LT, BuildFunc, ProbeFunc>::_build_output(ChunkPtr* chunk) {
 
         bool need_output = is_lazy ? hash_table_slot.need_lazy_materialize : hash_table_slot.need_output;
         if (need_output) {
-            ColumnPtr& column = _table_items->build_chunk->columns()[i];
+            auto& column = _table_items->build_chunk->columns()[i];
             if (!column->is_nullable()) {
                 _copy_build_column(column, chunk, slot, to_nullable);
             } else {
@@ -684,11 +686,10 @@ void JoinHashMap<LT, BuildFunc, ProbeFunc>::_copy_probe_nullable_column(ColumnPt
 }
 
 template <LogicalType LT, class BuildFunc, class ProbeFunc>
-void JoinHashMap<LT, BuildFunc, ProbeFunc>::_copy_build_column(const ColumnPtr& src_column, ChunkPtr* chunk,
+void JoinHashMap<LT, BuildFunc, ProbeFunc>::_copy_build_column(const SegmentedColumnPtr& src_column, ChunkPtr* chunk,
                                                                const SlotDescriptor* slot, bool to_nullable) {
     if (to_nullable) {
-        auto data_column = src_column->clone_empty();
-        data_column->append_selective(*src_column, _probe_state->build_index.data(), 0, _probe_state->count);
+        auto data_column = src_column->clone_selective(_probe_state->build_index.data(), 0, _probe_state->count);
 
         // When left outer join is executed,
         // build_index[i] Equal to 0 means it is not found in the hash table,
@@ -704,18 +705,15 @@ void JoinHashMap<LT, BuildFunc, ProbeFunc>::_copy_build_column(const ColumnPtr&
         auto dest_column = NullableColumn::create(std::move(data_column), null_column);
         (*chunk)->append_column(std::move(dest_column), slot->id());
     } else {
-        auto dest_column = src_column->clone_empty();
-        dest_column->append_selective(*src_column, _probe_state->build_index.data(), 0, _probe_state->count);
-        (*chunk)->append_column(std::move(dest_column), slot->id());
+        auto data_column = src_column->clone_selective(_probe_state->build_index.data(), 0, _probe_state->count);
+        (*chunk)->append_column(std::move(data_column), slot->id());
     }
 }
 
 template <LogicalType LT, class BuildFunc, class ProbeFunc>
-void JoinHashMap<LT, BuildFunc, ProbeFunc>::_copy_build_nullable_column(const ColumnPtr& src_column, ChunkPtr* chunk,
-                                                                        const SlotDescriptor* slot) {
-    ColumnPtr dest_column = src_column->clone_empty();
-
-    dest_column->append_selective(*src_column, _probe_state->build_index.data(), 0, _probe_state->count);
+void JoinHashMap<LT, BuildFunc, ProbeFunc>::_copy_build_nullable_column(const SegmentedColumnPtr& src_column,
+                                                                        ChunkPtr* chunk, const SlotDescriptor* slot) {
+    ColumnPtr dest_column = src_column->clone_selective(_probe_state->build_index.data(), 0, _probe_state->count);
 
     // When left outer join is executed,
     // build_index[i] Equal to 0 means it is not found in the hash table,

be/src/exec/pipeline/hashjoin/spillable_hash_join_build_operator.cpp

@@ -259,7 +259,7 @@ StatusOr<std::function<StatusOr<ChunkPtr>()>> SpillableHashJoinBuildOperator::_c
             }
         }
 
-        auto chunk = _hash_table_build_chunk_slice.cutoff(runtime_state()->chunk_size());
+        ChunkPtr chunk = _hash_table_build_chunk_slice.cutoff(runtime_state()->chunk_size());
         RETURN_IF_ERROR(chunk->downgrade());
         RETURN_IF_ERROR(append_hash_columns(chunk));
         _join_builder->update_build_rows(chunk->num_rows());

be/src/exec/pipeline/hashjoin/spillable_hash_join_build_operator.h

@@ -66,7 +66,7 @@ class SpillableHashJoinBuildOperator final : public HashJoinBuildOperator {
 
     size_t _hash_table_iterate_idx = 0;
     std::vector<JoinHashTable*> _hash_tables;
-    ChunkSharedSlice _hash_table_build_chunk_slice;
+    SegmentedChunkSlice _hash_table_build_chunk_slice;
     std::function<StatusOr<ChunkPtr>()> _hash_table_slice_iterator;
     bool _is_first_time_spill = true;
     DECLARE_ONCE_DETECTOR(_set_finishing_once);

be/src/exec/spill/mem_table.cpp

@@ -167,7 +167,7 @@ Status OrderedMemTable::finalize(workgroup::YieldContext& yield_ctx, const Spill
             return Status::OK();
         }
         SCOPED_RAW_TIMER(&yield_ctx.time_spent_ns);
-        auto chunk = _chunk_slice.cutoff(_runtime_state->chunk_size());
+        ChunkPtr chunk = _chunk_slice.cutoff(_runtime_state->chunk_size());
         bool need_aligned = _runtime_state->spill_enable_direct_io();
 
         RETURN_IF_ERROR(serde->serialize(_runtime_state, serde_ctx, chunk, output, need_aligned));