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 hll;
22pub mod uddsketch;
23pub mod welford;
24
25pub(crate) struct ApproximateFunction;
26
27impl ApproximateFunction {
28    pub fn register(registry: &FunctionRegistry) {
29        let uddsketch_state = uddsketch::UddSketchState::state_udf_impl();
30        let uddsketch_merge = uddsketch::UddSketchState::merge_udf_impl();
31        let uddsketch_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
32            uddsketch_merge.clone(),
33            uddsketch::UDDSKETCH_STATE_NAME,
34            vec![DataType::Int64, DataType::Float64, DataType::Binary],
35            DataType::Binary,
36        ));
37        registry.register_aggr(uddsketch_state);
38        registry.register_aggr(uddsketch_merge);
39        registry.register_aggr(uddsketch_delta);
40
41        // hll
42        let hll_state = hll::HllState::state_udf_impl();
43        let hll_merge = hll::HllState::merge_udf_impl();
44        let hll_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
45            hll_merge.clone(),
46            hll::HLL_NAME,
47            vec![DataType::Binary],
48            DataType::Binary,
49        ));
50        registry.register_aggr(hll_state);
51        registry.register_aggr(hll_merge);
52        registry.register_aggr(hll_delta);
53
54        // welford
55        let welford_state = welford::WelfordAccumulator::state_udf_impl();
56        let welford_merge = welford::WelfordAccumulator::merge_udf_impl();
57        let welford_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new(
58            welford_merge.clone(),
59            welford::STDDEV_POP_STATE_NAME,
60            vec![DataType::Binary],
61            DataType::Binary,
62        ));
63        registry.register_aggr(welford_state);
64        registry.register_aggr(welford_merge);
65        registry.register_aggr(welford_delta);
66    }
67}