[feature](be) Add aggregate state finalizers and batch merging

be/src/exprs/aggregate/aggregate_function_null.h

@@ -244,13 +244,15 @@ class AggregateFunctionNullBaseInline : public IAggregateFunctionHelper<Derived>
                                             AggregateDataPtr rhs, const IColumn* column,
                                             Arena& arena, const size_t num_rows) const override {
         if (nested_function->is_trivial()) {
-            BufferReadable buf({column->get_data_at(0).data, 0});
+            // Selected rows need independent readers: skipping a destination must not leave
+            // a shared reader pointing at that row's serialized state for the next destination.
             size_t size_of_data = this->size_of_data();
             if constexpr (result_is_nullable) {
                 for (int i = 0; i != num_rows; ++i) {
                     if (!places[i]) {
                         continue;
                     }
+                    BufferReadable buf(column->get_data_at(i));
                     buf.read_binary(*(bool*)(rhs + size_of_data * i));
                     if (get_flag(rhs + size_of_data * i)) {
                         nested_function->deserialize(nested_place(rhs + size_of_data * i), buf,
@@ -267,6 +269,7 @@ class AggregateFunctionNullBaseInline : public IAggregateFunctionHelper<Derived>
             } else {
                 for (size_t i = 0; i != num_rows; ++i) {
                     if (places[i]) {
+                        BufferReadable buf(column->get_data_at(i));
                         nested_function->deserialize(rhs + size_of_data * i, buf, arena);
                     }
                 }

be/src/exprs/aggregate/aggregate_function_state_union.h

@@ -71,11 +71,12 @@ class AggregateStateUnion : public AggregateFunctionNonFinalBase,
 
     void add_batch(size_t batch_size, AggregateDataPtr* places, size_t place_offset,
                    const IColumn** columns, Arena& arena, bool agg_many) const override {
-        for (size_t i = 0; i < batch_size; ++i) {
-            //the range is [i, i]
-            _function->deserialize_and_merge_from_column_range(places[i] + place_offset,
-                                                               *columns[0], i, i, arena);
-        }
+        add_batch_impl<false>(batch_size, places, place_offset, columns, arena);
+    }
+
+    void add_batch_selected(size_t batch_size, AggregateDataPtr* places, size_t place_offset,
+                            const IColumn** columns, Arena& arena) const override {
+        add_batch_impl<true>(batch_size, places, place_offset, columns, arena);
     }
     void reset(AggregateDataPtr place) const override { _function->reset(place); }
 
@@ -107,6 +108,32 @@ class AggregateStateUnion : public AggregateFunctionNonFinalBase,
 
     size_t align_of_data() const override { return _function->align_of_data(); }
 
+private:
+    template <bool selected>
+    void add_batch_impl(size_t batch_size, AggregateDataPtr* places, size_t place_offset,
+                        const IColumn** columns, Arena& arena) const {
+        // Some batch implementations read the first serialized row before entering their loop.
+        if (batch_size == 0) {
+            return;
+        }
+        // Delegate the whole batch so native state columns (e.g. SUM) can merge directly,
+        // instead of dispatching through the range interface for every input row.
+        // Generic implementations need aligned scratch states and manage their construction,
+        // destruction and exception cleanup. Release only this storage after the batch;
+        // keep passing the caller's arena for allocations that merged states may retain.
+        // The function can be shared by evaluators, so scratch must not be a mutable member.
+        Arena scratch_arena;
+        auto* scratch = scratch_arena.aligned_alloc(batch_size * _function->size_of_data(),
+                                                    _function->align_of_data());
+        if constexpr (selected) {
+            _function->deserialize_and_merge_vec_selected(places, place_offset, scratch, columns[0],
+                                                          arena, batch_size);
+        } else {
+            _function->deserialize_and_merge_vec(places, place_offset, scratch, columns[0], arena,
+                                                 batch_size);
+        }
+    }
+
 protected:
     AggregateFunctionPtr _function;
     DataTypePtr _return_type;

be/src/exprs/function/function_agg_state_finalize.h

@@ -0,0 +1,110 @@
+// 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 <cstring>
+
+#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);
+        output->reserve(input_rows_count);
+        // Keep the reusable state storage separate from its per-row variable-length data.
+        Arena state_arena;
+        const auto state_size = _agg_function->size_of_data();
+        auto* place = state_arena.aligned_alloc(state_size, _agg_function->align_of_data());
+        Arena arena;
+        auto finalize_row = [&](size_t row) {
+            // Serialized states can use native columns as well as strings.
+            _agg_function->deserialize_and_merge_from_column_range(place, states, row, row, arena);
+            _agg_function->insert_result_into(place, *output);
+        };
+        auto finalize_rows = [&]<bool is_trivial>() {
+            for (size_t row = 0; row < input_rows_count; ++row) {
+                if (nullable && nullable->is_null_at(row)) {
+                    output->insert_default();
+                    continue;
+                }
+                if constexpr (is_trivial) {
+                    // Trivial states support zero-init and need no destruction.
+                    std::memset(place, 0, state_size);
+                    finalize_row(row);
+                } else {
+                    _agg_function->create(place);
+                    DEFER(_agg_function->destroy(place));
+                    finalize_row(row);
+                }
+                // Destroy states before reclaiming any variable-length data they own.
+                arena.clear();
+            }
+        };
+        if (_agg_function->is_trivial()) {
+            finalize_rows.template operator()<true>();
+        } else {
+            finalize_rows.template operator()<false>();
+        }
+        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.