[fix](be) Handle zero-argument count aggregate states

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<DataTypeInt64>();
+ } 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) {

be/test/core/data_type/data_type_agg_state_test.cpp

@@ -25,9 +25,11 @@
 #include <memory>
 
 #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<DataTypeAggState>(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<char> 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<AggregateDataPtr>(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
+} // namespace doris

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
+

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")
+ """
+}