This is an automated email from the ASF dual-hosted git repository.
HappenLee pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 5378480d492 [fix](be) Support complex array_agg state serialization
(#66601)
5378480d492 is described below
commit 5378480d49253aa8709739e1ee09766c67d53f0e
Author: HappenLee <[email protected]>
AuthorDate: Tue Aug 11 22:19:19 2026 +0800
[fix](be) Support complex array_agg state serialization (#66601)
Problem Summary: A two-phase `array_agg_foreach` query over nested
arrays in an ordinary OLAP table failed with `array_agg not support
write`. The generic complex-type `array_agg` state did not implement
serialization, deserialization, or merge, so it could not cross
aggregation phases or grow the ForEach state array. This change
serializes complex states with the native DataType format using the
active BE execution version and merges state columns by value. The query
now completes without using Table Stream.
### Release note
Fix `array_agg` and `array_agg_foreach` failures for complex element
types in multi-phase aggregation.
### Check List (For Author)
- Test: Unit Test and Regression test
- Unit Test: `GLIBC_COMPATIBILITY=OFF ./run-be-ut.sh -j 48 --run
--filter='AggregateFunctionArrayAggTest.*'` (7 tests passed)
- Regression test: `./run-regression-test.sh --run -d function_p0 -s
test_agg_foreach` (1 suite passed)
- ASAN BE build passed
- FE build and Checkstyle passed
- Behavior changed: Yes. Complex `array_agg` states now support
serialization and merge.
- Does this need documentation: No
---
.../exprs/aggregate/aggregate_function_array_agg.h | 68 +++++--
.../exprs/aggregate/aggregate_function_foreach.h | 19 +-
be/test/exprs/aggregate/agg_array_agg_test.cpp | 219 ++++++++++++++++++++-
.../aggregate_function_exception_test.cpp | 70 ++++++-
.../suites/function_p0/test_agg_foreach.groovy | 2 +-
5 files changed, 345 insertions(+), 33 deletions(-)
diff --git a/be/src/exprs/aggregate/aggregate_function_array_agg.h
b/be/src/exprs/aggregate/aggregate_function_array_agg.h
index f2c7a1feadb..002e11cdd13 100644
--- a/be/src/exprs/aggregate/aggregate_function_array_agg.h
+++ b/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); }
@@ -247,9 +249,7 @@ struct AggregateFunctionArrayAggData<T> {
const auto& to_nested_col = to_arr.get_data();
auto start = to_arr.get_offsets()[row_num - 1];
auto end = start + to_arr.get_offsets()[row_num] -
to_arr.get_offsets()[row_num - 1];
- for (auto i = start; i < end; ++i) {
- column_data->insert_from(to_nested_col, i);
- }
+ column_data->insert_range_from(to_nested_col, start, end - start);
}
void reset() { column_data->clear(); }
@@ -257,23 +257,42 @@ struct AggregateFunctionArrayAggData<T> {
void insert_result_into(IColumn& to) const {
auto& to_arr = assert_cast<ColumnArray&,
TypeCheckOnRelease::DISABLE>(to);
auto& to_nested_col = to_arr.get_data();
- size_t num_rows = column_data->size();
- for (size_t i = 0; i < num_rows; ++i) {
- to_nested_col.insert_from(*column_data, i);
- }
+ to_nested_col.insert_range_from(*column_data, 0, column_data->size());
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 +304,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 +325,20 @@ public:
}
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 {
diff --git a/be/src/exprs/aggregate/aggregate_function_foreach.h
b/be/src/exprs/aggregate/aggregate_function_foreach.h
index 28313dee8b6..d6324cd7694 100644
--- a/be/src/exprs/aggregate/aggregate_function_foreach.h
+++ b/be/src/exprs/aggregate/aggregate_function_foreach.h
@@ -97,24 +97,25 @@ protected:
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]);
}
diff --git a/be/test/exprs/aggregate/agg_array_agg_test.cpp
b/be/test/exprs/aggregate/agg_array_agg_test.cpp
index d65565a4c99..84a3d823f68 100644
--- a/be/test/exprs/aggregate/agg_array_agg_test.cpp
+++ b/be/test/exprs/aggregate/agg_array_agg_test.cpp
@@ -17,13 +17,14 @@
#include <gtest/gtest-message.h>
#include <gtest/gtest-test-part.h>
-#include <stddef.h>
-#include <stdint.h>
+#include <cstddef>
+#include <cstdint>
#include <memory>
#include <ostream>
#include <string>
+#include "agent/be_exec_version_manager.h"
#include "common/logging.h"
#include "core/arena.h"
#include "core/column/column_array.h"
@@ -34,16 +35,21 @@
#include "core/data_type/data_type_date.h"
#include "core/data_type/data_type_date_time.h"
#include "core/data_type/data_type_decimal.h"
+#include "core/data_type/data_type_map.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h"
#include "core/data_type/data_type_string.h"
+#include "core/data_type/data_type_struct.h"
#include "core/string_buffer.hpp"
#include "core/types.h"
#include "exprs/aggregate/agg_function_test.h"
#include "exprs/aggregate/aggregate_function.h"
+#include "exprs/aggregate/aggregate_function_array_agg.h"
#include "exprs/aggregate/aggregate_function_simple_factory.h"
#include "exprs/aggregate/aggregate_function_sort.h"
#include "gtest/gtest_pred_impl.h"
+#include "runtime/memory/mem_tracker_limiter.h"
+#include "runtime/thread_context.h"
namespace doris {
class IColumn;
@@ -53,6 +59,83 @@ namespace doris {
struct AggregateFunctionArrayAggTest : public AggregateFunctiontest {};
+namespace {
+
+Field array_field(std::initializer_list<Field> values) {
+ return Field::create_field<TYPE_ARRAY>(Array(values));
+}
+
+Field map_field(std::initializer_list<Field> keys,
std::initializer_list<Field> values) {
+ return Field::create_field<TYPE_MAP>(Map {array_field(keys),
array_field(values)});
+}
+
+void add_column_to_state(const IAggregateFunction& function, AggregateDataPtr
state,
+ const IColumn& column, Arena& arena) {
+ const IColumn* columns[] = {&column};
+ for (size_t row = 0; row < column.size(); ++row) {
+ function.add(state, columns, row, arena);
+ }
+}
+
+void check_complex_array_agg_state(const DataTypePtr& data_type,
+ std::initializer_list<Field> source_values,
+ std::initializer_list<Field> rhs_values) {
+ SCOPED_TRACE(data_type->get_name());
+ const auto nullable_type = make_nullable(data_type);
+ auto function = AggregateFunctionSimpleFactory::instance().get(
+ "array_agg", {nullable_type}, nullptr, false,
+ BeExecVersionManager::get_newest_version());
+ ASSERT_NE(function, nullptr);
+ function->set_version(BeExecVersionManager::get_newest_version());
+
+ auto source_column = nullable_type->create_column();
+ for (const auto& value : source_values) {
+ source_column->insert(value);
+ }
+ auto rhs_column = nullable_type->create_column();
+ for (const auto& value : rhs_values) {
+ rhs_column->insert(value);
+ }
+
+ Arena arena;
+ AggregateFunctionGuard source(function.get());
+ AggregateFunctionGuard restored(function.get());
+ AggregateFunctionGuard rhs(function.get());
+ add_column_to_state(*function, source.data(), *source_column, arena);
+ add_column_to_state(*function, rhs.data(), *rhs_column, arena);
+
+ auto serialized_column = ColumnString::create();
+ BufferWritable writer(*serialized_column);
+ function->serialize(source.data(), writer);
+ writer.commit();
+ ASSERT_EQ(serialized_column->size(), 1);
+
+ auto serialized_data = serialized_column->get_data_at(0);
+ BufferReadable reader(serialized_data);
+ function->deserialize(restored.data(), reader, arena);
+ function->merge(restored.data(), rhs.data(), arena);
+
+ auto result_column = function->get_return_type()->create_column();
+ function->insert_result_into(restored.data(), *result_column);
+ auto expected_column = function->get_return_type()->create_column();
+ Array expected_values(source_values);
+ expected_values.insert(expected_values.end(), rhs_values.begin(),
rhs_values.end());
+ expected_column->insert(Field::create_field<TYPE_ARRAY>(expected_values));
+ EXPECT_TRUE(ColumnHelper::column_equal(std::move(result_column),
std::move(expected_column)));
+}
+
+class ThrowOnSerializeDataType final : public DataTypeString {
+public:
+ char* serialize(const IColumn&, char* buf, int) const override {
+ *buf = 1;
+ throw Exception(ErrorCode::MEM_ALLOC_FAILED, "mock serialize
allocation failure");
+ }
+};
+
+static_assert(sizeof(AggregateFunctionArrayAggData<INVALID_TYPE>) ==
sizeof(MutableColumnPtr));
+
+} // namespace
+
TEST_F(AggregateFunctionArrayAggTest, test_array_agg_aint64) {
create_agg("array_agg", false, {std::make_shared<DataTypeInt64>()},
std::make_shared<DataTypeInt64>());
@@ -201,6 +284,138 @@ TEST_F(AggregateFunctionArrayAggTest,
test_array_agg_aint64_foreach) {
ColumnWithTypeAndName(std::move(array_array_column),
array_array_data_type, "column"));
}
+TEST_F(AggregateFunctionArrayAggTest,
complex_type_state_serialize_deserialize_and_merge) {
+ auto nullable_int = make_nullable(std::make_shared<DataTypeInt32>());
+ auto nullable_string = make_nullable(std::make_shared<DataTypeString>());
+
+ Array streamvbyte_values;
+ for (int32_t value = 0; value < 65; ++value) {
+ streamvbyte_values.emplace_back(Field::create_field<TYPE_INT>(value));
+ }
+
check_complex_array_agg_state(std::make_shared<DataTypeArray>(nullable_int),
+
{Field::create_field<TYPE_ARRAY>(std::move(streamvbyte_values))},
+ {});
+
+ check_complex_array_agg_state(
+ std::make_shared<DataTypeArray>(nullable_int),
+ {array_field({Field::create_field<TYPE_INT>(1), Field()}),
Field()},
+ {array_field({Field::create_field<TYPE_INT>(2),
Field::create_field<TYPE_INT>(3)})});
+
+ check_complex_array_agg_state(
+ std::make_shared<DataTypeStruct>(DataTypes {nullable_int,
nullable_string}),
+ {Field::create_field<TYPE_STRUCT>(Struct
{Field::create_field<TYPE_INT>(1),
+
Field::create_field<TYPE_STRING>("one")}),
+ Field()},
+ {Field::create_field<TYPE_STRUCT>(
+ Struct {Field(),
Field::create_field<TYPE_STRING>("two")})});
+
+
check_complex_array_agg_state(std::make_shared<DataTypeMap>(nullable_string,
nullable_int),
+
{map_field({Field::create_field<TYPE_STRING>("one"),
+
Field::create_field<TYPE_STRING>("null")},
+
{Field::create_field<TYPE_INT>(1), Field()}),
+ Field()},
+
{map_field({Field::create_field<TYPE_STRING>("two")},
+
{Field::create_field<TYPE_INT>(2)})});
+}
+
+TEST_F(AggregateFunctionArrayAggTest,
foreach_complex_type_state_growth_and_round_trip) {
+ auto nullable_int = make_nullable(std::make_shared<DataTypeInt32>());
+ auto nullable_inner_array =
make_nullable(std::make_shared<DataTypeArray>(nullable_int));
+ auto input_type = std::make_shared<DataTypeArray>(nullable_inner_array);
+ auto function = AggregateFunctionSimpleFactory::instance().get(
+ "array_agg_foreach", {input_type}, input_type, false,
+ BeExecVersionManager::get_newest_version(), {.is_foreach = true,
.column_names = {}});
+ ASSERT_NE(function, nullptr);
+ function->set_version(BeExecVersionManager::get_newest_version());
+
+ auto input_column = input_type->create_column();
+
input_column->insert(array_field({array_field({Field::create_field<TYPE_INT>(1)})}));
+ input_column->insert(array_field(
+ {array_field({Field::create_field<TYPE_INT>(2)}),
+ array_field({Field::create_field<TYPE_INT>(3),
Field::create_field<TYPE_INT>(4)}),
+ Field()}));
+
+ Arena arena;
+ AggregateFunctionGuard source(function.get());
+ AggregateFunctionGuard restored(function.get());
+ AggregateFunctionGuard merged(function.get());
+ add_column_to_state(*function, source.data(), *input_column, arena);
+
+ auto serialized_column = ColumnString::create();
+ BufferWritable writer(*serialized_column);
+ function->serialize(source.data(), writer);
+ writer.commit();
+
+ auto serialized_data = serialized_column->get_data_at(0);
+ BufferReadable reader(serialized_data);
+ function->deserialize(restored.data(), reader, arena);
+ function->merge(merged.data(), restored.data(), arena);
+
+ auto result_column = function->get_return_type()->create_column();
+ function->insert_result_into(merged.data(), *result_column);
+ auto expected_column = function->get_return_type()->create_column();
+ expected_column->insert(
+
array_field({array_field({array_field({Field::create_field<TYPE_INT>(1)}),
+
array_field({Field::create_field<TYPE_INT>(2)})}),
+
array_field({array_field({Field::create_field<TYPE_INT>(3),
+
Field::create_field<TYPE_INT>(4)})}),
+ array_field({Field()})}));
+ EXPECT_TRUE(ColumnHelper::column_equal(std::move(result_column),
std::move(expected_column)));
+}
+
+TEST_F(AggregateFunctionArrayAggTest,
complex_type_state_write_rolls_back_on_failure) {
+ auto data_type = std::make_shared<ThrowOnSerializeDataType>();
+ AggregateFunctionArrayAggData<INVALID_TYPE> data(DataTypes {data_type});
+ auto serialized_column = ColumnString::create();
+ BufferWritable writer(*serialized_column);
+ writer.write_c_string("prefix");
+
+ EXPECT_THROW(data.write(writer, *data_type,
BeExecVersionManager::get_newest_version()),
+ Exception);
+ EXPECT_EQ(serialized_column->get_chars().size(), 6);
+
+ writer.commit();
+ EXPECT_EQ(serialized_column->get_data_at(0).to_string(), "prefix");
+}
+
+TEST_F(AggregateFunctionArrayAggTest,
large_complex_type_state_serialization_is_tracked) {
+ constexpr size_t PAYLOAD_BYTES = 2 * 1024 * 1024;
+ auto nullable_string = make_nullable(std::make_shared<DataTypeString>());
+ auto nullable_inner_array =
make_nullable(std::make_shared<DataTypeArray>(nullable_string));
+ auto input_type = std::make_shared<DataTypeArray>(nullable_inner_array);
+ auto function = AggregateFunctionSimpleFactory::instance().get(
+ "array_agg_foreach", {input_type}, input_type, false,
+ BeExecVersionManager::get_newest_version(), {.is_foreach = true,
.column_names = {}});
+ ASSERT_NE(function, nullptr);
+ function->set_version(BeExecVersionManager::get_newest_version());
+
+ auto input_column = input_type->create_column();
+ input_column->insert(array_field(
+
{array_field({Field::create_field<TYPE_STRING>(String(PAYLOAD_BYTES, 'x'))})}));
+ Arena arena;
+ AggregateFunctionGuard state(function.get());
+ add_column_to_state(*function, state.data(), *input_column, arena);
+
+ auto tracker =
MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER,
+
"ArrayAggLargeStateSerializationTest");
+ auto switch_tracker = SwitchThreadMemTrackerLimiter(tracker);
+ thread_context()->thread_mem_tracker_mgr->flush_untracked_mem();
+ const auto baseline = tracker->consumption();
+ {
+ auto serialized_column = ColumnString::create();
+ BufferWritable writer(*serialized_column);
+ function->serialize(state.data(), writer);
+ writer.commit();
+ thread_context()->thread_mem_tracker_mgr->flush_untracked_mem();
+
+ const auto tracked_bytes = tracker->consumption() - baseline;
+ EXPECT_GT(tracked_bytes, PAYLOAD_BYTES);
+ EXPECT_GE(tracked_bytes, serialized_column->allocated_bytes());
+ }
+ thread_context()->thread_mem_tracker_mgr->flush_untracked_mem();
+ EXPECT_EQ(tracker->consumption(), baseline);
+}
+
TEST(AggregateFunctionSortDataTest, merge_does_not_share_rhs_block) {
auto data_type = std::make_shared<DataTypeInt64>();
Block prototype({ColumnWithTypeAndName(data_type->create_column(),
data_type, "value"),
diff --git a/be/test/exprs/aggregate/aggregate_function_exception_test.cpp
b/be/test/exprs/aggregate/aggregate_function_exception_test.cpp
index 21ee64dba4a..d571e8c0320 100644
--- a/be/test/exprs/aggregate/aggregate_function_exception_test.cpp
+++ b/be/test/exprs/aggregate/aggregate_function_exception_test.cpp
@@ -17,10 +17,14 @@
#include <gtest/gtest.h>
+#include <memory>
#include <vector>
#include "core/arena.h"
+#include "core/data_type/data_type_array.h"
+#include "core/data_type/data_type_string.h"
#include "exprs/aggregate/aggregate_function.h"
+#include "exprs/aggregate/aggregate_function_foreach.h"
namespace doris {
@@ -31,14 +35,17 @@ struct TrackingAggregateState {
static void reset_counters() {
construct_count = 0;
destroy_count = 0;
+ merge_count = 0;
}
static int construct_count;
static int destroy_count;
+ static int merge_count;
};
int TrackingAggregateState::construct_count = 0;
int TrackingAggregateState::destroy_count = 0;
+int TrackingAggregateState::merge_count = 0;
class ThrowOnDeserializeAggregateFunction final
: public IAggregateFunctionDataHelper<TrackingAggregateState,
@@ -73,6 +80,34 @@ public:
void insert_result_into(ConstAggregateDataPtr, IColumn&) const override {}
};
+class ThrowOnSecondMergeAggregateFunction final
+ : public IAggregateFunctionDataHelper<TrackingAggregateState,
+
ThrowOnSecondMergeAggregateFunction> {
+public:
+ ThrowOnSecondMergeAggregateFunction()
+ : IAggregateFunctionDataHelper<TrackingAggregateState,
+
ThrowOnSecondMergeAggregateFunction>(
+ DataTypes {std::make_shared<DataTypeString>()}) {}
+
+ String get_name() const override { return "throw_on_second_merge"; }
+
+ DataTypePtr get_return_type() const override { return
std::make_shared<DataTypeString>(); }
+
+ void add(AggregateDataPtr, const IColumn**, ssize_t, Arena&) const
override {}
+
+ void merge(AggregateDataPtr, ConstAggregateDataPtr, Arena&) const override
{
+ if (++TrackingAggregateState::merge_count == 2) {
+ throw Exception(ErrorCode::MEM_ALLOC_FAILED, "mock merge
allocation failure");
+ }
+ }
+
+ void serialize(ConstAggregateDataPtr, BufferWritable&) const override {}
+
+ void deserialize(AggregateDataPtr, BufferReadable&, Arena&) const override
{}
+
+ void insert_result_into(ConstAggregateDataPtr, IColumn&) const override {}
+};
+
class AggregateFunctionExceptionTest : public testing::Test {
protected:
void SetUp() override { TrackingAggregateState::reset_counters(); }
@@ -159,4 +194,37 @@ TEST_F(AggregateFunctionExceptionTest,
EXPECT_EQ(TrackingAggregateState::construct_count,
TrackingAggregateState::destroy_count);
}
-} // namespace doris
\ No newline at end of file
+TEST_F(AggregateFunctionExceptionTest,
ForEachGrowthPreservesOldStatesWhenMergeThrows) {
+ auto nested_function =
std::make_shared<ThrowOnSecondMergeAggregateFunction>();
+ auto input_type =
std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>());
+ AggregateFunctionForEach foreach_function(nested_function, DataTypes
{input_type});
+ auto input_column = input_type->create_column();
+ input_column->insert(
+ Field::create_field<TYPE_ARRAY>(Array
{Field::create_field<TYPE_STRING>(String("a")),
+
Field::create_field<TYPE_STRING>(String("b"))}));
+ input_column->insert(
+ Field::create_field<TYPE_ARRAY>(Array
{Field::create_field<TYPE_STRING>(String("c")),
+
Field::create_field<TYPE_STRING>(String("d")),
+
Field::create_field<TYPE_STRING>(String("e"))}));
+ const IColumn* columns[] = {input_column.get()};
+
+ {
+ AggregateFunctionGuard state(&foreach_function);
+ foreach_function.add(state.data(), columns, 0, arena);
+
+ try {
+ foreach_function.add(state.data(), columns, 1, arena);
+ FAIL() << "Expected doris::Exception";
+ } catch (const doris::Exception& e) {
+ EXPECT_EQ(e.code(), doris::ErrorCode::MEM_ALLOC_FAILED);
+ }
+
+ EXPECT_EQ(TrackingAggregateState::merge_count, 2);
+ EXPECT_EQ(TrackingAggregateState::construct_count, 5);
+ EXPECT_EQ(TrackingAggregateState::destroy_count, 3);
+ }
+
+ EXPECT_EQ(TrackingAggregateState::construct_count,
TrackingAggregateState::destroy_count);
+}
+
+} // namespace doris
diff --git a/regression-test/suites/function_p0/test_agg_foreach.groovy
b/regression-test/suites/function_p0/test_agg_foreach.groovy
index d2280e0c9ce..72fea2f4cfc 100644
--- a/regression-test/suites/function_p0/test_agg_foreach.groovy
+++ b/regression-test/suites/function_p0/test_agg_foreach.groovy
@@ -119,7 +119,7 @@ suite("test_agg_foreach") {
qt_sql """select array_agg_foreach(s) from foreach_table;"""
qt_array_agg_nested """
- select /*+ SET_VAR(parallel_pipeline_task_num=1) */
+ select /*+ SET_VAR(agg_phase=2, parallel_pipeline_task_num=1) */
size(array_agg_foreach(b)),
size(element_at(array_agg_foreach(b), 1)),
size(element_at(array_agg_foreach(b), 2)),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
