Skip to main content

common_function/aggrs/
approximate.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use datafusion_expr::AggregateUDF;
16use datatypes::arrow::datatypes::DataType;
17
18use crate::aggrs::aggr_wrapper::DeltaMergeWrapper;
19use crate::function_registry::FunctionRegistry;
20
21pub mod avg;
22pub mod hll;
23pub mod uddsketch;
24pub mod welford;
25
26pub(crate) struct ApproximateFunction;
27
28impl ApproximateFunction {
29    pub fn register(registry: &FunctionRegistry) {
30        let avg_merge = avg::AvgAccumulator::merge_udf_impl();
31        registry.register_aggr(avg::AvgAccumulator::state_udf_impl());
32        registry.register_aggr(avg_merge.clone());
33        registry.register_aggr(AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
34            avg_merge.clone(),
35            avg::AVG_STATE_NAME,
36            vec![DataType::Binary],
37            DataType::Binary,
38        )));
39
40        let uddsketch_state = uddsketch::UddSketchState::state_udf_impl();
41        let uddsketch_merge = uddsketch::UddSketchState::merge_udf_impl();
42        let uddsketch_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
43            uddsketch_merge.clone(),
44            uddsketch::UDDSKETCH_STATE_NAME,
45            vec![DataType::Int64, DataType::Float64, DataType::Binary],
46            DataType::Binary,
47        ));
48        registry.register_aggr(uddsketch_state);
49        registry.register_aggr(uddsketch_merge);
50        registry.register_aggr(uddsketch_delta);
51
52        // hll
53        let hll_state = hll::HllState::state_udf_impl();
54        let hll_merge = hll::HllState::merge_udf_impl();
55        let hll_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
56            hll_merge.clone(),
57            hll::HLL_NAME,
58            vec![DataType::Binary],
59            DataType::Binary,
60        ));
61        registry.register_aggr(hll_state);
62        registry.register_aggr(hll_merge);
63        registry.register_aggr(hll_delta);
64
65        // welford
66        let welford_state = welford::WelfordAccumulator::state_udf_impl();
67        let welford_merge = welford::WelfordAccumulator::merge_udf_impl();
68        let welford_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
69            welford_merge.clone(),
70            welford::STDDEV_POP_STATE_NAME,
71            vec![DataType::Binary],
72            DataType::Binary,
73        ));
74        registry.register_aggr(welford_state);
75        registry.register_aggr(welford_merge);
76        registry.register_aggr(welford_delta);
77    }
78}