From 094f6499cbcb58acfaf99f1d0f942454c871b9eb Mon Sep 17 00:00:00 2001 From: happenlee Date: Wed, 9 Sep 2026 01:32:00 +0800 Subject: [PATCH] [fix](be) Handle zero-argument count aggregate states ### What problem does this PR solve? Issue Number: N/A Related PR: #58031 Problem Summary: `count_union(count_state())` can crash the BE while initializing an aggregate evaluator. FE represents zero-argument count as `count(*)` and sends an AggState with an empty subtype list, but `DataTypeAggState` unconditionally reads the first subtype to infer the nested result type. Handle zero-argument count with its explicit Int64 result type while preserving the empty argument signature and reusing the existing count implementation. Reject unsupported zero-argument AggState functions with `INVALID_ARGUMENT` before invoking creators that require an input type. Parameterized aggregate handling and serialization formats remain unchanged. ### Release note Fix a BE crash when constructing zero-argument count aggregate states, including queries using `count_union(count_state())` and `count_merge(count_state())`. ### Check List (For Author) - Test: - [x] Unit Test: 20 ASAN tests passed (`DataTypeAggStateZeroArgumentTest.*`, `Params/DataTypeAggStateTest.*`, `AggregateFunctionCountTest.*`). New coverage includes empty/one/three/8193-row state serialization and merge, plus invalid empty-argument functions. - [x] Regression test: `test_count_state_zero_arguments` passed on an isolated ASAN BE cluster. Generated expected output with `-forceGenOut`, then passed a normal comparison run. Covers the original crash, empty input, nullable input, 10001 rows, grouped union, and parameterized Decimal aggregates. BE remained alive afterward. - BE and FE build passed with `./build.sh --be --fe -j 48`; clang-format 16 and build hygiene checks passed. - clang-tidy reported no diagnostics on changed lines. The script could not pass because the base branch has an unmatched `NOLINTEND` in `be/src/core/types.h`; other emitted diagnostics also refer to unchanged code. - Behavior changed: - [x] Yes. Zero-argument count AggState construction succeeds instead of crashing; unsupported empty-argument states raise an error. - Does this need documentation? - [x] No. --- be/src/core/data_type/data_type_agg_state.h | 14 +++- .../data_type/data_type_agg_state_test.cpp | 58 +++++++++++++- .../test_count_state_zero_arguments.out | 28 +++++++ .../test_count_state_zero_arguments.groovy | 78 +++++++++++++++++++ 4 files changed, 175 insertions(+), 3 deletions(-) create mode 100644 regression-test/data/datatype_p0/agg_state/test_count_state_zero_arguments.out create mode 100644 regression-test/suites/datatype_p0/agg_state/test_count_state_zero_arguments.groovy diff --git a/be/src/core/data_type/data_type_agg_state.h b/be/src/core/data_type/data_type_agg_state.h index 378315f9512e84..2864a94d4c2db0 100644 --- a/be/src/core/data_type/data_type_agg_state.h +++ b/be/src/core/data_type/data_type_agg_state.h @@ -26,6 +26,7 @@ #include "core/data_type/data_type.h" #include "core/data_type/data_type_factory.hpp" #include "core/data_type/data_type_fixed_length_object.h" +#include "core/data_type/data_type_number.h" #include "core/data_type/data_type_string.h" #include "core/data_type/define_primitive_type.h" #include "core/data_type_serde/data_type_string_serde.h" @@ -43,8 +44,17 @@ class DataTypeAggState : public DataTypeString { _function_name(std::move(function_name)), _be_exec_version(be_exec_version) { DataTypePtr result_type; - auto arg_primitive_type = _sub_types[0]->get_primitive_type(); - if (is_decimalv3(arg_primitive_type)) { + if (_sub_types.empty()) { + // count(*) has no input type from which to infer its result type. + if (_function_name != "count") { + throw Exception( + ErrorCode::INVALID_ARGUMENT, + "Aggregate function {} does not support AggState with zero arguments", + _function_name); + } + result_type = std::make_shared(); + } else if (auto arg_primitive_type = _sub_types[0]->get_primitive_type(); + is_decimalv3(arg_primitive_type)) { // TODO: handle decimal256 correctly according to session var enable_decimal256 int precision = 0; if (arg_primitive_type == PrimitiveType::TYPE_DECIMAL256) { diff --git a/be/test/core/data_type/data_type_agg_state_test.cpp b/be/test/core/data_type/data_type_agg_state_test.cpp index 45920202ae23b1..71c5f845a30a03 100644 --- a/be/test/core/data_type/data_type_agg_state_test.cpp +++ b/be/test/core/data_type/data_type_agg_state_test.cpp @@ -25,9 +25,11 @@ #include #include "agent/be_exec_version_manager.h" +#include "core/arena.h" #include "core/assert_cast.h" #include "core/column/column.h" #include "core/column/column_fixed_length_object.h" +#include "core/custom_allocator.h" #include "core/data_type/common_data_type_serder_test.h" #include "core/data_type/common_data_type_test.h" #include "core/data_type/data_type.h" @@ -38,6 +40,7 @@ #include "core/field.h" #include "core/types.h" #include "exec/common/variant_util.h" +#include "exprs/aggregate/aggregate_function_state_merge.h" // 1. datatype meta info: // get_type_id, get_type_as_type_descriptor, get_storage_field_type, have_subtypes, get_pdata_type (const IDataType *data_type), to_pb_column_meta (PColumnMeta *col_meta) @@ -253,6 +256,59 @@ TEST_P(DataTypeAggStateTest, SerializeDeserializeTest2) { std::cout << "finish serialize deserialize test2" << std::endl; } +TEST(DataTypeAggStateZeroArgumentTest, CountSerializeAndMerge) { + const int version = BeExecVersionManager::get_newest_version(); + auto state_type = std::make_shared(DataTypes {}, false, "count", version); + auto count_function = state_type->get_nested_function(); + EXPECT_TRUE(state_type->get_sub_types().empty()); + EXPECT_TRUE(count_function->get_argument_types().empty()); + EXPECT_EQ(count_function->get_return_type()->get_primitive_type(), TYPE_BIGINT); + + for (size_t rows : {0, 1, 3, 8193}) { + SCOPED_TRACE(rows); + Arena arena; + auto states = state_type->create_column(); + count_function->streaming_agg_serialize_to_column(nullptr, states, rows, arena); + ASSERT_EQ(states->size(), rows); + + DorisVector buffer(state_type->get_uncompressed_serialized_bytes(*states, version)); + auto* end = state_type->serialize(*states, buffer.data(), version); + auto restored = state_type->create_column(); + EXPECT_EQ(state_type->deserialize(buffer.data(), &restored, version), end); + ASSERT_EQ(restored->size(), rows); + + auto merge_function = AggregateStateMerge::create(count_function, DataTypes {state_type}, + count_function->get_return_type()); + auto* place = + reinterpret_cast(arena.alloc(merge_function->size_of_data())); + merge_function->create(place); + const IColumn* columns[] = {restored.get()}; + for (size_t row = 0; row < rows; ++row) { + merge_function->add(place, columns, row, arena); + } + auto result = count_function->get_return_type()->create_column(); + merge_function->insert_result_into(place, *result); + merge_function->destroy(place); + ASSERT_EQ(result->size(), 1); + EXPECT_EQ(result->get_int(0), rows); + } +} + +TEST(DataTypeAggStateZeroArgumentTest, RejectOtherFunctions) { + for (const auto* name : {"sum", "avg", "unknown"}) { + SCOPED_TRACE(name); + try { + DataTypeAggState state_type(DataTypes {}, true, name, + BeExecVersionManager::get_newest_version()); + FAIL() << "Expected an invalid-argument exception"; + } catch (const Exception& e) { + EXPECT_EQ(e.code(), ErrorCode::INVALID_ARGUMENT); + EXPECT_NE(e.to_string().find("does not support AggState with zero arguments"), + std::string::npos); + } + } +} + INSTANTIATE_TEST_SUITE_P(Params, DataTypeAggStateTest, ::testing::Values(0, 1, 31)); -} // namespace doris \ No newline at end of file +} // namespace doris diff --git a/regression-test/data/datatype_p0/agg_state/test_count_state_zero_arguments.out b/regression-test/data/datatype_p0/agg_state/test_count_state_zero_arguments.out new file mode 100644 index 00000000000000..c862fd67cf1237 --- /dev/null +++ b/regression-test/data/datatype_p0/agg_state/test_count_state_zero_arguments.out @@ -0,0 +1,28 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !single_row -- +1 + +-- !three_rows -- +3 + +-- !union_merge -- +3 + +-- !empty_input -- +0 + +-- !empty_union -- +0 + +-- !nullable_input -- +3 2 3 + +-- !multiple_batches -- +10001 + +-- !grouped_union -- +10001 + +-- !decimal_arguments -- +45.00 4.5000 + diff --git a/regression-test/suites/datatype_p0/agg_state/test_count_state_zero_arguments.groovy b/regression-test/suites/datatype_p0/agg_state/test_count_state_zero_arguments.groovy new file mode 100644 index 00000000000000..5891daa1c67c11 --- /dev/null +++ b/regression-test/suites/datatype_p0/agg_state/test_count_state_zero_arguments.groovy @@ -0,0 +1,78 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +suite("test_count_state_zero_arguments") { + sql "set enable_agg_state=true" + + // Exercise the original crash path, whose result is an opaque aggregate state. + sql """ + SELECT count_union(count_state()) + FROM (SELECT 1 AS v UNION ALL SELECT 2 AS v UNION ALL SELECT 3 AS v) t + """ + + order_qt_single_row "SELECT count_merge(count_state())" + + order_qt_three_rows """ + SELECT count_merge(count_state()) + FROM (SELECT 1 AS v UNION ALL SELECT 2 AS v UNION ALL SELECT 3 AS v) t + """ + + order_qt_union_merge """ + SELECT count_merge(s) + FROM ( + SELECT count_union(count_state()) AS s + FROM (SELECT 1 AS v UNION ALL SELECT 2 AS v UNION ALL SELECT 3 AS v) t + ) u + """ + + order_qt_empty_input """ + SELECT count_merge(count_state()) + FROM numbers("number"="10") WHERE number < 0 + """ + + order_qt_empty_union """ + SELECT count_merge(s) + FROM ( + SELECT count_union(count_state()) AS s + FROM numbers("number"="10") WHERE number < 0 + ) t + """ + + order_qt_nullable_input """ + SELECT count_merge(count_state()), count_merge(count_state(v)), + count_merge(count_state(1)) + FROM (SELECT CAST(NULL AS INT) AS v UNION ALL SELECT 1 UNION ALL SELECT 2) t + """ + + order_qt_multiple_batches """ + SELECT count_merge(count_state()) FROM numbers("number"="10001") + """ + + order_qt_grouped_union """ + SELECT count_merge(s) + FROM ( + SELECT number % 17 AS k, count_union(count_state()) AS s + FROM numbers("number"="10001") GROUP BY k + ) t + """ + + order_qt_decimal_arguments """ + SELECT sum_merge(sum_state(CAST(number AS DECIMAL(18, 2)))), + avg_merge(avg_state(CAST(number AS DECIMAL(18, 2)))) + FROM numbers("number"="10") + """ +}