From e39631ab4d0af73a0f9783360bcf25376bc99d2a Mon Sep 17 00:00:00 2001 From: happenlee Date: Mon, 21 Sep 2026 13:01:32 +0800 Subject: [PATCH 1/3] [feature](be) Add scalar aggregate state finalize functions ### What problem does this PR solve? Problem Summary: A query that has already computed one aggregate state per key still needs an aggregate operator to obtain its final values through `_merge`. This repeats aggregation for the finest grouping when states are also reused for coarser rollups. Add the scalar `_finalize(state)` combinator. It returns one result for each input state, using the existing aggregate implementation and serialized-state representation. For example: ```sql SELECT k, avg_finalize(s) FROM (SELECT k, avg_combine(v) AS s FROM t GROUP BY k) partial; ``` The same implementation supports aggregates such as `count`, `sum`, `min`, `max`, and `array_agg`. FE validates that the state's canonical aggregate name matches the finalizer, derives its result type, and treats the function as scalar. BE handles the underlying serialized column type, skips outer NULL payloads, and releases temporary state after each row. Constant inputs use the ordinary scalar constant path. Existing empty-state semantics and serialized formats are unchanged. This adds the scalar building block only; it does not change optimizer rollup rewrites. ### Release note Add `_finalize(AGG_STATE)` scalar functions to retrieve each aggregate state's result without merging rows. ### Check List (For Author) - Test - [x] Regression test: test_agg_state_finalize on a fresh local ASAN BE + FE cluster; output generated by the standard runner, checked against direct original aggregates, and verified by a normal comparison run. - [x] Unit Test: 33 FE tests passed across FinalizeCombinatorTest, StateCombinatorTest, CombineCombinatorTest, and FunctionRegistryTest. All 7 FunctionAggStateFinalizeTest cases passed under ASAN. - [x] Manual test: compare grouped and empty AVG/COUNT/SUM/MIN/MAX, decimal AVG and ARRAY_AGG with the original aggregates; verify rollup AVG is 14/3 and distinguish a missing outer-join state from COUNT's empty state. - [ ] No need to test or manual test. Explain why: - Behavior changed: - [ ] No. - [x] Yes. Add a family of scalar finalization functions; existing aggregate/state semantics are unchanged. - Does this need documentation? - [ ] No. - [x] Yes. Usage and semantics are included in the function-combinator README. ### Check List (For Reviewer who merge this PR) - [ ] Confirm the release note - [ ] Confirm test cases - [ ] Confirm document - [ ] Add branch pick label ### Additional validation - Standard ASAN BE + FE build, FE Checkstyle, clang-format 16, build hygiene and source whitespace checks passed. - clang-tidy was attempted with the repository script and the production/test translation units. New-code style and added cognitive-complexity warnings were fixed. A fully clean run remains blocked by existing header diagnostics and analyzer issues (including a test-helper array-bound path that does not connect its asserted unary arity to the input array); the local tool's crashing `modernize-use-scoped-lock` check was disabled only for the supplemental analysis. No repository analysis configuration was changed. --- .../function/function_agg_state_finalize.h | 93 ++++++++ be/src/exprs/vectorized_fn_call.cpp | 34 ++- .../function_agg_state_finalize_test.cpp | 221 ++++++++++++++++++ .../doris/catalog/FunctionRegistry.java | 3 +- .../glue/translator/ExpressionTranslator.java | 8 + .../functions/AggCombinerFunctionBuilder.java | 14 +- .../combinator/FinalizeCombinator.java | 98 ++++++++ .../functions/combinator/readme.md | 22 ++ .../visitor/ScalarFunctionVisitor.java | 5 + .../combinator/FinalizeCombinatorTest.java | 127 ++++++++++ .../agg_state/test_agg_state_finalize.out | 60 +++++ .../agg_state/test_agg_state_finalize.groovy | 104 +++++++++ 12 files changed, 783 insertions(+), 6 deletions(-) create mode 100644 be/src/exprs/function/function_agg_state_finalize.h create mode 100644 be/test/exprs/function/function_agg_state_finalize_test.cpp create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinatorTest.java create mode 100644 regression-test/data/datatype_p0/agg_state/test_agg_state_finalize.out create mode 100644 regression-test/suites/datatype_p0/agg_state/test_agg_state_finalize.groovy diff --git a/be/src/exprs/function/function_agg_state_finalize.h b/be/src/exprs/function/function_agg_state_finalize.h new file mode 100644 index 00000000000000..dd268d169f8c9b --- /dev/null +++ b/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( + std::make_shared(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(*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 diff --git a/be/src/exprs/vectorized_fn_call.cpp b/be/src/exprs/vectorized_fn_call.cpp index de20bd6e74a256..5fc898975c6633 100644 --- a/be/src/exprs/vectorized_fn_call.cpp +++ b/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 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(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(_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. diff --git a/be/test/exprs/function/function_agg_state_finalize_test.cpp b/be/test/exprs/function/function_agg_state_finalize_test.cpp new file mode 100644 index 00000000000000..355afa72932059 --- /dev/null +++ b/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 + +#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 state_type(const std::string& name, + const DataTypePtr& argument_type, + bool result_nullable = true) { + return std::make_shared(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(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()); + auto states = type->create_column(); + append_state(type->get_nested_function(), ColumnHelper::create_column({1, 3}), + *states); + append_state(type->get_nested_function(), ColumnHelper::create_column({10}), + *states); + ColumnPtr serialized = std::move(states); + auto result = finalize(type, serialized); + const auto& values = assert_cast(*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()); + auto states = type->create_column(); + append_state(type->get_nested_function(), ColumnHelper::create_column({}), + *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(), name != std::string("count")); + auto states = type->create_column(); + append_state(type->get_nested_function(), + ColumnHelper::create_column({2, 4}), *states); + append_state(type->get_nested_function(), ColumnHelper::create_column({}), + *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()), + name != std::string("count")); + auto states = type->create_column(); + append_state(type->get_nested_function(), + ColumnHelper::create_nullable_column({2, 0, 4}, {0, 1, 0}), + *states); + append_state(type->get_nested_function(), + ColumnHelper::create_nullable_column({0}, {1}), *states); + append_state(type->get_nested_function(), + ColumnHelper::create_nullable_column({}, {}), *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())); + auto states = type->create_column(); + append_state(type->get_nested_function(), + ColumnHelper::create_nullable_column({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({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(), 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({large, std::to_string(group)}), + *states); + } + auto result = finalize(type, std::move(states)); + const auto& arrays = assert_cast(*result); + const auto& strings = assert_cast( + assert_cast(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()); + auto result = finalize(type, type->create_column()); + EXPECT_EQ(result->size(), 0); +} + +} // namespace doris diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/FunctionRegistry.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/FunctionRegistry.java index 0a82f786c2207a..fbc3e9d84af333 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/FunctionRegistry.java +++ b/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); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java index 162bece0228660..c2ea0d402a5b36 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java @@ -90,6 +90,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.Count; import org.apache.doris.nereids.trees.expressions.functions.agg.NotNullableAggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.combinator.CombineCombinator; +import org.apache.doris.nereids.trees.expressions.functions.combinator.FinalizeCombinator; import org.apache.doris.nereids.trees.expressions.functions.combinator.ForEachCombinator; import org.apache.doris.nereids.trees.expressions.functions.combinator.MergeCombinator; import org.apache.doris.nereids.trees.expressions.functions.combinator.StateCombinator; @@ -822,6 +823,13 @@ public Expr visitIsNull(IsNull isNull, PlanTranslatorContext context) { return isNullPredicate; } + @Override + public Expr visitFinalizeCombinator(FinalizeCombinator combinator, PlanTranslatorContext context) { + FunctionCallExpr functionCallExpr = (FunctionCallExpr) visitScalarFunction(combinator, context); + functionCallExpr.getFn().setBinaryType(Function.BinaryType.AGG_STATE); + return functionCallExpr; + } + @Override public Expr visitStateCombinator(StateCombinator combinator, PlanTranslatorContext context) { List arguments = combinator.getArguments().stream() diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/AggCombinerFunctionBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/AggCombinerFunctionBuilder.java index c32162ca095501..90ba14b1a2ffcf 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/AggCombinerFunctionBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/AggCombinerFunctionBuilder.java @@ -23,6 +23,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.agg.NotSupportAggState; import org.apache.doris.nereids.trees.expressions.functions.combinator.CombineCombinator; +import org.apache.doris.nereids.trees.expressions.functions.combinator.FinalizeCombinator; import org.apache.doris.nereids.trees.expressions.functions.combinator.ForEachCombinator; import org.apache.doris.nereids.trees.expressions.functions.combinator.MergeCombinator; import org.apache.doris.nereids.trees.expressions.functions.combinator.StateCombinator; @@ -43,12 +44,14 @@ public class AggCombinerFunctionBuilder extends FunctionBuilder { public static final String COMBINATOR_LINKER = "_"; public static final String STATE = "state"; public static final String MERGE = "merge"; + public static final String FINALIZE = "finalize"; public static final String UNION = "union"; public static final String COMBINE = "combine"; public static final String FOREACH = "foreach"; public static final String STATE_SUFFIX = COMBINATOR_LINKER + STATE; public static final String MERGE_SUFFIX = COMBINATOR_LINKER + MERGE; + public static final String FINALIZE_SUFFIX = COMBINATOR_LINKER + FINALIZE; public static final String UNION_SUFFIX = COMBINATOR_LINKER + UNION; public static final String COMBINE_SUFFIX = COMBINATOR_LINKER + COMBINE; public static final String FOREACH_SUFFIX = COMBINATOR_LINKER + FOREACH; @@ -121,7 +124,7 @@ private AggregateFunction buildForEach(String nestedName, List return (AggregateFunction) nestedBuilder.build(nestedName, forEachargs).first; } - private AggregateFunction buildMergeOrUnion(String nestedName, List arguments) { + private AggregateFunction buildFromState(String nestedName, List arguments) { if (arguments.size() != 1 || !(arguments.get(0) instanceof Expression) || !((Expression) arguments.get(0)).getDataType().isAggStateType()) { String argString = arguments.stream().map(arg -> { @@ -162,10 +165,13 @@ public Pair build(String name, List argumen AggregateFunction nestedFunction = buildState(nestedName, arguments); return Pair.of(new CombineCombinator((List) arguments, nestedFunction), nestedFunction); } else if (combinatorSuffix.equalsIgnoreCase(MERGE)) { - AggregateFunction nestedFunction = buildMergeOrUnion(nestedName, arguments); + AggregateFunction nestedFunction = buildFromState(nestedName, arguments); return Pair.of(new MergeCombinator((List) arguments, nestedFunction), nestedFunction); + } else if (combinatorSuffix.equalsIgnoreCase(FINALIZE)) { + AggregateFunction nestedFunction = buildFromState(nestedName, arguments); + return Pair.of(new FinalizeCombinator((List) arguments, nestedFunction), nestedFunction); } else if (combinatorSuffix.equalsIgnoreCase(UNION)) { - AggregateFunction nestedFunction = buildMergeOrUnion(nestedName, arguments); + AggregateFunction nestedFunction = buildFromState(nestedName, arguments); return Pair.of(new UnionCombinator((List) arguments, nestedFunction), nestedFunction); } else if (combinatorSuffix.equalsIgnoreCase(FOREACH)) { AggregateFunction nestedFunction = buildForEach(nestedName, arguments); @@ -182,7 +188,7 @@ public String parameterDisplayString() { public static boolean isAggStateCombinator(String name) { return name.toLowerCase().endsWith(STATE_SUFFIX) || name.toLowerCase().endsWith(MERGE_SUFFIX) || name.toLowerCase().endsWith(UNION_SUFFIX) || name.toLowerCase().endsWith(COMBINE_SUFFIX) - || name.toLowerCase().endsWith(FOREACH_SUFFIX); + || name.toLowerCase().endsWith(FOREACH_SUFFIX) || name.toLowerCase().endsWith(FINALIZE_SUFFIX); } public static String getNestedName(String name) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java new file mode 100644 index 00000000000000..5191fcf61bb3ef --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java @@ -0,0 +1,98 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.combinator; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AggCombinerFunctionBuilder; +import org.apache.doris.nereids.trees.expressions.functions.ComputeNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ScalarFunction; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ScalarFunctionParams; +import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.AggStateType; + +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** Finalize each serialized aggregate state without aggregating rows. */ +public class FinalizeCombinator extends ScalarFunction + implements UnaryExpression, ExplicitlyCastableSignature, ComputeNullable, Combinator { + + private final AggregateFunction nested; + + public FinalizeCombinator(List arguments, AggregateFunction nested) { + super(nested.getName() + AggCombinerFunctionBuilder.FINALIZE_SUFFIX, arguments); + this.nested = nested; + checkStateFunction(); + } + + private FinalizeCombinator(ScalarFunctionParams functionParams, AggregateFunction nested) { + super(functionParams); + this.nested = nested; + checkStateFunction(); + } + + private void checkStateFunction() { + AggStateType inputType = (AggStateType) getArgument(0).getDataType(); + // AggStateType canonicalizes aliases; acceptsType alone accepts unrelated aggregate states. + AggStateType expected = new AggStateType(nested.getName(), inputType.getSubTypes(), + inputType.getSubTypeNullables(), nested.nullable()); + if (!inputType.getFunctionName().equals(expected.getFunctionName())) { + throw new AnalysisException(getName() + " requires a state of " + expected.getFunctionName() + + ", but got " + inputType.toSql()); + } + } + + @Override + public FinalizeCombinator withChildren(List children) { + return new FinalizeCombinator(getFunctionParams(children), nested); + } + + @Override + public List getSignatures() { + return ImmutableList.of(FunctionSignature.ret(nested.getDataType()).args(getArgument(0).getDataType())); + } + + @Override + public boolean nullable() { + return getArgument(0).nullable() || nested.nullable(); + } + + @Override + public AggregateFunction getNestedFunction() { + return nested; + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitFinalizeCombinator(this, context); + } + + @Override + public void checkLegalityBeforeTypeCoercion() { + // Stored states retain serialized parameters, not the original constant expressions. + if (getArgument(0) instanceof StateCombinator || getArgument(0) instanceof CombineCombinator) { + nested.checkLegalityBeforeTypeCoercion(); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/readme.md b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/readme.md index e861156605a873..28d06aad7cfa06 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/readme.md +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/readme.md @@ -5,3 +5,25 @@ The function combiner is used to automatically create some functions based on th Currently we can generate the combinator function we want by implementing a new `FunctionBuilder`. For example `AggStateFunctionBuilder` can generate combiner functions created by the `_state`/`_merge` combiner, so that we can use functions such as `sum_state` and `sum_merge` in SQL statements. When we input `sum_state`, if `BuiltinFunctionBuilder` does not find a function named `sum_state`, then it will try to use `AggStateFunctionBuilder` to construct `StateCombinator`, so we get a `StateCombinator` object with nested function sum, and use it to generate the `sum_state` function. + +## Finalize one aggregate state + +`_finalize(state)` is a scalar combinator that returns the result of each +serialized aggregate state independently. It accepts exactly one `AGG_STATE` of +the matching aggregate (including its aliases), and uses the same result type and +state semantics as `_merge`. An outer SQL NULL state returns NULL; +an empty serialized state follows the nested aggregate's existing behavior. + +For example, after pre-aggregating each key, the finest grouping can read its AVG +without another aggregate operator: + +```sql +SET enable_agg_state = true; +SELECT k, avg_finalize(s) +FROM (SELECT k, avg_combine(v) AS s FROM t GROUP BY k) partial; +``` + +`avg_finalize` returns one result per input row. In contrast, `avg_merge` merges +all input states in each group before computing one result. The same scalar +combinator applies to other aggregates supporting `AGG_STATE`, such as `sum`, +`count`, `min`, `max`, and `array_agg`; it does not change their serialized formats. diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java index 3ba1ab10352bac..9183d9317bed4a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java @@ -31,6 +31,7 @@ import org.apache.doris.nereids.trees.expressions.functions.ai.AISummarize; import org.apache.doris.nereids.trees.expressions.functions.ai.AITranslate; import org.apache.doris.nereids.trees.expressions.functions.ai.Embed; +import org.apache.doris.nereids.trees.expressions.functions.combinator.FinalizeCombinator; import org.apache.doris.nereids.trees.expressions.functions.combinator.StateCombinator; import org.apache.doris.nereids.trees.expressions.functions.scalar.Abs; import org.apache.doris.nereids.trees.expressions.functions.scalar.Acos; @@ -2847,6 +2848,10 @@ default R visitYearsDiff(YearsDiff yearsDiff, C context) { return visitScalarFunction(yearsDiff, context); } + default R visitFinalizeCombinator(FinalizeCombinator combinator, C context) { + return visitScalarFunction(combinator, context); + } + default R visitStateCombinator(StateCombinator combinator, C context) { return visitScalarFunction(combinator, context); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinatorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinatorTest.java new file mode 100644 index 00000000000000..080ba6347cfc5d --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinatorTest.java @@ -0,0 +1,127 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.combinator; + +import org.apache.doris.analysis.FunctionCallExpr; +import org.apache.doris.catalog.Function; +import org.apache.doris.catalog.FunctionRegistry; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.glue.translator.ExpressionTranslator; +import org.apache.doris.nereids.glue.translator.PlanTranslatorContext; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.FunctionBuilder; +import org.apache.doris.nereids.trees.expressions.functions.agg.Avg; +import org.apache.doris.nereids.trees.expressions.functions.agg.Count; +import org.apache.doris.nereids.types.AggStateType; +import org.apache.doris.nereids.types.DecimalV3Type; +import org.apache.doris.nereids.types.DoubleType; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.nereids.util.PlanChecker; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class FinalizeCombinatorTest { + private FinalizeCombinator build(String name, Expression argument) { + FunctionBuilder builder = new FunctionRegistry().findFunctionBuilder(name, argument); + return (FinalizeCombinator) builder.build(name, argument).first; + } + + @Test + void testScalarClassificationAndTranslation() { + FunctionRegistry registry = new FunctionRegistry(); + Assertions.assertTrue(registry.isBuiltinAggStateCombinator("avg_finalize")); + Assertions.assertFalse(registry.isAggregateFunction(null, "avg_finalize")); + Assertions.assertFalse(registry.isBuiltinAggStateCombinator("abs_finalize")); + SlotReference value = new SlotReference("v", IntegerType.INSTANCE, true); + StateCombinator state = StateCombinator.create(new Avg(value)); + FinalizeCombinator finalize = build("avg_finalize", state); + Assertions.assertEquals(DoubleType.INSTANCE, finalize.getDataType()); + Assertions.assertTrue(finalize.nullable()); + FunctionCallExpr translated = (FunctionCallExpr) ExpressionTranslator.translate( + finalize, new PlanTranslatorContext()); + Assertions.assertEquals("avg_finalize", translated.getFn().getFunctionName().getFunction()); + Assertions.assertEquals(Function.BinaryType.AGG_STATE, translated.getFn().getBinaryType()); + Assertions.assertEquals(Function.NullableMode.ALWAYS_NULLABLE, translated.getFn().getNullableMode()); + Assertions.assertEquals(state.getDataType().toCatalogDataType(), translated.getFn().getArgs()[0]); + } + + @Test + void testCountAndOuterNullability() { + SlotReference value = new SlotReference("v", IntegerType.INSTANCE, true); + StateCombinator state = StateCombinator.create(new Count(value)); + Assertions.assertFalse(build("count_finalize", state).nullable()); + SlotReference nullableState = new SlotReference("s", state.getDataType(), true); + Assertions.assertTrue(build("count_finalize", nullableState).nullable()); + } + + @Test + void testNonNullableStateUsesNestedResultType() { + SlotReference value = new SlotReference("v", IntegerType.INSTANCE, false); + StateCombinator state = StateCombinator.create(new Avg(value)); + Assertions.assertFalse(build("avg_finalize", state).nullable()); + } + + @Test + void testDecimalResultPrecision() { + MemoTestUtils.createConnectContext(); + SlotReference value = new SlotReference("v", DecimalV3Type.createDecimalV3Type(12, 2), true); + FinalizeCombinator finalize = build("avg_finalize", StateCombinator.create(new Avg(value))); + Assertions.assertEquals(DecimalV3Type.createDecimalV3Type(38, 4), finalize.getDataType()); + } + + @Test + void testStateFunctionMustMatch() { + for (String sql : ImmutableList.of("select sum_finalize(avg_state(1))", + "select avg_finalize(sum_state(1))")) { + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> PlanChecker.from(MemoTestUtils.createConnectContext()).analyze(sql)); + Assertions.assertTrue(exception.getMessage().contains("requires a state of")); + } + } + + @Test + void testAliasCanonicalization() { + AggStateType state = new AggStateType("var_pop", ImmutableList.of(DoubleType.INSTANCE), + ImmutableList.of(true), true); + FinalizeCombinator finalize = build("variance_finalize", new SlotReference("s", state, false)); + Assertions.assertEquals("variance_finalize", finalize.getName()); + Assertions.assertEquals(DoubleType.INSTANCE, finalize.getDataType()); + } + + @Test + void testInvalidArguments() { + for (String sql : ImmutableList.of("select avg_finalize(1)", "select avg_finalize()", + "select avg_finalize(avg_state(1), avg_state(2))", + "select avg_finalize(distinct avg_state(1))")) { + Assertions.assertThrows(AnalysisException.class, + () -> PlanChecker.from(MemoTestUtils.createConnectContext()).analyze(sql)); + } + } + + @Test + void testStoredParameterizedStateAndCombine() { + PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select topn_finalize(s) from (select topn_combine('x', 2) s) t"); + PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select avg_finalize(avg_combine(cast(1 as int)))"); + } +} diff --git a/regression-test/data/datatype_p0/agg_state/test_agg_state_finalize.out b/regression-test/data/datatype_p0/agg_state/test_agg_state_finalize.out new file mode 100644 index 00000000000000..763c14200737d4 --- /dev/null +++ b/regression-test/data/datatype_p0/agg_state/test_agg_state_finalize.out @@ -0,0 +1,60 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !rows -- +1 1 1 1 1 1 +1 3 1 3 3 3 +2 10 1 10 10 10 +3 \N 0 \N \N \N + +-- !grouped -- +1 2 2 4 1 3 +2 10 1 10 10 10 +3 \N 0 \N \N \N + +-- !empty -- +\N 0 \N + +-- !decimal -- +1 2.5000 +2 10.0000 +3 \N + +-- !array -- +1 ["a", "b"] +2 ["c"] +3 [null] + +-- !parameters -- +{"a":1} + +-- !const -- +0 7 0 +1 7 0 +2 7 0 + +-- !rollup -- +\N 4.666666666666667 +1 2 +2 10 +3 \N + +-- !stored -- +1 2 2 4 +2 10 1 10 +3 \N 0 \N + +-- !outer_null -- +0 \N \N +1 2 2 +2 10 1 +3 \N 0 +4 \N \N + +-- !union -- +4.666666666666667 + +-- !alias -- +\N \N +0 0 +0 0 +0 0 + diff --git a/regression-test/suites/datatype_p0/agg_state/test_agg_state_finalize.groovy b/regression-test/suites/datatype_p0/agg_state/test_agg_state_finalize.groovy new file mode 100644 index 00000000000000..5285a97d5545ef --- /dev/null +++ b/regression-test/suites/datatype_p0/agg_state/test_agg_state_finalize.groovy @@ -0,0 +1,104 @@ +// 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_agg_state_finalize") { + sql "set enable_agg_state=true" + sql "drop table if exists agg_finalize_input" + sql """ + create table agg_finalize_input (k int, v int, s string, d decimal(12, 2)) + duplicate key(k) distributed by hash(k) buckets 1 + properties("replication_num"="1") + """ + sql """ + insert into agg_finalize_input values + (1, 1, 'a', 1.25), (1, 3, 'b', 3.75), + (2, 10, 'c', 10.00), (3, null, null, null) + """ + order_qt_rows """ + select k, avg_finalize(avg_state(v)), count_finalize(count_state(v)), + sum_finalize(sum_state(v)), min_finalize(min_state(v)), max_finalize(max_state(v)) + from agg_finalize_input + """ + order_qt_grouped """ + select k, avg_finalize(a), count_finalize(c), sum_finalize(s), min_finalize(mi), max_finalize(ma) + from (select k, avg_combine(v) a, count_combine(v) c, sum_combine(v) s, + min_combine(v) mi, max_combine(v) ma from agg_finalize_input group by k) t + """ + order_qt_empty """ + select avg_finalize(a), count_finalize(c), sum_finalize(s) + from (select avg_combine(v) a, count_combine(v) c, sum_combine(v) s + from agg_finalize_input where k < 0) t + """ + order_qt_decimal """ + select k, avg_finalize(a) + from (select k, avg_combine(d) a from agg_finalize_input group by k) t + """ + order_qt_array """ + select k, array_sort(array_agg_finalize(a)) + from (select k, array_agg_combine(s) a from agg_finalize_input group by k) t + """ + order_qt_parameters """ + select topn_finalize(s) from (select topn_combine('a', 2) s) t + """ + order_qt_const """ + select number, avg_finalize(avg_state(7)), count_finalize(count_state(cast(null as int))) + from numbers("number"="3") + """ + // The finest groups need only a scalar projection; coarser groups merge their states. + order_qt_rollup """ + with partial as (select k, avg_combine(v) a from agg_finalize_input group by k) + select k, avg_finalize(a) from partial + union all + select null, avg_merge(a) from partial + """ + sql "drop table if exists agg_finalize_stored" + sql """ + create table agg_finalize_stored ( + k int, a agg_state generic, c agg_state generic, + s agg_state generic + ) aggregate key(k) distributed by hash(k) buckets 1 + properties("replication_num"="1") + """ + sql """ + insert into agg_finalize_stored + select k, avg_state(v), count_state(v), sum_state(v) from agg_finalize_input + """ + order_qt_stored """ + select k, avg_finalize(a), count_finalize(c), sum_finalize(s) from agg_finalize_stored + """ + order_qt_outer_null """ + select n.number, avg_finalize(t.a), count_finalize(t.c) + from numbers("number"="5") n left join agg_finalize_stored t on n.number=t.k + """ + order_qt_union """ + select avg_finalize(a) from (select avg_union(a) a from agg_finalize_stored) t + """ + order_qt_alias """ + select var_pop_finalize(variance_state(v)), variance_finalize(var_pop_state(v)) + from agg_finalize_input + """ + for (def query : ["select sum_finalize(avg_state(1))", "select avg_finalize(sum_state(1))"]) { + test { + sql query + exception "requires a state of" + } + } + test { + sql "select avg_finalize(1)" + exception "Can not found function" + } +} From c620bed661c583f717357f4bad6d16709b7c35a9 Mon Sep 17 00:00:00 2001 From: happenlee Date: Mon, 21 Sep 2026 15:32:49 +0800 Subject: [PATCH 2/3] [doc](fe) Document deferred aggregate state alias support ### What problem does this PR solve? Issue Number: N/A Related PR: #68312 Problem Summary: The finalizer's name-check comment implied that all aggregate aliases were canonicalized. Keep strict state/function name validation and add a TODO documenting that stored aliases missing from AggStateType's mapping, such as std, remain unsupported until FE and BE canonicalization is unified. ### Release note None ### Check List (For Author) - Test: Maven Checkstyle and git diff --check - No need to run runtime tests: comments only, with no behavior change - Behavior changed: No - Does this need documentation: No --- .../expressions/functions/combinator/FinalizeCombinator.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java index 5191fcf61bb3ef..10321d707d9e8f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java @@ -54,7 +54,9 @@ private FinalizeCombinator(ScalarFunctionParams functionParams, AggregateFunctio private void checkStateFunction() { AggStateType inputType = (AggStateType) getArgument(0).getDataType(); - // AggStateType canonicalizes aliases; acceptsType alone accepts unrelated aggregate states. + // Keep the name check because acceptsType alone accepts unrelated aggregate states. + // TODO: Unify aggregate alias canonicalization across FE and BE. Stored aliases absent + // from AggStateType's mapping (for example, std) remain unsupported by finalize for now. AggStateType expected = new AggStateType(nested.getName(), inputType.getSubTypes(), inputType.getSubTypeNullables(), nested.nullable()); if (!inputType.getFunctionName().equals(expected.getFunctionName())) { From 6e8121543b0548c42887f8bddf23b64ed00a4274 Mon Sep 17 00:00:00 2001 From: happenlee Date: Mon, 21 Sep 2026 16:35:57 +0800 Subject: [PATCH 3/3] [fix](fe) Preserve value expressions when matching aggregate state rollups ### What problem does this PR solve? Related PR: #68312 Problem Summary: MV rollup matching searched the entire argument tree for the last combinator. A scalar finalizer such as avg_finalize(avg_state(v)) inside SUM and MAX could make their outer aggregate states appear equivalent, and scalar wrappers such as ABS could be ignored. Matching plain aggregates with COMBINE states could also fail for otherwise identical value expressions. Follow only direct MERGE/UNION state chains for the same aggregate, stopping at STATE/COMBINE so their complete value expressions remain part of equality. Use the shared helper in both rollup handlers. Add FE tests and a regression suite covering invalid candidates, successful rollups, negative values and NULLs. Six of the nine new FE tests fail before the fix and all pass after it. ### Release note Fix materialized view rollup matching for aggregate states with nested scalar finalizers, preserving the outer aggregate and complete value expressions. ### Check List (For Author) - Test: 51 targeted FE unit tests passed; agg_state_finalize_roll_up and test_agg_state_finalize regression suites passed on local ASAN BE + rebuilt FE. Regression output generated with the standard runner and independently checked. Standard FE build, Checkstyle and source whitespace checks passed. - Behavior changed: Yes. Reject incompatible state rollups and retain valid rollups over identical finalized value expressions. - Does this need documentation: No. Correctness fix for existing MV rewrites. --- .../mv/rollup/AggFunctionRollUpHandler.java | 28 +++- .../rollup/BothCombinatorRollupHandler.java | 4 +- .../rollup/SingleCombinatorRollupHandler.java | 2 +- .../rollup/CombinatorRollupHandlerTest.java | 156 ++++++++++++++++++ .../agg_state_finalize_roll_up.out | 36 ++++ .../agg_state_finalize_roll_up.groovy | 92 +++++++++++ 6 files changed, 309 insertions(+), 9 deletions(-) create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/rollup/CombinatorRollupHandlerTest.java create mode 100644 regression-test/data/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.out create mode 100644 regression-test/suites/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.groovy diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/AggFunctionRollUpHandler.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/AggFunctionRollUpHandler.java index 190080424511c7..6d02b4379a2bbf 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/AggFunctionRollUpHandler.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/AggFunctionRollUpHandler.java @@ -23,6 +23,11 @@ import org.apache.doris.nereids.trees.expressions.functions.Function; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.agg.RollUpTrait; +import org.apache.doris.nereids.trees.expressions.functions.combinator.Combinator; +import org.apache.doris.nereids.trees.expressions.functions.combinator.CombineCombinator; +import org.apache.doris.nereids.trees.expressions.functions.combinator.MergeCombinator; +import org.apache.doris.nereids.trees.expressions.functions.combinator.StateCombinator; +import org.apache.doris.nereids.trees.expressions.functions.combinator.UnionCombinator; import com.google.common.collect.ImmutableList; @@ -73,12 +78,23 @@ protected static List extractArguments(Expression functionWithAny, E } /** - * Extract the target expression in actualFunction by targetClazz - * Such as actualFunction def is avg_merge(avg_union(c1)), target Clazz is Combinator - * after extracting, the return argument is avg_union(c1) + * Unwrap direct MERGE/UNION state chains for the same aggregate function. + * STATE and COMBINE consume values, so keep their complete argument expressions for comparison. + * For example, sum_combine(avg_finalize(avg_state(v))) must retain SUM and its finalized value. */ - protected static T extractLastExpression(Expression actualFunction, Class targetClazz) { - List expressions = actualFunction.collectToList(targetClazz::isInstance); - return targetClazz.cast(expressions.get(expressions.size() - 1)); + protected static Combinator extractRollupCombinator(Combinator current) { + while (current instanceof MergeCombinator || current instanceof UnionCombinator) { + Expression argument = current.getArguments().get(0); + if (!(argument instanceof StateCombinator || argument instanceof CombineCombinator + || argument instanceof UnionCombinator)) { + break; + } + Combinator next = (Combinator) argument; + if (!current.getNestedFunction().getName().equals(next.getNestedFunction().getName())) { + break; + } + current = next; + } + return current; } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/BothCombinatorRollupHandler.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/BothCombinatorRollupHandler.java index 38c1dedcefeb41..c4c047ce0a7c12 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/BothCombinatorRollupHandler.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/BothCombinatorRollupHandler.java @@ -47,8 +47,8 @@ public boolean canRollup(AggregateFunction queryAggregateFunction, return false; } if (queryAggregateFunction instanceof Combinator && viewFunction instanceof Combinator) { - Combinator queryCombinator = extractLastExpression(queryAggregateFunction, Combinator.class); - Combinator viewCombinator = extractLastExpression(viewFunction, Combinator.class); + Combinator queryCombinator = extractRollupCombinator((Combinator) queryAggregateFunction); + Combinator viewCombinator = extractRollupCombinator((Combinator) viewFunction); // construct actual aggregate function in combinator and compare return Objects.equals(queryCombinator.getNestedFunction().withChildren(queryCombinator.getArguments()), viewCombinator.getNestedFunction().withChildren(viewCombinator.getArguments())); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/SingleCombinatorRollupHandler.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/SingleCombinatorRollupHandler.java index 9a865d7a08c301..6d012a2f0c974b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/SingleCombinatorRollupHandler.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/rollup/SingleCombinatorRollupHandler.java @@ -56,7 +56,7 @@ public boolean canRollup(AggregateFunction queryAggregateFunction, if (!(queryAggregateFunction instanceof Combinator) && (viewFunction instanceof UnionCombinator || viewFunction instanceof StateCombinator || viewFunction instanceof CombineCombinator)) { - Combinator viewCombinator = extractLastExpression(viewFunction, Combinator.class); + Combinator viewCombinator = extractRollupCombinator((Combinator) viewFunction); return Objects.equals(queryAggregateFunction, viewCombinator.getNestedFunction().withChildren(viewCombinator.getArguments())); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/rollup/CombinatorRollupHandlerTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/rollup/CombinatorRollupHandlerTest.java new file mode 100644 index 00000000000000..4821dbef6e8689 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/rollup/CombinatorRollupHandlerTest.java @@ -0,0 +1,156 @@ +// 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. + +package org.apache.doris.nereids.rules.exploration.mv.rollup; + +import org.apache.doris.catalog.FunctionRegistry; +import org.apache.doris.common.Pair; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.Function; +import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; +import org.apache.doris.nereids.trees.expressions.functions.agg.Sum; +import org.apache.doris.nereids.trees.expressions.functions.combinator.StateCombinator; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Abs; +import org.apache.doris.nereids.types.DoubleType; +import org.apache.doris.nereids.util.MemoTestUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +class CombinatorRollupHandlerTest { + private final FunctionRegistry registry = new FunctionRegistry(); + private final SlotReference value = new SlotReference("v", DoubleType.INSTANCE, true); + + @BeforeAll + static void setUp() { + MemoTestUtils.createConnectContext(); + } + + private Expression build(String name, Expression argument) { + return registry.findFunctionBuilder(name, argument).build(name, argument).first; + } + + private Expression finalizeValue(boolean stored) { + Expression state = build("avg_state", value); + return build("avg_finalize", stored ? new SlotReference("s", state.getDataType(), true) : state); + } + + private boolean canRollup(AggFunctionRollUpHandler handler, Expression query, Expression view) { + SlotReference mvSlot = SlotReference.of("mv_state", view.getDataType()); + return handler.canRollup((AggregateFunction) query, query, Pair.of(view, mvSlot), + ImmutableMap.of(view, mvSlot)); + } + + private void assertRollup(AggFunctionRollUpHandler handler, Expression query, Expression view, String name) { + Assertions.assertTrue(canRollup(handler, query, view)); + SlotReference mvSlot = SlotReference.of("mv_state", view.getDataType()); + Function result = handler.doRollup((AggregateFunction) query, query, Pair.of(view, mvSlot), + ImmutableMap.of(view, mvSlot)); + Assertions.assertEquals(name, result.getName()); + Assertions.assertEquals(ImmutableList.of(mvSlot), result.getArguments()); + } + + @Test + void testDifferentOuterAggregatesOverFinalize() { + for (boolean stored : ImmutableList.of(false, true)) { + Expression finalized = finalizeValue(stored); + Assertions.assertFalse(canRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_combine", finalized), build("max_combine", finalized))); + } + } + + @Test + void testScalarArgumentMustMatch() { + Expression finalized = finalizeValue(false); + Assertions.assertFalse(canRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_combine", finalized), build("sum_combine", new Abs(finalized)))); + } + + @Test + void testSameOuterAggregateOverFinalize() { + Expression finalized = finalizeValue(false); + assertRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_combine", finalized), build("sum_combine", finalized), "sum_union"); + } + + @Test + void testMergeAndUnionStateChain() { + Expression state = build("sum_state", value); + assertRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_merge", state), build("sum_union", state), "sum_merge"); + assertRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_merge", build("sum_union", state)), build("sum_union", state), "sum_merge"); + } + + @Test + void testUnionStopsAtCombineValueArguments() { + Expression finalized = finalizeValue(false); + assertRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_merge", build("sum_combine", finalized)), + build("sum_union", build("sum_combine", finalized)), "sum_merge"); + Assertions.assertFalse(canRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_merge", build("sum_combine", finalized)), + build("max_union", build("max_combine", finalized)))); + } + + @Test + void testStateRetainsFinalizeArgument() { + Expression finalized = finalizeValue(false); + Assertions.assertFalse(canRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_merge", build("sum_state", finalized)), + build("max_union", build("max_state", finalized)))); + } + + @Test + void testStoredStateMustMatch() { + Expression state = build("sum_state", value); + SlotReference first = SlotReference.of("first_state", state.getDataType()); + SlotReference second = SlotReference.of("second_state", state.getDataType()); + assertRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_merge", first), build("sum_union", first), "sum_merge"); + Assertions.assertFalse(canRollup(BothCombinatorRollupHandler.INSTANCE, + build("sum_merge", first), build("sum_union", second))); + } + + @Test + void testSingleCombinatorOverFinalize() { + for (boolean stored : ImmutableList.of(false, true)) { + Expression finalized = finalizeValue(stored); + Sum query = new Sum(finalized); + assertRollup(SingleCombinatorRollupHandler.INSTANCE, + query, build("sum_combine", finalized), "sum_merge"); + Assertions.assertFalse(canRollup(SingleCombinatorRollupHandler.INSTANCE, + query, build("max_combine", finalized))); + Assertions.assertFalse(canRollup(SingleCombinatorRollupHandler.INSTANCE, + query, build("sum_combine", new Abs(finalized)))); + } + } + + @Test + void testSingleCombinatorStateChain() { + Expression finalized = finalizeValue(false); + Sum query = new Sum(finalized); + assertRollup(SingleCombinatorRollupHandler.INSTANCE, query, + build("sum_union", StateCombinator.create(query)), "sum_merge"); + assertRollup(SingleCombinatorRollupHandler.INSTANCE, query, + build("sum_union", build("sum_combine", finalized)), "sum_merge"); + } +} diff --git a/regression-test/data/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.out b/regression-test/data/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.out new file mode 100644 index 00000000000000..cd51bb4190f6e9 --- /dev/null +++ b/regression-test/data/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.out @@ -0,0 +1,36 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !combined_before -- +1 400 +2 500 +3 \N + +-- !plain_before -- +1 400 +2 500 +3 \N + +-- !different_aggregate -- +1 400 +2 500 +3 \N + +-- !different_argument -- +1 400 +2 500 +3 \N + +-- !combined_after -- +1 400 +2 500 +3 \N + +-- !plain_after -- +1 400 +2 500 +3 \N + +-- !merge_after -- +1 400 +2 500 +3 \N + diff --git a/regression-test/suites/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.groovy b/regression-test/suites/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.groovy new file mode 100644 index 00000000000000..cf57a756b49a6b --- /dev/null +++ b/regression-test/suites/nereids_rules_p0/mv/agg_with_roll_up/agg_state_finalize_roll_up.groovy @@ -0,0 +1,92 @@ +// 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("agg_state_finalize_roll_up") { + String db = context.config.getDbNameByFile(context.file) + sql "set enable_agg_state=true" + sql "set pre_materialized_view_rewrite_strategy=TRY_IN_RBO" + for (String mv : ["finalize_max_mv", "finalize_abs_mv", "finalize_sum_mv", "finalize_union_mv"]) { + sql "drop materialized view if exists ${mv}" + } + sql "drop table if exists agg_finalize_rollup_input" + sql """ + create table agg_finalize_rollup_input (k int, g int, v double) + duplicate key(k, g) distributed by hash(k) buckets 1 + properties("replication_num"="1") + """ + sql """ + insert into agg_finalize_rollup_input values + (1, 1, -4), (1, 1, 3), (1, 2, -2), (1, 2, 7), + (2, 1, 5), (2, 2, null), (3, 1, null) + """ + // Make the finer-group MV cheaper than scanning the base table so the test executes its rollup. + sql """ + insert into agg_finalize_rollup_input + select k, g, v from agg_finalize_rollup_input cross join numbers("number"="99") + """ + sql "analyze table agg_finalize_rollup_input with sync" + + String combinedQuery = """ + select k, sum_finalize(sum_combine(avg_finalize(avg_state(v)))) + from agg_finalize_rollup_input group by k + """ + String plainQuery = """ + select k, sum(avg_finalize(avg_state(v))) + from agg_finalize_rollup_input group by k + """ + order_qt_combined_before combinedQuery + order_qt_plain_before plainQuery + + // Sharing the inner AVG finalizer must not make SUM compatible with MAX. + create_async_mv(db, "finalize_max_mv", """ + select k, g, max_combine(avg_finalize(avg_state(v))) s + from agg_finalize_rollup_input group by k, g + """) + mv_rewrite_fail(combinedQuery, "finalize_max_mv") + mv_rewrite_fail(plainQuery, "finalize_max_mv") + order_qt_different_aggregate combinedQuery + + // The full value expression matters even when both outer aggregates are SUM. + create_async_mv(db, "finalize_abs_mv", """ + select k, g, sum_combine(abs(avg_finalize(avg_state(v)))) s + from agg_finalize_rollup_input group by k, g + """) + mv_rewrite_fail(combinedQuery, "finalize_abs_mv") + mv_rewrite_fail(plainQuery, "finalize_abs_mv") + order_qt_different_argument combinedQuery + + create_async_mv(db, "finalize_sum_mv", """ + select k, g, sum_combine(avg_finalize(avg_state(v))) s + from agg_finalize_rollup_input group by k, g + """) + mv_rewrite_success(combinedQuery, "finalize_sum_mv") + mv_rewrite_success(plainQuery, "finalize_sum_mv") + order_qt_combined_after combinedQuery + order_qt_plain_after plainQuery + + // Existing MERGE/UNION/STATE rollup remains available. + create_async_mv(db, "finalize_union_mv", """ + select k, g, sum_union(sum_state(v)) s + from agg_finalize_rollup_input group by k, g + """) + String mergeQuery = """ + select k, sum_merge(sum_state(v)) + from agg_finalize_rollup_input group by k + """ + mv_rewrite_success(mergeQuery, "finalize_union_mv") + order_qt_merge_after mergeQuery +}