From fa7d9f60115e414e8fa227489c94ba70c74f0169 Mon Sep 17 00:00:00 2001 From: lizhuoyu5 Date: Wed, 29 Jul 2026 20:50:54 +0800 Subject: [PATCH] [feature](function) add uniq_theta theta-sketch distinct-count aggregate Add uniq_theta, a theta-sketch based approximate distinct-count aggregate returning BIGINT, ported from ClickHouse uniqTheta and backed by the Apache DataSketches library (contrib/datasketches-cpp submodule added in #63143). The aggregate reuses the generic agg_state combinator framework so uniq_theta_state / uniq_theta_merge / uniq_theta_union are auto-derived with no extra wiring. Includes BE gtest, FE unit test, and regression suites. Signed-off-by: lizhuoyu5 --- .../aggregate_function_simple_factory.cpp | 2 + .../aggregate_function_uniq_theta.cpp | 48 ++++ .../aggregate/aggregate_function_uniq_theta.h | 217 ++++++++++++++++++ .../exprs/aggregate/agg_uniq_theta_test.cpp | 214 +++++++++++++++++ .../catalog/BuiltinAggregateFunctions.java | 2 + .../expressions/functions/agg/UniqTheta.java | 92 ++++++++ .../visitor/AggregateFunctionVisitor.java | 5 + .../functions/agg/UniqThetaTest.java | 151 ++++++++++++ .../test_agg_state_uniq_theta.groovy | 96 ++++++++ .../test_uniq_theta.groovy | 82 +++++++ 10 files changed, 909 insertions(+) create mode 100644 be/src/exprs/aggregate/aggregate_function_uniq_theta.cpp create mode 100644 be/src/exprs/aggregate/aggregate_function_uniq_theta.h create mode 100644 be/test/exprs/aggregate/agg_uniq_theta_test.cpp create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqTheta.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqThetaTest.java create mode 100644 regression-test/suites/datatype_p0/agg_state/uniq_theta/test_agg_state_uniq_theta.groovy create mode 100644 regression-test/suites/query_p0/sql_functions/aggregate_functions/test_uniq_theta.groovy diff --git a/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp b/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp index 2d1a94dc04c875..1249542a07a97b 100644 --- a/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp +++ b/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp @@ -51,6 +51,7 @@ void register_aggregate_function_stddev_variance_pop(AggregateFunctionSimpleFact void register_aggregate_function_stddev_variance_samp(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_topn(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_approx_count_distinct(AggregateFunctionSimpleFactory& factory); +void register_aggregate_function_uniq_theta(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_group_array_set_op(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_group_concat(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_percentile(AggregateFunctionSimpleFactory& factory); @@ -108,6 +109,7 @@ AggregateFunctionSimpleFactory& AggregateFunctionSimpleFactory::instance() { register_aggregate_function_stddev_variance_pop(instance); register_aggregate_function_topn(instance); register_aggregate_function_approx_count_distinct(instance); + register_aggregate_function_uniq_theta(instance); register_aggregate_function_percentile(instance); register_aggregate_function_percentile_old(instance); register_aggregate_function_percentile_approx(instance); diff --git a/be/src/exprs/aggregate/aggregate_function_uniq_theta.cpp b/be/src/exprs/aggregate/aggregate_function_uniq_theta.cpp new file mode 100644 index 00000000000000..31118f21222ff9 --- /dev/null +++ b/be/src/exprs/aggregate/aggregate_function_uniq_theta.cpp @@ -0,0 +1,48 @@ +// 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. +// This file is copied from +// https://github.com/ClickHouse/ClickHouse/blob/master/src/AggregateFunctions/AggregateFunctionUniq.cpp +// and modified by Doris + +#include "exprs/aggregate/aggregate_function_uniq_theta.h" + +#include "core/data_type/data_type.h" +#include "core/data_type/define_primitive_type.h" +#include "exprs/aggregate/helpers.h" + +namespace doris { +#include "common/compile_check_begin.h" + +AggregateFunctionPtr create_aggregate_function_uniq_theta(const std::string& name, + const DataTypes& argument_types, + const DataTypePtr& result_type, + const bool result_is_nullable, + const AggregateFunctionAttr& attr) { + return creator_with_type_list< + TYPE_BOOLEAN, TYPE_TINYINT, TYPE_SMALLINT, TYPE_INT, TYPE_BIGINT, TYPE_LARGEINT, + TYPE_FLOAT, TYPE_DOUBLE, TYPE_DECIMALV2, TYPE_DECIMAL32, TYPE_DECIMAL64, + TYPE_DECIMAL128I, TYPE_DECIMAL256, TYPE_VARCHAR, TYPE_DATEV2, TYPE_DATETIMEV2, TYPE_IPV4, + TYPE_IPV6, TYPE_TIMESTAMPTZ>::create(argument_types, + result_is_nullable, + attr); +} + +void register_aggregate_function_uniq_theta(AggregateFunctionSimpleFactory& factory) { + factory.register_function_both("uniq_theta", create_aggregate_function_uniq_theta); +} + +} // namespace doris diff --git a/be/src/exprs/aggregate/aggregate_function_uniq_theta.h b/be/src/exprs/aggregate/aggregate_function_uniq_theta.h new file mode 100644 index 00000000000000..20c8f903476b4d --- /dev/null +++ b/be/src/exprs/aggregate/aggregate_function_uniq_theta.h @@ -0,0 +1,217 @@ +// 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. +// This file is copied from +// https://github.com/ClickHouse/ClickHouse/blob/master/src/AggregateFunctions/ThetaSketchData.h +// and modified by Doris + +#pragma once + +#include +#include + +#include +#include +#include + +// roaring/portability.h (pulled in via the precompiled header) defines +// IS_BIG_ENDIAN as a function-like macro, which collides with the enum member +// of the same name in datasketches theta_sketch.hpp. This TU does not use +// roaring's macro, so drop it before including the sketch headers. +#ifdef IS_BIG_ENDIAN +#undef IS_BIG_ENDIAN +#endif +#include +#include + +#include "common/exception.h" +#include "common/status.h" +#include "core/assert_cast.h" +#include "core/column/column_decimal.h" +#include "core/data_type/data_type_number.h" +#include "core/string_ref.h" +#include "core/types.h" +#include "exprs/aggregate/aggregate_function.h" +#include "exprs/aggregate/aggregate_function_simple_factory.h" + +namespace doris { +#include "common/compile_check_begin.h" +class Arena; +class BufferReadable; +class BufferWritable; +class IColumn; +} // namespace doris + +namespace doris { + +// Mirrors ClickHouse ThetaSketchData: sk_update takes raw inserts; once a merge +// happens the accumulated sk_update is folded into sk_union and reset. The two +// datasketches types are distinct and cannot be unified. +struct AggregateFunctionUniqThetaData { + std::unique_ptr sk_update; + std::unique_ptr sk_union; + + AggregateFunctionUniqThetaData() = default; + ~AggregateFunctionUniqThetaData() = default; + AggregateFunctionUniqThetaData(const AggregateFunctionUniqThetaData&) = delete; + AggregateFunctionUniqThetaData& operator=(const AggregateFunctionUniqThetaData&) = delete; + +private: + datasketches::update_theta_sketch* get_sk_update() { + if (!sk_update) { + sk_update = std::make_unique( + datasketches::update_theta_sketch::builder().build()); + } + return sk_update.get(); + } + + datasketches::theta_union* get_sk_union() { + if (!sk_union) { + sk_union = std::make_unique( + datasketches::theta_union::builder().build()); + } + return sk_union.get(); + } + +public: + // Feed raw value bytes; datasketches hashes internally. + void add(const void* data_ptr, size_t len) { + get_sk_update()->update(data_ptr, len); + // add() may be called after a merge() (e.g. incremental aggregation); keep + // sk_union authoritative so those inserts are not lost by get()/serialize(). + if (sk_union) { + sk_union->update(*sk_update); + sk_update.reset(); + } + } + + void merge(const AggregateFunctionUniqThetaData& rhs) { + datasketches::theta_union* u = get_sk_union(); + if (sk_update) { + u->update(*sk_update); + sk_update.reset(); + } + if (rhs.sk_update) { + u->update(*rhs.sk_update); + } else if (rhs.sk_union) { + u->update(rhs.sk_union->get_result()); + } + } + + int64_t get() const { + if (sk_union) { + return static_cast(sk_union->get_result().get_estimate()); + } + if (sk_update) { + return static_cast(sk_update->get_estimate()); + } + return 0; + } + + void write(BufferWritable& buf) const { + std::string bytes; + if (sk_update) { + auto vec = sk_update->compact().serialize(); + bytes.assign(reinterpret_cast(vec.data()), vec.size()); + } else if (sk_union) { + auto vec = sk_union->get_result().serialize(); + bytes.assign(reinterpret_cast(vec.data()), vec.size()); + } + // Empty state serializes to a zero-length blob; never serialize() an empty + // sketch (that would emit an 8-byte preamble instead). + buf.write_binary(bytes); + } + + void read(BufferReadable& buf) { + std::string bytes; + buf.read_binary(bytes); + if (bytes.empty()) { + return; + } + try { + auto sk = datasketches::compact_theta_sketch::deserialize(bytes.data(), bytes.size()); + get_sk_union()->update(sk); + } catch (const std::bad_alloc&) { + throw; + } catch (const std::exception& e) { + throw Exception(ErrorCode::INTERNAL_ERROR, + "uniq_theta: cannot deserialize theta sketch: {}", e.what()); + } + } + + void reset() { + sk_update.reset(); + sk_union.reset(); + } +}; + +template +class AggregateFunctionUniqTheta final + : public IAggregateFunctionDataHelper>, + UnaryExpression, + NotNullableAggregateFunction { +public: + using ColumnDataType = typename PrimitiveTypeTraits::ColumnType; + String get_name() const override { return "uniq_theta"; } + + AggregateFunctionUniqTheta(const DataTypes& argument_types_) + : IAggregateFunctionDataHelper>(argument_types_) {} + + DataTypePtr get_return_type() const override { return std::make_shared(); } + + void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num, + Arena&) const override { + if constexpr (is_decimal(type) || is_int_or_bool(type) || is_ip(type) || + is_date_type(type) || is_timestamptz_type(type) || is_float_or_double(type) || + type == TYPE_TIMEV2) { + auto column = + assert_cast(columns[0]); + auto value = column->get_element(row_num); + this->data(place).add(&value, sizeof(value)); + } else { + auto value = assert_cast(columns[0]) + ->get_data_at(row_num); + this->data(place).add(value.data, value.size); + } + } + + void reset(AggregateDataPtr place) const override { this->data(place).reset(); } + + void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, + Arena&) const override { + this->data(place).merge(this->data(rhs)); + } + + void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override { + this->data(place).write(buf); + } + + void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf, + Arena&) const override { + this->data(place).read(buf); + } + + void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override { + auto& column = assert_cast(to); + column.get_data().push_back(this->data(place).get()); + } +}; + +} // namespace doris + +#include "common/compile_check_end.h" diff --git a/be/test/exprs/aggregate/agg_uniq_theta_test.cpp b/be/test/exprs/aggregate/agg_uniq_theta_test.cpp new file mode 100644 index 00000000000000..180013911ac7ca --- /dev/null +++ b/be/test/exprs/aggregate/agg_uniq_theta_test.cpp @@ -0,0 +1,214 @@ +// 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 + +#include +#include +#include + +#include "core/arena.h" +#include "core/column/column_string.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type.h" +#include "core/data_type/data_type_number.h" +#include "core/data_type/data_type_string.h" +#include "core/string_buffer.hpp" +#include "exprs/aggregate/aggregate_function.h" +#include "exprs/aggregate/aggregate_function_simple_factory.h" + +namespace doris { + +void register_aggregate_function_uniq_theta(AggregateFunctionSimpleFactory& factory); + +class VAggUniqThetaTest : public testing::Test { +public: + void SetUp() override { + AggregateFunctionSimpleFactory factory = AggregateFunctionSimpleFactory::instance(); + register_aggregate_function_uniq_theta(factory); + } + void TearDown() override {} + + AggregateFunctionPtr get_int64_fn() { + DataTypes types = {std::make_shared()}; + auto fn = AggregateFunctionSimpleFactory::instance().get("uniq_theta", types, nullptr, + false, -1); + EXPECT_NE(fn, nullptr); + return fn; + } + + // Build a place holding the int64 values [begin, end). + struct Place { + AggregateFunctionPtr fn; + std::unique_ptr mem; + AggregateDataPtr ptr = nullptr; + Place(AggregateFunctionPtr f) : fn(std::move(f)) { + mem.reset(new char[fn->size_of_data()]); + ptr = mem.get(); + fn->create(ptr); + } + ~Place() { fn->destroy(ptr); } + }; + + void add_int64_range(const AggregateFunctionPtr& fn, AggregateDataPtr place, int64_t begin, + int64_t end) { + auto col = ColumnInt64::create(); + for (int64_t v = begin; v < end; ++v) { + col->insert_value(v); + } + const IColumn* cols[1] = {col.get()}; + for (size_t i = 0; i < col->size(); ++i) { + fn->add(place, cols, i, _arena); + } + } + + int64_t result_of(const AggregateFunctionPtr& fn, AggregateDataPtr place) { + auto out = ColumnInt64::create(); + fn->insert_result_into(place, *out); + EXPECT_EQ(out->size(), 1); + return out->get_data()[0]; + } + + Arena _arena; +}; + +TEST_F(VAggUniqThetaTest, empty_state) { + auto fn = get_int64_fn(); + Place p(fn); + EXPECT_EQ(result_of(fn, p.ptr), 0); +} + +TEST_F(VAggUniqThetaTest, small_exact) { + auto fn = get_int64_fn(); + Place p(fn); + add_int64_range(fn, p.ptr, 1, 6); // 1,2,3,4,5 + EXPECT_EQ(result_of(fn, p.ptr), 5); +} + +TEST_F(VAggUniqThetaTest, duplicates_deduped) { + auto fn = get_int64_fn(); + Place p(fn); + auto col = ColumnInt64::create(); + for (int64_t v : {1, 1, 2, 2, 3}) { + col->insert_value(v); + } + const IColumn* cols[1] = {col.get()}; + for (size_t i = 0; i < col->size(); ++i) { + fn->add(p.ptr, cols, i, _arena); + } + EXPECT_EQ(result_of(fn, p.ptr), 3); +} + +TEST_F(VAggUniqThetaTest, large_cardinality_within_error_bound) { + auto fn = get_int64_fn(); + Place p(fn); + const int64_t n = 100000; + add_int64_range(fn, p.ptr, 0, n); + int64_t est = result_of(fn, p.ptr); + // Theta sketch relative error ~3.125% at 95% confidence; allow a small margin. + EXPECT_LE(std::abs(est - n), static_cast(n * 0.04)); +} + +TEST_F(VAggUniqThetaTest, serialize_deserialize_roundtrip) { + auto fn = get_int64_fn(); + Place p(fn); + add_int64_range(fn, p.ptr, 0, 10); + int64_t before = result_of(fn, p.ptr); + + ColumnString buf; + VectorBufferWriter writer(buf); + fn->serialize(p.ptr, writer); + writer.commit(); + + Place p2(fn); + VectorBufferReader reader(buf.get_data_at(0)); + fn->deserialize(p2.ptr, reader, _arena); + EXPECT_EQ(result_of(fn, p2.ptr), before); +} + +TEST_F(VAggUniqThetaTest, merge_disjoint) { + auto fn = get_int64_fn(); + Place a(fn); + Place b(fn); + add_int64_range(fn, a.ptr, 1, 4); // 1,2,3 + add_int64_range(fn, b.ptr, 4, 7); // 4,5,6 + fn->merge(a.ptr, b.ptr, _arena); + EXPECT_EQ(result_of(fn, a.ptr), 6); +} + +TEST_F(VAggUniqThetaTest, merge_overlapping) { + auto fn = get_int64_fn(); + Place a(fn); + Place b(fn); + add_int64_range(fn, a.ptr, 1, 5); // 1,2,3,4 + add_int64_range(fn, b.ptr, 3, 7); // 3,4,5,6 + fn->merge(a.ptr, b.ptr, _arena); + EXPECT_EQ(result_of(fn, a.ptr), 6); +} + +TEST_F(VAggUniqThetaTest, add_after_merge_not_lost) { + auto fn = get_int64_fn(); + Place a(fn); + Place b(fn); + add_int64_range(fn, a.ptr, 1, 4); // 1,2,3 + add_int64_range(fn, b.ptr, 4, 7); // 4,5,6 -> forces sk_union on a + fn->merge(a.ptr, b.ptr, _arena); + add_int64_range(fn, a.ptr, 7, 9); // 7,8 added after union exists + EXPECT_EQ(result_of(fn, a.ptr), 8); +} + +TEST_F(VAggUniqThetaTest, string_type) { + DataTypes types = {std::make_shared()}; + auto fn = AggregateFunctionSimpleFactory::instance().get("uniq_theta", types, nullptr, false, + -1); + ASSERT_NE(fn, nullptr); + std::unique_ptr mem(new char[fn->size_of_data()]); + AggregateDataPtr place = mem.get(); + fn->create(place); + + auto col = ColumnString::create(); + for (const auto* s : {"a", "b", "c", "a", "d"}) { + col->insert_data(s, strlen(s)); + } + const IColumn* cols[1] = {col.get()}; + for (size_t i = 0; i < col->size(); ++i) { + fn->add(place, cols, i, _arena); + } + auto out = ColumnInt64::create(); + fn->insert_result_into(place, *out); + EXPECT_EQ(out->get_data()[0], 4); + fn->destroy(place); +} + +TEST_F(VAggUniqThetaTest, streaming_state_roundtrip) { + // Exercises the _state path: one serialized sketch per row, then merge them all. + auto fn = get_int64_fn(); + auto col = ColumnInt64::create(); + for (int64_t v : {10, 20, 20, 30, 40}) { + col->insert_value(v); + } + const IColumn* cols[1] = {col.get()}; + MutableColumnPtr dst = ColumnString::create(); + fn->streaming_agg_serialize_to_column(cols, dst, col->size(), _arena); + EXPECT_EQ(dst->size(), col->size()); + + Place merged(fn); + fn->deserialize_and_merge_from_column(merged.ptr, *dst, _arena); + EXPECT_EQ(result_of(fn, merged.ptr), 4); +} + +} // namespace doris diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java index 1f60a228c11b2c..55f1e9d6dea4be 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java @@ -103,6 +103,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.TopN; import org.apache.doris.nereids.trees.expressions.functions.agg.TopNArray; import org.apache.doris.nereids.trees.expressions.functions.agg.TopNWeighted; +import org.apache.doris.nereids.trees.expressions.functions.agg.UniqTheta; import org.apache.doris.nereids.trees.expressions.functions.agg.Variance; import org.apache.doris.nereids.trees.expressions.functions.agg.VarianceSamp; import org.apache.doris.nereids.trees.expressions.functions.agg.WindowFunnel; @@ -182,6 +183,7 @@ private BuiltinAggregateFunctions() { agg(MultiDistinctSum.class, "multi_distinct_sum"), agg(MultiDistinctSum0.class, "multi_distinct_sum0"), agg(Ndv.class, "approx_count_distinct", "ndv"), + agg(UniqTheta.class, "uniq_theta"), agg(OrthogonalBitmapExprCalculate.class, "orthogonal_bitmap_expr_calculate"), agg(OrthogonalBitmapExprCalculateCount.class, "orthogonal_bitmap_expr_calculate_count"), agg(OrthogonalBitmapIntersect.class, "orthogonal_bitmap_intersect"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqTheta.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqTheta.java new file mode 100644 index 00000000000000..6dc60c3f590c9c --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqTheta.java @@ -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. + +package org.apache.doris.nereids.trees.expressions.functions.agg; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.catalog.Type; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; +import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.BigIntType; +import org.apache.doris.nereids.types.coercion.AnyDataType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * AggregateFunction 'uniq_theta'. Theta-sketch based approximate distinct count. + */ +public class UniqTheta extends NotNullableAggregateFunction + implements UnaryExpression, ExplicitlyCastableSignature { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(BigIntType.INSTANCE).args(AnyDataType.INSTANCE_WITHOUT_INDEX) + ); + + /** + * constructor with 1 argument. + */ + public UniqTheta(Expression arg) { + super("uniq_theta", arg); + } + + public UniqTheta(boolean distinct, Expression arg) { + this(arg); + } + + /** constructor for withChildren and reuse signature */ + private UniqTheta(AggregateFunctionParams functionParams) { + super(functionParams); + } + + @Override + public void checkLegalityBeforeTypeCoercion() { + if (getArgumentType(0).isOnlyMetricType()) { + throw new AnalysisException(Type.OnlyMetricTypeErrorMsg); + } + } + + /** + * withDistinctAndChildren. + */ + @Override + public UniqTheta withDistinctAndChildren(boolean distinct, List children) { + Preconditions.checkArgument(children.size() == 1); + return new UniqTheta(getFunctionParams(distinct, children)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitUniqTheta(this, context); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } + + @Override + public Expression resultForEmptyInput() { + return new BigIntLiteral(0); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java index 81b0830be66c81..910fd9289e847d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java @@ -96,6 +96,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.TopN; import org.apache.doris.nereids.trees.expressions.functions.agg.TopNArray; import org.apache.doris.nereids.trees.expressions.functions.agg.TopNWeighted; +import org.apache.doris.nereids.trees.expressions.functions.agg.UniqTheta; import org.apache.doris.nereids.trees.expressions.functions.agg.Variance; import org.apache.doris.nereids.trees.expressions.functions.agg.VarianceSamp; import org.apache.doris.nereids.trees.expressions.functions.agg.WindowFunnel; @@ -299,6 +300,10 @@ default R visitNdv(Ndv ndv, C context) { return visitAggregateFunction(ndv, context); } + default R visitUniqTheta(UniqTheta uniqTheta, C context) { + return visitAggregateFunction(uniqTheta, context); + } + default R visitOrthogonalBitmapIntersect(OrthogonalBitmapIntersect function, C context) { return visitAggregateFunction(function, context); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqThetaTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqThetaTest.java new file mode 100644 index 00000000000000..2e3ea6b676bed3 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/UniqThetaTest.java @@ -0,0 +1,151 @@ +// 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.agg; + +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; +import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionVisitor; +import org.apache.doris.nereids.types.BigIntType; +import org.apache.doris.nereids.types.HllType; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.VarcharType; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Direct expression unit tests for {@link UniqTheta} (no FE service needed). + */ +public class UniqThetaTest { + + private final SlotReference intSlot = new SlotReference("col_int", IntegerType.INSTANCE); + private final SlotReference strSlot = new SlotReference("col_str", VarcharType.SYSTEM_DEFAULT); + private final SlotReference hllSlot = new SlotReference("col_hll", HllType.INSTANCE); + + @Test + public void testOneArgConstructor() { + UniqTheta func = new UniqTheta(intSlot); + Assertions.assertEquals("uniq_theta", func.getName()); + Assertions.assertEquals(1, func.arity()); + Assertions.assertEquals(intSlot, func.getArgument(0)); + Assertions.assertFalse(func.isDistinct()); + Assertions.assertEquals(BigIntType.INSTANCE, func.getDataType()); + } + + @Test + public void testDistinctConstructorIgnoresFlag() { + UniqTheta func = new UniqTheta(true, intSlot); + Assertions.assertEquals("uniq_theta", func.getName()); + Assertions.assertEquals(intSlot, func.getArgument(0)); + // distinct is absorbed (theta is already an approximate distinct) + Assertions.assertFalse(func.isDistinct()); + } + + @Test + public void testNotNullable() { + UniqTheta func = new UniqTheta(intSlot); + Assertions.assertFalse(func.nullable()); + } + + @Test + public void testGetSignatures() { + UniqTheta func = new UniqTheta(intSlot); + Assertions.assertEquals(1, func.getSignatures().size()); + Assertions.assertSame(UniqTheta.SIGNATURES, func.getSignatures()); + Assertions.assertEquals(BigIntType.INSTANCE, func.getSignatures().get(0).returnType); + } + + @Test + public void testStaticSignatures() { + Assertions.assertEquals(1, UniqTheta.SIGNATURES.size()); + Assertions.assertEquals(BigIntType.INSTANCE, UniqTheta.SIGNATURES.get(0).returnType); + Assertions.assertEquals(1, UniqTheta.SIGNATURES.get(0).argumentsTypes.size()); + } + + @Test + public void testWithDistinctAndChildren() { + UniqTheta func = new UniqTheta(intSlot); + UniqTheta replaced = func.withDistinctAndChildren(false, ImmutableList.of(strSlot)); + Assertions.assertEquals(1, replaced.arity()); + Assertions.assertEquals(strSlot, replaced.getArgument(0)); + Assertions.assertEquals(BigIntType.INSTANCE, replaced.getDataType()); + } + + @Test + public void testWithChildren() { + UniqTheta func = new UniqTheta(intSlot); + Expression replaced = func.withChildren(ImmutableList.of(strSlot)); + Assertions.assertTrue(replaced instanceof UniqTheta); + Assertions.assertEquals(strSlot, ((UniqTheta) replaced).getArgument(0)); + } + + @Test + public void testWithDistinctAndChildrenIllegalArity() { + UniqTheta func = new UniqTheta(intSlot); + Assertions.assertThrows(IllegalArgumentException.class, + () -> func.withDistinctAndChildren(false, ImmutableList.of(intSlot, strSlot))); + } + + @Test + public void testResultForEmptyInput() { + UniqTheta func = new UniqTheta(intSlot); + Expression empty = func.resultForEmptyInput(); + Assertions.assertTrue(empty instanceof BigIntLiteral); + Assertions.assertEquals(0L, ((BigIntLiteral) empty).getValue().longValue()); + } + + @Test + public void testCheckLegalityAcceptsNonMetricType() { + UniqTheta func = new UniqTheta(intSlot); + Assertions.assertDoesNotThrow(func::checkLegalityBeforeTypeCoercion); + } + + @Test + public void testCheckLegalityRejectsMetricType() { + UniqTheta func = new UniqTheta(hllSlot); + Assertions.assertThrows(AnalysisException.class, func::checkLegalityBeforeTypeCoercion); + } + + @Test + public void testAcceptVisitorDispatch() { + UniqTheta func = new UniqTheta(intSlot); + DefaultExpressionVisitor visitor = new DefaultExpressionVisitor() { + @Override + public String visitUniqTheta(UniqTheta f, Void ctx) { + return "visited_uniq_theta"; + } + }; + Assertions.assertEquals("visited_uniq_theta", func.accept(visitor, null)); + } + + @Test + public void testAcceptVisitorDefaultDispatch() { + // Exercises the default visitUniqTheta -> visitAggregateFunction path. + UniqTheta func = new UniqTheta(intSlot); + DefaultExpressionVisitor visitor = new DefaultExpressionVisitor() { + @Override + public String visitAggregateFunction(AggregateFunction f, Void ctx) { + return "visited_agg_function"; + } + }; + Assertions.assertEquals("visited_agg_function", func.accept(visitor, null)); + } +} diff --git a/regression-test/suites/datatype_p0/agg_state/uniq_theta/test_agg_state_uniq_theta.groovy b/regression-test/suites/datatype_p0/agg_state/uniq_theta/test_agg_state_uniq_theta.groovy new file mode 100644 index 00000000000000..912b1f6e26b300 --- /dev/null +++ b/regression-test/suites/datatype_p0/agg_state/uniq_theta/test_agg_state_uniq_theta.groovy @@ -0,0 +1,96 @@ +// 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_uniq_theta") { + sql "set enable_agg_state=true" + + sql "DROP TABLE IF EXISTS uniq_theta_src" + sql """ + CREATE TABLE uniq_theta_src ( + k1 int null, + v_int int null, + v_str varchar(64) null + ) + DUPLICATE KEY(k1) + DISTRIBUTED BY HASH(k1) BUCKETS 3 + PROPERTIES("replication_num" = "1") + """ + sql """ + INSERT INTO uniq_theta_src VALUES + (1, 10, 'a'), (1, 20, 'b'), (1, 20, 'b'), (1, null, null), + (2, 30, 'c'), (2, 40, 'd'), (2, 50, 'e') + """ + + // state -> merge round-trip equals direct uniq_theta (nulls excluded) + qt_state_merge_basic """ + SELECT uniq_theta_merge(uniq_theta_state(v_int)) FROM uniq_theta_src + """ + qt_direct """ + SELECT uniq_theta(v_int) FROM uniq_theta_src + """ + + // persist state into an AGGREGATE table, then merge per key. + // Note: INSERT ... SELECT ... GROUP BY into an agg_state column is not + // supported by the planner (same limitation as ndv_state), so use + // value-style inserts like datatype_p0/agg_state/hll/hll.groovy. + sql "DROP TABLE IF EXISTS uniq_theta_agg" + sql """ + CREATE TABLE uniq_theta_agg ( + k1 int null, + s agg_state generic + ) + AGGREGATE KEY(k1) + DISTRIBUTED BY HASH(k1) BUCKETS 3 + PROPERTIES("replication_num" = "1") + """ + // key 1 -> {10, 20, 20} = 2 distinct; multi-insert same key auto-merges state + sql "INSERT INTO uniq_theta_agg VALUES (1, uniq_theta_state(cast(10 as int)))" + sql "INSERT INTO uniq_theta_agg VALUES (1, uniq_theta_state(cast(20 as int)))" + sql "INSERT INTO uniq_theta_agg VALUES (1, uniq_theta_state(cast(20 as int)))" + // key 2 -> {30, 40} = 2 distinct + sql "INSERT INTO uniq_theta_agg VALUES (2, uniq_theta_state(cast(30 as int)))" + sql "INSERT INTO uniq_theta_agg VALUES (2, uniq_theta_state(cast(40 as int)))" + + qt_agg_merge_by_key """ + SELECT k1, uniq_theta_merge(s) FROM uniq_theta_agg GROUP BY k1 ORDER BY k1 + """ + + // union across partitions then merge to global cardinality (4 distinct total) + qt_union_then_merge """ + SELECT uniq_theta_merge(us) FROM ( + SELECT uniq_theta_union(s) AS us FROM uniq_theta_agg GROUP BY k1 + ) t + """ + + // string state column + sql "DROP TABLE IF EXISTS uniq_theta_agg_str" + sql """ + CREATE TABLE uniq_theta_agg_str ( + k1 int null, + s agg_state generic + ) + AGGREGATE KEY(k1) + DISTRIBUTED BY HASH(k1) BUCKETS 3 + PROPERTIES("replication_num" = "1") + """ + sql "INSERT INTO uniq_theta_agg_str VALUES (1, uniq_theta_state(cast('a' as varchar(64))))" + sql "INSERT INTO uniq_theta_agg_str VALUES (1, uniq_theta_state(cast('b' as varchar(64))))" + sql "INSERT INTO uniq_theta_agg_str VALUES (2, uniq_theta_state(cast('c' as varchar(64))))" + qt_str_state_merge """ + SELECT uniq_theta_merge(s) FROM uniq_theta_agg_str + """ +} diff --git a/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_uniq_theta.groovy b/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_uniq_theta.groovy new file mode 100644 index 00000000000000..0fd939017e3710 --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_uniq_theta.groovy @@ -0,0 +1,82 @@ +// 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_uniq_theta") { + sql "DROP TABLE IF EXISTS test_uniq_theta" + sql """ + CREATE TABLE test_uniq_theta ( + id int, + grp int, + v_int int, + v_bigint bigint, + v_str varchar(64), + v_date date, + v_dec decimal(10, 2), + v_null int + ) + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 3 + PROPERTIES("replication_num" = "1") + """ + + // empty table -> 0 + qt_empty "SELECT uniq_theta(v_int) FROM test_uniq_theta" + + sql """ + INSERT INTO test_uniq_theta VALUES + (1, 1, 10, 1000, 'aaa', '2024-01-01', 1.10, 7), + (2, 1, 20, 2000, 'bbb', '2024-01-02', 2.20, null), + (3, 1, 30, 3000, 'ccc', '2024-01-03', 3.30, 7), + (4, 2, 10, 1000, 'aaa', '2024-01-01', 1.10, null), + (5, 2, 40, 4000, 'ddd', '2024-01-04', 4.40, 8), + (6, 2, 50, 5000, 'eee', '2024-01-05', 5.50, 9) + """ + + // exact at small cardinality across types + qt_int "SELECT uniq_theta(v_int) FROM test_uniq_theta" + qt_bigint "SELECT uniq_theta(v_bigint) FROM test_uniq_theta" + qt_str "SELECT uniq_theta(v_str) FROM test_uniq_theta" + qt_date "SELECT uniq_theta(v_date) FROM test_uniq_theta" + qt_dec "SELECT uniq_theta(v_dec) FROM test_uniq_theta" + + // consistent with count(distinct) on small data + qt_vs_count_distinct """ + SELECT uniq_theta(v_int), count(distinct v_int) FROM test_uniq_theta + """ + // consistent with approx_count_distinct on small data + qt_vs_ndv """ + SELECT uniq_theta(v_str), approx_count_distinct(v_str) FROM test_uniq_theta + """ + + // nulls excluded; column has 3 distinct non-null values (7,8,9) + qt_nulls "SELECT uniq_theta(v_null) FROM test_uniq_theta" + + // group by + qt_group_by """ + SELECT grp, uniq_theta(v_int) FROM test_uniq_theta GROUP BY grp ORDER BY grp + """ + + // all-null column -> 0 + sql "DROP TABLE IF EXISTS test_uniq_theta_allnull" + sql """ + CREATE TABLE test_uniq_theta_allnull (id int, v int) + DUPLICATE KEY(id) DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES("replication_num" = "1") + """ + sql "INSERT INTO test_uniq_theta_allnull VALUES (1, null), (2, null)" + qt_all_null "SELECT uniq_theta(v) FROM test_uniq_theta_allnull" +}