From 8afc5c20a051a1370229aa9279d29aaed8c928f4 Mon Sep 17 00:00:00 2001
From: Murphy <mofei@starrocks.com>
Date: Fri, 20 Sep 2024 19:05:27 +0800
Subject: [PATCH] fix

Signed-off-by: Murphy <mofei@starrocks.com>
---
 be/src/storage/chunk_helper.cpp | 32 +++++++++++++++++++++++++-------
 be/src/storage/chunk_helper.h   |  3 ++-
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/be/src/storage/chunk_helper.cpp b/be/src/storage/chunk_helper.cpp
index de291dc68e9ffb..5c2338c95c2232 100644
--- a/be/src/storage/chunk_helper.cpp
+++ b/be/src/storage/chunk_helper.cpp
@@ -608,6 +608,19 @@ bool ChunkPipelineAccumulator::is_finished() const {
     return _finalized && _out_chunk == nullptr && _in_chunk == nullptr;
 }
 
+template <class T>
+inline constexpr bool is_object = false;
+template <>
+inline constexpr bool is_object<ArrayColumn> = true;
+template <>
+inline constexpr bool is_object<StructColumn> = true;
+template <>
+inline constexpr bool is_object<MapColumn> = true;
+template <>
+inline constexpr bool is_object<JsonColumn> = true;
+template <class T>
+inline constexpr bool is_object<ObjectColumn<T>> = true;
+
 // Selective-copy data from SegmentedColumn according to provided index
 class SegmentedColumnSelectiveCopy final : public ColumnVisitorAdapter<SegmentedColumnSelectiveCopy> {
 public:
@@ -624,6 +637,7 @@ class SegmentedColumnSelectiveCopy final : public ColumnVisitorAdapter<Segmented
         using ColumnT = FixedLengthColumnBase<T>;
         using ContainerT = typename ColumnT::Container*;
 
+        _result = column.clone_empty();
         auto output = ColumnHelper::as_column<ColumnT>(_result);
         auto& columns = _segment_column->columns();
         size_t segment_size = _segment_column->segment_size();
@@ -646,6 +660,7 @@ class SegmentedColumnSelectiveCopy final : public ColumnVisitorAdapter<Segmented
         using ColumnT = BinaryColumnBase<T>;
         using ContainerT = typename ColumnT::Container*;
 
+        _result = column.clone_empty();
         auto output = ColumnHelper::as_column<ColumnT>(_result);
         auto& columns = _segment_column->columns();
         std::vector<ContainerT> buffers;
@@ -663,7 +678,8 @@ class SegmentedColumnSelectiveCopy final : public ColumnVisitorAdapter<Segmented
 
     // Fallback implementation, it's usually used for Array/Struct/Map/Json
     template <class ColumnT>
-    Status do_visit(const ColumnT& column) {
+    typename std::enable_if_t<is_object<ColumnT>, Status> do_visit(const ColumnT& column) {
+        _result = column.clone_empty();
         auto output = ColumnHelper::as_column<ColumnT>(_result);
         auto& columns = _segment_column->columns();
         for (uint32_t i = _from; i < _size; i++) {
@@ -675,7 +691,6 @@ class SegmentedColumnSelectiveCopy final : public ColumnVisitorAdapter<Segmented
     }
 
     Status do_visit(const NullableColumn& column) {
-        auto output = ColumnHelper::as_column<NullableColumn>(_result);
         std::vector<ColumnPtr> data_columns, null_columns;
         for (auto& column : _segment_column->columns()) {
             NullableColumn::Ptr nullable = ColumnHelper::as_column<NullableColumn>(column);
@@ -683,12 +698,13 @@ class SegmentedColumnSelectiveCopy final : public ColumnVisitorAdapter<Segmented
             null_columns.push_back(nullable->null_column());
         }
 
-        auto segmented_data_column = std::make_shared<SegmentedColumn>(data_columns);
+        auto segmented_data_column = std::make_shared<SegmentedColumn>(data_columns, _segment_column->segment_size());
         SegmentedColumnSelectiveCopy copy_data(segmented_data_column, _indexes, _from, _size);
         (void)data_columns[0]->accept(&copy_data);
-        auto segmented_null_column = std::make_shared<SegmentedColumn>(null_columns);
+        auto segmented_null_column = std::make_shared<SegmentedColumn>(null_columns, _segment_column->segment_size());
         SegmentedColumnSelectiveCopy copy_null(segmented_null_column, _indexes, _from, _size);
         (void)null_columns[0]->accept(&copy_null);
+        _result = NullableColumn::create(copy_data.result(), ColumnHelper::as_column<NullColumn>(copy_null.result()));
 
         return {};
     }
@@ -712,13 +728,15 @@ class SegmentedColumnSelectiveCopy final : public ColumnVisitorAdapter<Segmented
     uint32_t _size;
 };
 
-SegmentedColumn::SegmentedColumn(SegmentedChunkPtr chunk, size_t column_index) : _chunk(std::move(chunk)) {
+SegmentedColumn::SegmentedColumn(SegmentedChunkPtr chunk, size_t column_index)
+        : _chunk(std::move(chunk)), _segment_size(_chunk->segment_size()) {
     for (auto& segment : _chunk->segments()) {
         _columns.push_back(segment->get_column_by_index(column_index));
     }
 }
 
-SegmentedColumn::SegmentedColumn(std::vector<ColumnPtr> columns) : _columns(std::move(columns)) {}
+SegmentedColumn::SegmentedColumn(std::vector<ColumnPtr> columns, size_t segment_size)
+        : _columns(std::move(columns)), _segment_size(segment_size) {}
 
 ColumnPtr SegmentedColumn::clone_selective(const uint32_t* indexes, uint32_t from, uint32_t size) {
     SegmentedColumnSelectiveCopy visitor(shared_from_this(), indexes, from, size);
@@ -727,7 +745,7 @@ ColumnPtr SegmentedColumn::clone_selective(const uint32_t* indexes, uint32_t fro
 }
 
 size_t SegmentedColumn::segment_size() const {
-    return _chunk->segment_size();
+    return _segment_size;
 }
 
 size_t SegmentedChunk::segment_size() const {
diff --git a/be/src/storage/chunk_helper.h b/be/src/storage/chunk_helper.h
index a11646a9c6df89..28a5f3b5efe159 100644
--- a/be/src/storage/chunk_helper.h
+++ b/be/src/storage/chunk_helper.h
@@ -153,7 +153,7 @@ class ChunkPipelineAccumulator {
 class SegmentedColumn final : public std::enable_shared_from_this<SegmentedColumn> {
 public:
     SegmentedColumn(SegmentedChunkPtr chunk, size_t column_index);
-    SegmentedColumn(std::vector<ColumnPtr> columns);
+    SegmentedColumn(std::vector<ColumnPtr> columns, size_t segment_size);
     ~SegmentedColumn() = default;
 
     ColumnPtr clone_selective(const uint32_t* indexes, uint32_t from, uint32_t size);
@@ -168,6 +168,7 @@ class SegmentedColumn final : public std::enable_shared_from_this<SegmentedColum
 private:
     SegmentedChunkPtr _chunk;        // The chunk it belongs to
     std::vector<ColumnPtr> _columns; // All segmented columns
+    size_t _segment_size;
 };
 
 // A big-chunk would be segmented into multi small ones, to avoid allocating large-continuous memory