[Enhancement] split chunk of HashTable (#51175)

Signed-off-by: Murphy <[email protected]>

be/src/bench/shuffle_chunk_bench.cpp

@@ -16,14 +16,14 @@
 #include <testutil/assert.h>
 
 #include <memory>
+#include <random>
 
 #include "column/chunk.h"
 #include "column/column_helper.h"
-#include "column/datum_tuple.h"
-#include "common/config.h"
-#include "runtime/chunk_cursor.h"
-#include "runtime/runtime_state.h"
+#include "column/vectorized_fwd.h"
 #include "runtime/types.h"
+#include "storage/chunk_helper.h"
+#include "types/logical_type.h"
 
 namespace starrocks {
 
@@ -196,6 +196,181 @@ static void bench_func(benchmark::State& state) {
  perf.do_bench(state);
 }
 
+// Benchmark SegmentedColumn::clone_selective && Chunk::append_selective function
+class SegmentedChunkPerf {
+public:
+ SegmentedChunkPerf() = default;
+
+ void prepare_bench_segmented_chunk_clone(benchmark::State& state) {
+ // std::cerr << "chunk_size: " << _dest_chunk_size << std::endl;
+ // std::cerr << "segment_size: " << _segment_size << std::endl;
+ // std::cerr << "segmented_chunk_size: " << _segment_chunk_size << std::endl;
+ SegmentedChunkPtr seg_chunk = prepare_chunk();
+ CHECK_EQ(seg_chunk->num_rows(), _segment_chunk_size);
+
+ // random select
+ random_select(select, _dest_chunk_size, seg_chunk->num_rows());
+ }
+
+ void prepare_bench_chunk_clone(benchmark::State& state) {
+ ChunkPtr chunk = build_chunk(_segment_size);
+ CHECK_EQ(chunk->num_rows(), _segment_size);
+ random_select(select, _dest_chunk_size, chunk->num_rows());
+ }
+
+ void prepare(benchmark::State& state) {
+ state.PauseTiming();
+
+ _column_count = state.range(0);
+ _data_type = state.range(1);
+ _num_segments = state.range(2);
+ _types.clear();
+
+ prepare_bench_chunk_clone(state);
+ prepare_bench_segmented_chunk_clone(state);
+
+ state.ResumeTiming();
+ }
+
+ void do_bench_segmented_chunk_clone(benchmark::State& state) {
+ SegmentedChunkPtr seg_chunk = prepare_chunk();
+ // clone_selective
+ size_t items = 0;
+ for (auto _ : state) {
+ for (auto& column : seg_chunk->columns()) {
+ auto cloned = column->clone_selective(select.data(), 0, select.size());
+ }
+ items += select.size();
+ }
+ state.SetItemsProcessed(items);
+ }
+
+ void do_bench_chunk_clone(benchmark::State& state) {
+ ChunkPtr chunk = prepare_big_chunk();
+ size_t items = 0;
+ for (auto _ : state) {
+ ChunkPtr empty = chunk->clone_empty();
+ empty->append_selective(*chunk, select.data(), 0, select.size());
+ items += select.size();
+ }
+ state.SetItemsProcessed(items);
+ }
+
+ ChunkPtr prepare_big_chunk() {
+ if (_big_chunk) {
+ return _big_chunk;
+ }
+ _big_chunk = build_chunk(_segment_chunk_size);
+ return _big_chunk;
+ }
+
+ SegmentedChunkPtr prepare_chunk() {
+ if (_seg_chunk) {
+ return _seg_chunk;
+ }
+ ChunkPtr chunk = build_chunk(_dest_chunk_size);
+
+ for (int i = 0; i < (_segment_chunk_size / _dest_chunk_size); i++) {
+ if (!_seg_chunk) {
+ _seg_chunk = SegmentedChunk::create(_segment_size);
+ ChunkPtr chunk = build_chunk(_dest_chunk_size);
+ auto map = chunk->get_slot_id_to_index_map();
+ for (auto entry : map) {
+ _seg_chunk->append_column(chunk->get_column_by_slot_id(entry.first), entry.first);
+ }
+ _seg_chunk->build_columns();
+ } else {
+ // std::cerr << " append " << chunk->num_rows() << "rows, become " << _seg_chunk->num_rows() << std::endl;
+ _seg_chunk->append_chunk(chunk);
+ }
+ }
+ return _seg_chunk;
+ }
+
+ void random_select(std::vector<uint32_t>& select, size_t count, size_t range) {
+ select.resize(count);
+ std::random_device rd;
+ std::mt19937 gen(rd());
+ std::uniform_int_distribution<> dis(0, range - 1);
+ std::generate(select.begin(), select.end(), [&]() { return dis(gen); });
+ }
+
+ ChunkPtr build_chunk(size_t chunk_size) {
+ if (_types.empty()) {
+ for (int i = 0; i < _column_count; i++) {
+ if (_data_type == 0) {
+ _types.emplace_back(TypeDescriptor::create_varchar_type(128));
+ } else if (_data_type == 1) {
+ _types.emplace_back(LogicalType::TYPE_INT);
+ } else {
+ CHECK(false) << "data type not supported: " << _data_type;
+ }
+ }
+ }
+
+ auto chunk = std::make_unique<Chunk>();
+ for (int i = 0; i < _column_count; i++) {
+ auto col = init_dest_column(_types[i], chunk_size);
+ chunk->append_column(col, i);
+ }
+ return chunk;
+ }
+
+ ColumnPtr init_dest_column(const TypeDescriptor& type, size_t chunk_size) {
+ auto c1 = ColumnHelper::create_column(type, true);
+ c1->reserve(chunk_size);
+ for (int i = 0; i < chunk_size; i++) {
+ if (type.is_string_type()) {
+ std::string str = fmt::format("str{}", i);
+ c1->append_datum(Slice(str));
+ } else if (type.is_integer_type()) {
+ c1->append_datum(i);
+ } else {
+ CHECK(false) << "data type not supported";
+ }
+ }
+ return c1;
+ }
+
+private:
+ int _column_count = 4;
+ int _data_type = 0;
+ size_t _dest_chunk_size = 4096;
+ size_t _segment_size = 65536;
+ size_t _num_segments = 10;
+ size_t _segment_chunk_size = _segment_size * _num_segments;
+
+ SegmentedChunkPtr _seg_chunk;
+ ChunkPtr _big_chunk;
+ std::vector<uint32_t> select;
+ std::vector<TypeDescriptor> _types;
+};
+
+static void BenchSegmentedChunkClone(benchmark::State& state) {
+ google::InstallFailureSignalHandler();
+ auto perf = std::make_unique<SegmentedChunkPerf>();
+ perf->prepare(state);
+ perf->do_bench_segmented_chunk_clone(state);
+}
+
+static void BenchChunkClone(benchmark::State& state) {
+ google::InstallFailureSignalHandler();
+ auto perf = std::make_unique<SegmentedChunkPerf>();
+ perf->prepare(state);
+ perf->do_bench_chunk_clone(state);
+}
+
+static std::vector<std::vector<int64_t>> chunk_clone_args() {
+ return {
+ {1, 2, 3, 4}, // num columns
+ {0, 1}, // data type
+ {1, 4, 16, 64} // num_segments
+ };
+}
+
+BENCHMARK(BenchSegmentedChunkClone)->ArgsProduct(chunk_clone_args());
+BENCHMARK(BenchChunkClone)->ArgsProduct(chunk_clone_args());
+
 static void process_args(benchmark::internal::Benchmark* b) {
  // chunk_count, column_count, node_count, src_chunk_size, null percent
  b->Args({400, 400, 140, 4096, 80});

be/src/column/binary_column.cpp

@@ -46,8 +46,16 @@ void BinaryColumnBase<T>::check_or_die() const {
  }
 }
 
+template <typename T>
+void BinaryColumnBase<T>::append(const Slice& str) {
+ _bytes.insert(_bytes.end(), str.data, str.data + str.size);
+ _offsets.emplace_back(_bytes.size());
+ _slices_cache = false;
+}
+
 template <typename T>
 void BinaryColumnBase<T>::append(const Column& src, size_t offset, size_t count) {
+ DCHECK(offset + count <= src.size());
  const auto& b = down_cast<const BinaryColumnBase<T>&>(src);
  const unsigned char* p = &b._bytes[b._offsets[offset]];
  const unsigned char* e = &b._bytes[b._offsets[offset + count]];

be/src/column/binary_column.h

@@ -32,7 +32,7 @@ class BinaryColumnBase final : public ColumnFactory<Column, BinaryColumnBase<T>>
 
  using Offset = T;
  using Offsets = Buffer<T>;
-
+ using Byte = uint8_t;
  using Bytes = starrocks::raw::RawVectorPad16<uint8_t, ColumnAllocator<uint8_t>>;
 
  struct BinaryDataProxyContainer {
@@ -172,11 +172,7 @@ class BinaryColumnBase final : public ColumnFactory<Column, BinaryColumnBase<T>>
  // No complain about the overloaded-virtual for this function
  DIAGNOSTIC_PUSH
  DIAGNOSTIC_IGNORE("-Woverloaded-virtual")
- void append(const Slice& str) {
- _bytes.insert(_bytes.end(), str.data, str.data + str.size);
- _offsets.emplace_back(_bytes.size());
- _slices_cache = false;
- }
+ void append(const Slice& str);
  DIAGNOSTIC_POP
 
  void append_datum(const Datum& datum) override {

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"
@@ -465,7 +466,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);
@@ -478,7 +479,31 @@ 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
+ segment_id = offset / chunk->segment_size();
+
+ 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

@@ -536,12 +536,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);
 };
 
@@ -571,5 +572,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/const_column.h

@@ -26,6 +26,8 @@ class ConstColumn final : public ColumnFactory<Column, ConstColumn> {
  friend class ColumnFactory<Column, ConstColumn>;
 
 public:
+ using ValueType = void;
+
  explicit ConstColumn(ColumnPtr data_column);
  ConstColumn(ColumnPtr data_column, size_t size);
 

be/src/column/nullable_column.h

@@ -33,6 +33,8 @@ class NullableColumn : public ColumnFactory<Column, NullableColumn> {
  friend class ColumnFactory<Column, NullableColumn>;
 
 public:
+ using ValueType = void;
+
  inline static ColumnPtr wrap_if_necessary(ColumnPtr column) {
  if (column->is_nullable()) {
  return column;

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>>;