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..6eb813f12ebdd1 --- /dev/null +++ b/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 + +#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); + 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 = [&]() { + 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()(); + } else { + finalize_rows.template operator()(); + } + 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..1b0e4210951a4c --- /dev/null +++ b/be/test/exprs/function/function_agg_state_finalize_test.cpp @@ -0,0 +1,280 @@ +// 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, TrivialStatesResetAcrossOuterNulls) { + auto type = state_type("sum", std::make_shared()); + auto function = type->get_nested_function(); + ASSERT_TRUE(function->is_trivial()); + auto states = type->create_column(); + append_state(function, ColumnHelper::create_column({2, 4}), *states); + states->insert_default(); + append_state(function, ColumnHelper::create_column({}), *states); + append_state(function, ColumnHelper::create_column({-7}), *states); + append_state(function, ColumnHelper::create_column({10}), *states); + ColumnPtr nullable_states = ColumnNullable::create( + std::move(states), ColumnHelper::create_column({0, 1, 0, 0, 0})); + + auto result = finalize(make_nullable(type), nullable_states); + const auto& values = assert_cast(*result).get_nested_column(); + ASSERT_EQ(result->size(), 5); + EXPECT_EQ(values.get_int(0), 6); + EXPECT_TRUE(result->is_null_at(1)); + EXPECT_FALSE(result->is_null_at(2)); + EXPECT_EQ(values.get_int(2), 0); + EXPECT_EQ(values.get_int(3), -7); + EXPECT_EQ(values.get_int(4), 10); + + auto repeated = finalize(make_nullable(type), nullable_states); + for (size_t row = 0; row < result->size(); ++row) { + EXPECT_EQ(result->compare_at(row, row, *repeated, 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); +} + +TEST_F(FunctionAggStateFinalizeTest, VariableLengthStatesWithOuterNulls) { + auto type = state_type("collect_set", std::make_shared(), false); + auto states = type->create_column(); + auto null_map = ColumnUInt8::create(); + const std::string large(16384, 'x'); + for (size_t group = 0; group < 16; ++group) { + // Valid rows exceed the initial arena chunk; outer NULL payloads must be skipped. + states->insert_default(); + null_map->insert_value(1); + append_state(type->get_nested_function(), + ColumnHelper::create_column({large + std::to_string(group)}), + *states); + null_map->insert_value(0); + } + auto result = finalize(make_nullable(type), + ColumnNullable::create(std::move(states), std::move(null_map))); + const auto& arrays = assert_cast( + assert_cast(*result).get_nested_column()); + const auto& strings = assert_cast( + assert_cast(arrays.get_data()).get_nested_column()); + ASSERT_EQ(result->size(), 32); + for (size_t group = 0; group < 16; ++group) { + EXPECT_TRUE(result->is_null_at(group * 2)); + EXPECT_FALSE(result->is_null_at(group * 2 + 1)); + EXPECT_EQ(arrays.get_offsets()[group * 2], group); + EXPECT_EQ(arrays.get_offsets()[group * 2 + 1], group + 1); + EXPECT_EQ(strings.get_data_at(group).to_string(), large + std::to_string(group)); + } +} + +} // 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/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/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..10321d707d9e8f --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/combinator/FinalizeCombinator.java @@ -0,0 +1,100 @@ +// 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(); + // 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())) { + 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/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/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/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/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" + } +} 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 +}