[feature](be) Add scalar aggregate state finalize functions

be/src/exprs/function/function_agg_state_finalize.h

@@ -0,0 +1,93 @@
+// 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.
+
+#pragma once
+
+#include "core/arena.h"
+#include "core/block/block.h"
+#include "core/column/column_nullable.h"
+#include "exprs/aggregate/aggregate_function.h"
+#include "exprs/function/function.h"
+#include "util/defer_op.h"
+
+namespace doris {
+
+class FunctionAggStateFinalize : public IFunction {
+public:
+    FunctionAggStateFinalize(DataTypePtr return_type, AggregateFunctionPtr agg_function)
+            : _return_type(std::move(return_type)), _agg_function(std::move(agg_function)) {}
+
+    static FunctionBasePtr create(const DataTypes& argument_types, const DataTypePtr& return_type,
+                                  const AggregateFunctionPtr& agg_function) {
+        return std::make_shared<DefaultFunction>(
+                std::make_shared<FunctionAggStateFinalize>(return_type, agg_function),
+                argument_types, return_type);
+    }
+
+    String get_name() const override { return _agg_function->get_name() + "_finalize"; }
+
+    size_t get_number_of_arguments() const override { return 1; }
+
+    // An outer NULL can have an empty, invalid serialized payload. Do not deserialize it.
+    bool use_default_implementation_for_nulls() const override { return false; }
+
+    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
+        return _return_type;
+    }
+
+    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
+                        uint32_t result, size_t input_rows_count) const override {
+        const auto input =
+                block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
+        const auto* nullable = check_and_get_column<ColumnNullable>(*input);
+        const auto& states = nullable ? nullable->get_nested_column() : *input;
+        auto output = _agg_function->get_return_type()->create_column();
+        _agg_function->check_result_column_type(*output);
+        Arena arena;
+        const auto state_size = _agg_function->size_of_data();
+        const auto state_alignment = _agg_function->align_of_data();
+        for (size_t row = 0; row < input_rows_count; ++row) {
+            if (nullable && nullable->is_null_at(row)) {
+                output->insert_default();
+                continue;
+            }
+            {
+                auto* place = arena.aligned_alloc(state_size, state_alignment);
+                _agg_function->create(place);
+                DEFER(_agg_function->destroy(place));
+                // The serialized column can be numeric, fixed-length, string or a complex column.
+                _agg_function->deserialize_and_merge_from_column_range(place, states, row, row,
+                                                                       arena);
+                _agg_function->insert_result_into(place, *output);
+            }
+            // States may own variable-length data. Destroy them before reclaiming their arena.
+            arena.clear();
+        }
+        ColumnPtr result_column = std::move(output);
+        if (_return_type->is_nullable()) {
+            result_column = wrap_in_nullable(result_column, block, arguments, input_rows_count);
+        }
+        block.replace_by_position(result, std::move(result_column));
+        return Status::OK();
+    }
+
+private:
+    DataTypePtr _return_type;
+    AggregateFunctionPtr _agg_function;
+};
+
+} // namespace doris

be/src/exprs/vectorized_fn_call.cpp

@@ -49,6 +49,7 @@
 #include "exec/pipeline/pipeline_task.h"
 #include "exprs/function/array/function_array_distance.h"
 #include "exprs/function/function_agg_state.h"
+#include "exprs/function/function_agg_state_finalize.h"
 #include "exprs/function/function_fake.h"
 #include "exprs/function/function_java_udf.h"
 #include "exprs/function/function_python_udf.h"
