[fix](be) Support complex array_agg state serialization

be/src/exprs/aggregate/aggregate_function_array_agg.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include "common/check.h"
 #include "core/assert_cast.h"
 #include "core/column/column.h"
 #include "core/column/column_array.h"
@@ -37,6 +38,7 @@ class Arena;
 template <PrimitiveType T>
 struct AggregateFunctionArrayAggData {
     static constexpr PrimitiveType PType = T;
+    static constexpr bool use_native_serde = false;
     using ElementType = typename PrimitiveTypeTraits<T>::CppType;
     using ColVecType = typename PrimitiveTypeTraits<T>::ColumnType;
     using Self = AggregateFunctionArrayAggData<T>;
@@ -136,6 +138,7 @@ template <PrimitiveType T>
     requires(is_string_type(T))
 struct AggregateFunctionArrayAggData<T> {
     static constexpr PrimitiveType PType = T;
+    static constexpr bool use_native_serde = false;
     using ElementType = StringRef;
     using ColVecType = ColumnString;
     using Self = AggregateFunctionArrayAggData<T>;
@@ -231,14 +234,13 @@ template <PrimitiveType T>
              !is_date_type(T) && !is_ip(T))
 struct AggregateFunctionArrayAggData<T> {
     static constexpr PrimitiveType PType = T;
+    static constexpr bool use_native_serde = true;
     using ElementType = StringRef;
     using Self = AggregateFunctionArrayAggData<T>;
     MutableColumnPtr column_data;
 
-    AggregateFunctionArrayAggData(const DataTypes& argument_types) {
-        DataTypePtr column_type = argument_types[0];
-        column_data = column_type->create_column();
-    }
+    AggregateFunctionArrayAggData(const DataTypes& argument_types)
+            : column_data(argument_types[0]->create_column()) {}
 
     void add(const IColumn& column, size_t row_num) { column_data->insert_from(column, row_num); }
 
@@ -264,16 +266,38 @@ struct AggregateFunctionArrayAggData<T> {
         to_arr.get_offsets().push_back(to_nested_col.size());
     }
 
-    void write(BufferWritable& buf) const {
-        throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, "array_agg not support write");
+    void write(BufferWritable& buf, const IDataType& column_type, int be_exec_version) const {
+        const auto max_serialized_bytes = cast_set<size_t>(
+                column_type.get_uncompressed_serialized_bytes(*column_data, be_exec_version));
+        buf.resize(sizeof(UInt64) + max_serialized_bytes);
+        auto* serialized_data = buf.data() + sizeof(UInt64);
+        const char* end = nullptr;
+        try {
+            end = column_type.serialize(*column_data, serialized_data, be_exec_version);
+        } catch (...) {
+            buf.resize(0);
+            throw;
+        }
+        DORIS_CHECK_LE(end, serialized_data + max_serialized_bytes);
+        const auto serialized_bytes = static_cast<UInt64>(end - serialized_data);
+        memcpy(buf.data(), &serialized_bytes, sizeof(serialized_bytes));
+        const auto frame_bytes = sizeof(serialized_bytes) + cast_set<size_t>(serialized_bytes);
+        buf.resize(frame_bytes);
+        buf.add_offset(frame_bytes);
     }
 
-    void read(BufferReadable& buf) {
-        throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, "array_agg not support read");
+    void read(BufferReadable& buf, const IDataType& column_type, int be_exec_version) {
+        DORIS_CHECK(column_data->empty());
+        UInt64 serialized_bytes = 0;
+        buf.read_binary(serialized_bytes);
+        const auto* serialized_data = buf.data();
+        const auto* end = column_type.deserialize(serialized_data, &column_data, be_exec_version);
+        DORIS_CHECK_EQ(end, serialized_data + serialized_bytes);
+        buf.add_offset(serialized_bytes);
     }
 
     void merge(const Self& rhs) {
-        throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, "array_agg not support merge");
+        column_data->insert_range_from(*rhs.column_data, 0, rhs.column_data->size());
     }
 };
 
@@ -285,9 +309,10 @@ class AggregateFunctionArrayAgg final
           UnaryExpression,
           NotNullableAggregateFunction {
 public:
+    using Base = IAggregateFunctionDataHelper<Data, AggregateFunctionArrayAgg<Data>, true>;
+
     AggregateFunctionArrayAgg(const DataTypes& argument_types_)
-            : IAggregateFunctionDataHelper<Data, AggregateFunctionArrayAgg<Data>, true>(
-                      {argument_types_}),
+            : Base({argument_types_}),
               return_type(std::make_shared<DataTypeArray>(make_nullable(argument_types_[0]))) {}
 
     std::string get_name() const override { return "array_agg"; }
@@ -305,12 +330,20 @@ class AggregateFunctionArrayAgg final
     }
 
     void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override {
-        this->data(place).write(buf);
+        if constexpr (Data::use_native_serde) {
+            this->data(place).write(buf, *this->argument_types[0], this->version);
+        } else {
+            this->data(place).write(buf);
+        }
     }
 
     void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,
                      Arena&) const override {
-        this->data(place).read(buf);
+        if constexpr (Data::use_native_serde) {
+            this->data(place).read(buf, *this->argument_types[0], this->version);
+        } else {
+            this->data(place).read(buf);
+        }
     }
 
     void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {

be/src/exprs/aggregate/aggregate_function_foreach.h

@@ -97,24 +97,25 @@ class AggregateFunctionForEach : public AggregateFunctionNonFinalBase,
             char* new_state =
                     arena.aligned_alloc(allocation_size, nested_function->align_of_data());
 
-            size_t i;
+            size_t num_created = 0;
             try {
-                for (i = 0; i < new_size; ++i) {
-                    nested_function->create(&new_state[i * nested_size_of_data]);
+                for (; num_created < new_size; ++num_created) {
+                    nested_function->create(&new_state[num_created * nested_size_of_data]);
                 }
-            } catch (...) {
-                size_t cleanup_size = i;
 
-                for (i = 0; i < cleanup_size; ++i) {
+                for (size_t i = 0; i < old_size; ++i) {
+                    nested_function->merge(&new_state[i * nested_size_of_data],
+                                           &old_state[i * nested_size_of_data], arena);
+                }
+            } catch (...) {
+                for (size_t i = 0; i < num_created; ++i) {
                     nested_function->destroy(&new_state[i * nested_size_of_data]);
                 }
 
                 throw;
             }
 
-            for (i = 0; i < old_size; ++i) {
-                nested_function->merge(&new_state[i * nested_size_of_data],
-                                       &old_state[i * nested_size_of_data], arena);
+            for (size_t i = 0; i < old_size; ++i) {
                 nested_function->destroy(&old_state[i * nested_size_of_data]);
             }