datafusion_extra_functions/
kurtosis.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use datafusion::arrow::array::{Float64Array, UInt64Array};
19use datafusion::{arrow, common, error, logical_expr, scalar};
20use std::{fmt, mem};
21
22make_udaf_expr_and_func!(
23    KurtosisFunction,
24    kurtosis,
25    x,
26    "Calculates the excess kurtosis (Fisher’s definition) with bias correction according to the sample size.",
27    kurtosis_udaf
28);
29
30#[derive(Eq, Hash, PartialEq)]
31pub struct KurtosisFunction {
32    signature: logical_expr::Signature,
33}
34
35impl fmt::Debug for KurtosisFunction {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        f.debug_struct("KurtosisFunction")
38            .field("signature", &self.signature)
39            .finish()
40    }
41}
42
43impl Default for KurtosisFunction {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl KurtosisFunction {
50    pub fn new() -> Self {
51        Self {
52            signature: logical_expr::Signature::exact(
53                vec![arrow::datatypes::DataType::Float64],
54                logical_expr::Volatility::Immutable,
55            ),
56        }
57    }
58}
59
60impl logical_expr::AggregateUDFImpl for KurtosisFunction {
61    fn name(&self) -> &str {
62        "kurtosis"
63    }
64
65    fn signature(&self) -> &logical_expr::Signature {
66        &self.signature
67    }
68
69    fn return_type(
70        &self,
71        _arg_types: &[arrow::datatypes::DataType],
72    ) -> error::Result<arrow::datatypes::DataType> {
73        Ok(arrow::datatypes::DataType::Float64)
74    }
75
76    fn accumulator(
77        &self,
78        _acc_args: logical_expr::function::AccumulatorArgs,
79    ) -> error::Result<Box<dyn logical_expr::Accumulator>> {
80        Ok(Box::new(KurtosisAccumulator::new()))
81    }
82
83    fn state_fields(
84        &self,
85        _args: logical_expr::function::StateFieldsArgs,
86    ) -> error::Result<Vec<arrow::datatypes::FieldRef>> {
87        Ok(vec![
88            arrow::datatypes::Field::new("count", arrow::datatypes::DataType::UInt64, true).into(),
89            arrow::datatypes::Field::new("sum", arrow::datatypes::DataType::Float64, true).into(),
90            arrow::datatypes::Field::new("sum_sqr", arrow::datatypes::DataType::Float64, true)
91                .into(),
92            arrow::datatypes::Field::new("sum_cub", arrow::datatypes::DataType::Float64, true)
93                .into(),
94            arrow::datatypes::Field::new("sum_four", arrow::datatypes::DataType::Float64, true)
95                .into(),
96        ])
97    }
98}
99
100/// Accumulator for calculating the excess kurtosis (Fisher’s definition) with bias correction according to the sample size.
101/// This implementation follows the [DuckDB implementation]:
102/// <https://github.com/duckdb/duckdb/blob/main/src/core_functions/aggregate/distributive/kurtosis.cpp>
103#[derive(Debug, Default)]
104pub struct KurtosisAccumulator {
105    count: u64,
106    sum: f64,
107    sum_sqr: f64,
108    sum_cub: f64,
109    sum_four: f64,
110}
111
112impl KurtosisAccumulator {
113    pub fn new() -> Self {
114        Self {
115            count: 0,
116            sum: 0.0,
117            sum_sqr: 0.0,
118            sum_cub: 0.0,
119            sum_four: 0.0,
120        }
121    }
122}
123
124impl logical_expr::Accumulator for KurtosisAccumulator {
125    fn update_batch(&mut self, values: &[arrow::array::ArrayRef]) -> error::Result<()> {
126        let array = common::cast::as_float64_array(&values[0])?;
127        for value in array.iter().flatten() {
128            self.count += 1;
129            self.sum += value;
130            self.sum_sqr += value.powi(2);
131            self.sum_cub += value.powi(3);
132            self.sum_four += value.powi(4);
133        }
134        Ok(())
135    }
136
137    fn merge_batch(&mut self, states: &[arrow::array::ArrayRef]) -> error::Result<()> {
138        let counts = common::downcast_value!(states[0], UInt64Array);
139        let sums = common::downcast_value!(states[1], Float64Array);
140        let sum_sqrs = common::downcast_value!(states[2], Float64Array);
141        let sum_cubs = common::downcast_value!(states[3], Float64Array);
142        let sum_fours = common::downcast_value!(states[4], Float64Array);
143
144        for i in 0..counts.len() {
145            let c = counts.value(i);
146            if c == 0 {
147                continue;
148            }
149            self.count += c;
150            self.sum += sums.value(i);
151            self.sum_sqr += sum_sqrs.value(i);
152            self.sum_cub += sum_cubs.value(i);
153            self.sum_four += sum_fours.value(i);
154        }
155
156        Ok(())
157    }
158
159    fn evaluate(&mut self) -> error::Result<scalar::ScalarValue> {
160        if self.count <= 3 {
161            return Ok(scalar::ScalarValue::Float64(None));
162        }
163
164        let count_64 = 1_f64 / self.count as f64;
165        let m4 = count_64
166            * (self.sum_four - 4.0 * self.sum_cub * self.sum * count_64
167                + 6.0 * self.sum_sqr * self.sum.powi(2) * count_64.powi(2)
168                - 3.0 * self.sum.powi(4) * count_64.powi(3));
169
170        let m2 = (self.sum_sqr - self.sum.powi(2) * count_64) * count_64;
171        if m2 <= 0.0 {
172            return Ok(scalar::ScalarValue::Float64(None));
173        }
174
175        let count = self.count as f64;
176        let numerator = (count - 1.0) * ((count + 1.0) * m4 / m2.powi(2) - 3.0 * (count - 1.0));
177        let denominator = (count - 2.0) * (count - 3.0);
178
179        let target = numerator / denominator;
180
181        Ok(scalar::ScalarValue::Float64(Some(target)))
182    }
183
184    fn size(&self) -> usize {
185        mem::size_of_val(self)
186    }
187
188    fn state(&mut self) -> error::Result<Vec<scalar::ScalarValue>> {
189        Ok(vec![
190            scalar::ScalarValue::from(self.count),
191            scalar::ScalarValue::from(self.sum),
192            scalar::ScalarValue::from(self.sum_sqr),
193            scalar::ScalarValue::from(self.sum_cub),
194            scalar::ScalarValue::from(self.sum_four),
195        ])
196    }
197}
datafusion_extra_functions/kurtosis.rs

datafusion_extra_functions/
kurtosis.rs