@@ -79,6 +80,7 @@ class TExprNode;
 namespace doris {
 
 const std::string AGG_STATE_SUFFIX = "_state";
+const std::string AGG_FINALIZE_SUFFIX = "_finalize";
 
 // Now left child is a function call, we need to check if it is a distance function
 const static std::set<std::string> DISTANCE_FUNCS = {L2DistanceApproximate::name,
@@ -215,6 +217,33 @@ size_t raw_comparison_value_size(PrimitiveType primitive_type) {
     }
 }
 
+Status create_agg_state_finalize_function(const std::string& function_name,
+                                          const DataTypes& argument_types,
+                                          const DataTypePtr& return_type,
+                                          FunctionBasePtr& function) {
+    if (argument_types.size() != 1 ||
+        remove_nullable(argument_types[0])->get_primitive_type() != TYPE_AGG_STATE) {
+        return Status::InternalError("Finalize function requires one AGG_STATE argument");
+    }
+    const auto state_type = remove_nullable(argument_types[0]);
+    const auto* agg_state = assert_cast<const DataTypeAggState*>(state_type.get());
+    if (agg_state->get_function_name() + AGG_FINALIZE_SUFFIX != function_name) {
+        return Status::InternalError("{} does not match function {}", state_type->get_name(),
+                                     function_name);
+    }
+    const auto& nested = agg_state->get_nested_function();
+    auto expected_type = nested->get_return_type();
+    if (argument_types[0]->is_nullable()) {
+        expected_type = make_nullable(expected_type);
+    }
+    if (!expected_type->equals(*return_type)) {
+        return Status::InternalError("{} expects return type {}, but got {}", function_name,
+                                     expected_type->get_name(), return_type->get_name());
+    }
+    function = FunctionAggStateFinalize::create(argument_types, return_type, nested);
+    return Status::OK();
+}
+
 } // namespace
 
 VectorizedFnCall::VectorizedFnCall(const TExprNode& node) : VExpr(node) {
@@ -293,8 +322,11 @@ Status VectorizedFnCall::prepare(RuntimeState* state, const RowDescriptor& desc,
             _function = FunctionAggState::create(
                     argument_types, _data_type,
                     assert_cast<const DataTypeAggState*>(_data_type.get())->get_nested_function());
+        } else if (match_suffix(_fn.name.function_name, AGG_FINALIZE_SUFFIX)) {
+            RETURN_IF_ERROR(create_agg_state_finalize_function(
+                    _fn.name.function_name, argument_types, _data_type, _function));
         } else {
-            return Status::InternalError("Function {} is not endwith '_state'", _fn.signature);
+            return Status::InternalError("Unsupported AggState function {}", _fn.signature);
         }
     } else {
         // get the function. won't prepare function.

be/test/exprs/function/function_agg_state_finalize_test.cpp

@@ -0,0 +1,221 @@
+// 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.
+
+#include "exprs/function/function_agg_state_finalize.h"
+
+#include <gtest/gtest.h>
+
+#include "agent/be_exec_version_manager.h"
+#include "core/column/column_array.h"
+#include "core/column/column_const.h"
+#include "core/column/column_string.h"
+#include "core/data_type/data_type_agg_state.h"
+#include "core/data_type/data_type_number.h"
+#include "core/data_type/data_type_string.h"
+#include "exprs/aggregate/aggregate_function_state_merge.h"
+#include "testutil/column_helper.h"
+
+namespace doris {
+
+class FunctionAggStateFinalizeTest : public testing::Test {
+protected:
+    static std::shared_ptr<DataTypeAggState> state_type(const std::string& name,
+                                                        const DataTypePtr& argument_type,
+                                                        bool result_nullable = true) {
+        return std::make_shared<DataTypeAggState>(DataTypes {argument_type}, result_nullable, name,
+                                                  BeExecVersionManager::get_newest_version());
+    }
+
+    static void append_state(const AggregateFunctionPtr& function, const ColumnPtr& input,
+                             IColumn& states) {
+        ASSERT_EQ(function->get_argument_types().size(), 1);
+        Arena arena;
+        auto* place = arena.aligned_alloc(function->size_of_data(), function->align_of_data());
+        function->create(place);
+        DEFER(function->destroy(place));
+        const IColumn* columns[] = {input.get()};
+        function->check_input_columns_type(columns);
+        function->add_batch_single_place(input->size(), place, columns, arena);
+        // Some aggregates resize their no-key output instead of appending to it.
+        auto serialized = function->create_serialize_column();
+        function->serialize_without_key_to_column(place, *serialized);
+        ASSERT_EQ(serialized->size(), 1);
+        states.insert_range_from(*serialized, 0, 1);
+    }
+
+    static ColumnPtr finalize(const DataTypePtr& type, const ColumnPtr& states) {
+        auto nested = assert_cast<const DataTypeAggState*>(remove_nullable(type).get())
+                              ->get_nested_function();
+        auto result_type = nested->get_return_type();
+        if (type->is_nullable()) {
+            result_type = make_nullable(result_type);
+        }
+        auto function = FunctionAggStateFinalize::create({type}, result_type, nested);
+        Block block {{states, type, "state"}, {nullptr, result_type, "result"}};
+        EXPECT_TRUE(function->execute(nullptr, block, {0}, 1, states->size()).ok());
+        return block.get_by_position(1).column;
+    }
+};
+
+TEST_F(FunctionAggStateFinalizeTest, FinalizesEachRowIndependently) {
+    auto type = state_type("avg", std::make_shared<DataTypeInt64>());
+    auto states = type->create_column();
+    append_state(type->get_nested_function(), ColumnHelper::create_column<DataTypeInt64>({1, 3}),
+                 *states);
+    append_state(type->get_nested_function(), ColumnHelper::create_column<DataTypeInt64>({10}),
+                 *states);
+    ColumnPtr serialized = std::move(states);
+    auto result = finalize(type, serialized);
+    const auto& values = assert_cast<const ColumnFloat64&>(*result).get_data();
+    ASSERT_EQ(values.size(), 2);
+    EXPECT_DOUBLE_EQ(values[0], 2);
+    EXPECT_DOUBLE_EQ(values[1], 10);
+    // Repeated execution must not consume or mutate the serialized state.
+    EXPECT_EQ(result->compare_at(0, 0, *finalize(type, serialized), 1), 0);
+}
+
+TEST_F(FunctionAggStateFinalizeTest, EmptyNonNullableAvgMatchesMerge) {
+    auto type = state_type("avg", std::make_shared<DataTypeInt64>());
+    auto states = type->create_column();
+    append_state(type->get_nested_function(), ColumnHelper::create_column<DataTypeInt64>({}),
+                 *states);
+    ColumnPtr serialized = std::move(states);
+    auto result = finalize(type, serialized);
+    auto nested = type->get_nested_function();
+    auto merge = AggregateStateMerge::create(nested, {type}, nested->get_return_type());
+    Arena arena;
+    auto* place = arena.aligned_alloc(merge->size_of_data(), merge->align_of_data());
+    merge->create(place);
+    DEFER(merge->destroy(place));
+    const IColumn* columns[] = {serialized.get()};
+    merge->add(place, columns, 0, arena);
+    auto expected = nested->get_return_type()->create_column();
+    merge->check_result_column_type(*expected);
+    merge->insert_result_into(place, *expected);
+    EXPECT_EQ(result->compare_at(0, 0, *expected, 1), 0);
+}
+
+TEST_F(FunctionAggStateFinalizeTest, NativeSerializedColumnsAndEmptyCount) {
+    for (const auto& name : {"count", "sum", "min", "max"}) {
+        SCOPED_TRACE(name);
+        auto type =
+                state_type(name, std::make_shared<DataTypeInt64>(), name != std::string("count"));
+        auto states = type->create_column();
+        append_state(type->get_nested_function(),
+                     ColumnHelper::create_column<DataTypeInt64>({2, 4}), *states);
+        append_state(type->get_nested_function(), ColumnHelper::create_column<DataTypeInt64>({}),
+                     *states);
+        auto result = finalize(type, std::move(states));
+        EXPECT_EQ(result->get_int(0), name == std::string("count") ? 2
+                                      : name == std::string("sum") ? 6
+                                      : name == std::string("min") ? 2
+                                                                   : 4);
+        // Empty non-nullable states have the existing aggregate's identity value.
+        auto expected = type->get_nested_function()->get_return_type()->create_column();
+        Arena arena;
+        auto function = type->get_nested_function();
+        auto* place = arena.aligned_alloc(function->size_of_data(), function->align_of_data());
+        function->create(place);
+        DEFER(function->destroy(place));
+        function->check_result_column_type(*expected);
+        function->insert_result_into(place, *expected);
+        EXPECT_EQ(result->compare_at(1, 0, *expected, 1), 0);
+    }
+}
+
+// GTest assertion macros inflate complexity in this table-driven check.
+// NOLINTNEXTLINE(readability-function-cognitive-complexity)
+TEST_F(FunctionAggStateFinalizeTest, NullableInputsAndEmptyStates) {
+    for (const auto& name : {"avg", "sum", "min", "max", "count"}) {
+        SCOPED_TRACE(name);
+        auto type = state_type(name, make_nullable(std::make_shared<DataTypeInt64>()),
+                               name != std::string("count"));
+        auto states = type->create_column();
+        append_state(type->get_nested_function(),
+                     ColumnHelper::create_nullable_column<DataTypeInt64>({2, 0, 4}, {0, 1, 0}),
+                     *states);
+        append_state(type->get_nested_function(),
+                     ColumnHelper::create_nullable_column<DataTypeInt64>({0}, {1}), *states);
+        append_state(type->get_nested_function(),
+                     ColumnHelper::create_nullable_column<DataTypeInt64>({}, {}), *states);
+        auto result = finalize(type, std::move(states));
+        ASSERT_EQ(result->size(), 3);
+        EXPECT_FALSE(result->is_null_at(0));
+        if (name == std::string("count")) {
+            EXPECT_EQ(result->get_int(0), 2);
+            EXPECT_EQ(result->get_int(1), 0);
+            EXPECT_EQ(result->get_int(2), 0);
+        } else {
+            EXPECT_TRUE(result->is_null_at(1));
+            EXPECT_TRUE(result->is_null_at(2));
+        }
+    }
+}
+
+TEST_F(FunctionAggStateFinalizeTest, OuterNullAndConstantStates) {
+    auto type = state_type("avg", make_nullable(std::make_shared<DataTypeInt64>()));
+    auto states = type->create_column();
+    append_state(type->get_nested_function(),
+                 ColumnHelper::create_nullable_column<DataTypeInt64>({2, 4}, {0, 0}), *states);
+    ColumnPtr constant = ColumnConst::create(states->clone_resized(1), 5);
+    auto result = finalize(type, constant);
+    EXPECT_TRUE(is_column_const(*result));
+    EXPECT_EQ(result->size(), 5);
+    // NULL string payload is deliberately empty, and must not be deserialized.
+    states->insert_default();
+    ColumnPtr nullable_states = ColumnNullable::create(
+            std::move(states), ColumnHelper::create_column<DataTypeUInt8>({0, 1}));
+    result = finalize(make_nullable(type), nullable_states);
+    ASSERT_EQ(result->size(), 2);
+    EXPECT_FALSE(result->is_null_at(0));
+    EXPECT_TRUE(result->is_null_at(1));
+    auto null_state = nullable_states->clone_empty();
+    null_state->insert_from(*nullable_states, 1);
+    ColumnPtr null_constant = ColumnConst::create(std::move(null_state), 7);
+    result = finalize(make_nullable(type), null_constant);
+    EXPECT_EQ(result->size(), 7);
+    EXPECT_TRUE(result->is_null_at(0));
+}
+
+TEST_F(FunctionAggStateFinalizeTest, VariableLengthArrayResultsOwnTheirData) {
+    auto type = state_type("array_agg", std::make_shared<DataTypeString>(), false);
+    auto states = type->create_column();
+    const std::string large(8192, 'x');
+    for (size_t group = 0; group < 16; ++group) {
+        append_state(type->get_nested_function(),
+                     ColumnHelper::create_column<DataTypeString>({large, std::to_string(group)}),
+                     *states);
+    }
+    auto result = finalize(type, std::move(states));
+    const auto& arrays = assert_cast<const ColumnArray&>(*result);
+    const auto& strings = assert_cast<const ColumnString&>(
+            assert_cast<const ColumnNullable&>(arrays.get_data()).get_nested_column());
+    ASSERT_EQ(arrays.size(), 16);
+    for (size_t group = 0; group < 16; ++group) {
+        EXPECT_EQ(arrays.get_offsets()[group], (group + 1) * 2);
+        EXPECT_EQ(strings.get_data_at(group * 2).to_string(), large);
+        EXPECT_EQ(strings.get_data_at(group * 2 + 1).to_string(), std::to_string(group));
+    }
+}
+
+TEST_F(FunctionAggStateFinalizeTest, EmptyBlock) {
+    auto type = state_type("avg", std::make_shared<DataTypeInt64>());
+    auto result = finalize(type, type->create_column());
+    EXPECT_EQ(result->size(), 0);
+}
+
+} // namespace doris

fe/fe-core/src/main/java/org/apache/doris/catalog/FunctionRegistry.java

@@ -115,7 +115,8 @@ public boolean isAggregateFunction(String dbName, String name) {
                 return containsAggregateFunction(functionBuilders);
             }
             if (isBuiltinAggStateCombinator(name)) {
-                return !name.endsWith(AggCombinerFunctionBuilder.STATE_SUFFIX);
+                return !name.endsWith(AggCombinerFunctionBuilder.STATE_SUFFIX)
+                        && !name.endsWith(AggCombinerFunctionBuilder.FINALIZE_SUFFIX);
             }
         }