datafusion_functions_aggregate/
stddev.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
18//! Defines physical expressions that can evaluated at runtime during query execution
19
20use std::fmt::Debug;
21use std::hash::Hash;
22use std::mem::align_of_val;
23use std::sync::Arc;
24
25use arrow::array::{BooleanArray, Float64Array};
26use arrow::datatypes::FieldRef;
27use arrow::{array::ArrayRef, datatypes::DataType, datatypes::Field};
28use datafusion_common::ScalarValue;
29use datafusion_common::{Result, internal_err, not_impl_err};
30use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
31use datafusion_expr::utils::format_state_name;
32use datafusion_expr::{
33    Accumulator, AggregateUDFImpl, Documentation, GroupsAccumulator, Signature,
34    Volatility,
35};
36use datafusion_functions_aggregate_common::stats::StatsType;
37use datafusion_macros::user_doc;
38
39use crate::variance::{VarianceAccumulator, VarianceGroupsAccumulator};
40
41make_udaf_expr_and_func!(
42    Stddev,
43    stddev,
44    expression,
45    "Compute the standard deviation of a set of numbers",
46    stddev_udaf
47);
48
49#[user_doc(
50    doc_section(label = "Statistical Functions"),
51    description = "Returns the standard deviation of a set of numbers.",
52    syntax_example = "stddev(expression)",
53    sql_example = r#"```sql
54> SELECT stddev(column_name) FROM table_name;
55+----------------------+
56| stddev(column_name)   |
57+----------------------+
58| 12.34                |
59+----------------------+
60```"#,
61    standard_argument(name = "expression",)
62)]
63/// STDDEV and STDDEV_SAMP (standard deviation) aggregate expression
64#[derive(PartialEq, Eq, Hash, Debug)]
65pub struct Stddev {
66    signature: Signature,
67    alias: Vec<String>,
68}
69
70impl Default for Stddev {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl Stddev {
77    /// Create a new STDDEV aggregate function
78    pub fn new() -> Self {
79        Self {
80            signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable),
81            alias: vec!["stddev_samp".to_string()],
82        }
83    }
84}
85
86impl AggregateUDFImpl for Stddev {
87    fn name(&self) -> &str {
88        "stddev"
89    }
90
91    fn signature(&self) -> &Signature {
92        &self.signature
93    }
94
95    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
96        Ok(DataType::Float64)
97    }
98
99    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
100        Ok(vec![
101            Field::new(
102                format_state_name(args.name, "count"),
103                DataType::UInt64,
104                true,
105            ),
106            Field::new(
107                format_state_name(args.name, "mean"),
108                DataType::Float64,
109                true,
110            ),
111            Field::new(format_state_name(args.name, "m2"), DataType::Float64, true),
112        ]
113        .into_iter()
114        .map(Arc::new)
115        .collect())
116    }
117
118    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
119        if acc_args.is_distinct {
120            return not_impl_err!("STDDEV_POP(DISTINCT) aggregations are not available");
121        }
122        Ok(Box::new(StddevAccumulator::try_new(StatsType::Sample)?))
123    }
124
125    fn aliases(&self) -> &[String] {
126        &self.alias
127    }
128
129    fn groups_accumulator_supported(&self, acc_args: AccumulatorArgs) -> bool {
130        !acc_args.is_distinct
131    }
132
133    fn create_groups_accumulator(
134        &self,
135        _args: AccumulatorArgs,
136    ) -> Result<Box<dyn GroupsAccumulator>> {
137        Ok(Box::new(StddevGroupsAccumulator::new(StatsType::Sample)))
138    }
139
140    fn documentation(&self) -> Option<&Documentation> {
141        self.doc()
142    }
143}
144
145make_udaf_expr_and_func!(
146    StddevPop,
147    stddev_pop,
148    expression,
149    "Compute the population standard deviation of a set of numbers",
150    stddev_pop_udaf
151);
152
153#[user_doc(
154    doc_section(label = "Statistical Functions"),
155    description = "Returns the population standard deviation of a set of numbers.",
156    syntax_example = "stddev_pop(expression)",
157    sql_example = r#"```sql
158> SELECT stddev_pop(column_name) FROM table_name;
159+--------------------------+
160| stddev_pop(column_name)   |
161+--------------------------+
162| 10.56                    |
163+--------------------------+
164```"#,
165    standard_argument(name = "expression",)
166)]
167/// STDDEV_POP population aggregate expression
168#[derive(PartialEq, Eq, Hash, Debug)]
169pub struct StddevPop {
170    signature: Signature,
171}
172
173impl Default for StddevPop {
174    fn default() -> Self {
175        Self::new()
176    }
177}
178
179impl StddevPop {
180    /// Create a new STDDEV_POP aggregate function
181    pub fn new() -> Self {
182        Self {
183            signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable),
184        }
185    }
186}
187
188impl AggregateUDFImpl for StddevPop {
189    fn name(&self) -> &str {
190        "stddev_pop"
191    }
192
193    fn signature(&self) -> &Signature {
194        &self.signature
195    }
196
197    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
198        Ok(vec![
199            Field::new(
200                format_state_name(args.name, "count"),
201                DataType::UInt64,
202                true,
203            ),
204            Field::new(
205                format_state_name(args.name, "mean"),
206                DataType::Float64,
207                true,
208            ),
209            Field::new(format_state_name(args.name, "m2"), DataType::Float64, true),
210        ]
211        .into_iter()
212        .map(Arc::new)
213        .collect())
214    }
215
216    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
217        if acc_args.is_distinct {
218            return not_impl_err!("STDDEV_POP(DISTINCT) aggregations are not available");
219        }
220        Ok(Box::new(StddevAccumulator::try_new(StatsType::Population)?))
221    }
222
223    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
224        Ok(DataType::Float64)
225    }
226
227    fn groups_accumulator_supported(&self, acc_args: AccumulatorArgs) -> bool {
228        !acc_args.is_distinct
229    }
230
231    fn create_groups_accumulator(
232        &self,
233        _args: AccumulatorArgs,
234    ) -> Result<Box<dyn GroupsAccumulator>> {
235        Ok(Box::new(StddevGroupsAccumulator::new(
236            StatsType::Population,
237        )))
238    }
239
240    fn documentation(&self) -> Option<&Documentation> {
241        self.doc()
242    }
243}
244
245/// An accumulator to compute the average
246#[derive(Debug)]
247pub struct StddevAccumulator {
248    variance: VarianceAccumulator,
249}
250
251impl StddevAccumulator {
252    /// Creates a new `StddevAccumulator`
253    pub fn try_new(s_type: StatsType) -> Result<Self> {
254        Ok(Self {
255            variance: VarianceAccumulator::try_new(s_type)?,
256        })
257    }
258
259    pub fn get_m2(&self) -> f64 {
260        self.variance.get_m2()
261    }
262}
263
264impl Accumulator for StddevAccumulator {
265    fn state(&mut self) -> Result<Vec<ScalarValue>> {
266        Ok(vec![
267            ScalarValue::from(self.variance.get_count()),
268            ScalarValue::from(self.variance.get_mean()),
269            ScalarValue::from(self.variance.get_m2()),
270        ])
271    }
272
273    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
274        self.variance.update_batch(values)
275    }
276
277    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
278        self.variance.retract_batch(values)
279    }
280
281    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
282        self.variance.merge_batch(states)
283    }
284
285    fn evaluate(&mut self) -> Result<ScalarValue> {
286        let variance = self.variance.evaluate()?;
287        match variance {
288            ScalarValue::Float64(None) => Ok(ScalarValue::Float64(None)),
289            ScalarValue::Float64(Some(f)) => Ok(ScalarValue::Float64(Some(f.sqrt()))),
290            _ => internal_err!("Variance should be f64"),
291        }
292    }
293
294    fn size(&self) -> usize {
295        align_of_val(self) - align_of_val(&self.variance) + self.variance.size()
296    }
297
298    fn supports_retract_batch(&self) -> bool {
299        self.variance.supports_retract_batch()
300    }
301}
302
303#[derive(Debug)]
304pub struct StddevGroupsAccumulator {
305    variance: VarianceGroupsAccumulator,
306}
307
308impl StddevGroupsAccumulator {
309    pub fn new(s_type: StatsType) -> Self {
310        Self {
311            variance: VarianceGroupsAccumulator::new(s_type),
312        }
313    }
314}
315
316impl GroupsAccumulator for StddevGroupsAccumulator {
317    fn update_batch(
318        &mut self,
319        values: &[ArrayRef],
320        group_indices: &[usize],
321        opt_filter: Option<&BooleanArray>,
322        total_num_groups: usize,
323    ) -> Result<()> {
324        self.variance
325            .update_batch(values, group_indices, opt_filter, total_num_groups)
326    }
327
328    fn merge_batch(
329        &mut self,
330        values: &[ArrayRef],
331        group_indices: &[usize],
332        total_num_groups: usize,
333    ) -> Result<()> {
334        self.variance
335            .merge_batch(values, group_indices, total_num_groups)
336    }
337
338    fn evaluate(&mut self, emit_to: datafusion_expr::EmitTo) -> Result<ArrayRef> {
339        let (mut variances, nulls) = self.variance.variance(emit_to);
340        variances.iter_mut().for_each(|v| *v = v.sqrt());
341        Ok(Arc::new(Float64Array::new(variances.into(), Some(nulls))))
342    }
343
344    fn state(&mut self, emit_to: datafusion_expr::EmitTo) -> Result<Vec<ArrayRef>> {
345        self.variance.state(emit_to)
346    }
347
348    fn convert_to_state(
349        &self,
350        values: &[ArrayRef],
351        opt_filter: Option<&BooleanArray>,
352    ) -> Result<Vec<ArrayRef>> {
353        self.variance.convert_to_state(values, opt_filter)
354    }
355    fn size(&self) -> usize {
356        self.variance.size()
357    }
358}
359
360#[cfg(test)]
361mod tests {
362    use super::*;
363    use arrow::{array::*, datatypes::*};
364    use datafusion_expr::AggregateUDF;
365    use datafusion_functions_aggregate_common::utils::get_accum_scalar_values_as_arrays;
366    use datafusion_physical_expr::expressions::col;
367
368    #[test]
369    fn stddev_f64_merge_1() -> Result<()> {
370        let a = Arc::new(Float64Array::from(vec![1_f64, 2_f64, 3_f64]));
371        let b = Arc::new(Float64Array::from(vec![4_f64, 5_f64]));
372
373        let schema = Schema::new(vec![Field::new("a", DataType::Float64, false)]);
374
375        let batch1 = RecordBatch::try_new(Arc::new(schema.clone()), vec![a])?;
376        let batch2 = RecordBatch::try_new(Arc::new(schema.clone()), vec![b])?;
377
378        let agg1 = stddev_pop_udaf();
379        let agg2 = stddev_pop_udaf();
380
381        let actual = merge(&batch1, &batch2, agg1, agg2, &schema)?;
382        assert_eq!(actual, ScalarValue::from(std::f64::consts::SQRT_2));
383
384        Ok(())
385    }
386
387    #[test]
388    fn stddev_f64_merge_2() -> Result<()> {
389        let a = Arc::new(Float64Array::from(vec![1_f64, 2_f64, 3_f64, 4_f64, 5_f64]));
390        let b = Arc::new(Float64Array::from(vec![None]));
391
392        let schema = Schema::new(vec![Field::new("a", DataType::Float64, true)]);
393
394        let batch1 = RecordBatch::try_new(Arc::new(schema.clone()), vec![a])?;
395        let batch2 = RecordBatch::try_new(Arc::new(schema.clone()), vec![b])?;
396
397        let agg1 = stddev_pop_udaf();
398        let agg2 = stddev_pop_udaf();
399
400        let actual = merge(&batch1, &batch2, agg1, agg2, &schema)?;
401        assert_eq!(actual, ScalarValue::from(std::f64::consts::SQRT_2));
402
403        Ok(())
404    }
405
406    fn merge(
407        batch1: &RecordBatch,
408        batch2: &RecordBatch,
409        agg1: Arc<AggregateUDF>,
410        agg2: Arc<AggregateUDF>,
411        schema: &Schema,
412    ) -> Result<ScalarValue> {
413        let expr = col("a", schema)?;
414        let expr_field = expr.return_field(schema)?;
415
416        let args1 = AccumulatorArgs {
417            return_field: Field::new("f", DataType::Float64, true).into(),
418            schema,
419            expr_fields: &[Arc::clone(&expr_field)],
420            ignore_nulls: false,
421            order_bys: &[],
422            name: "a",
423            is_distinct: false,
424            is_reversed: false,
425            exprs: &[Arc::clone(&expr)],
426        };
427
428        let args2 = AccumulatorArgs {
429            return_field: Field::new("f", DataType::Float64, true).into(),
430            schema,
431            expr_fields: &[expr_field],
432            ignore_nulls: false,
433            order_bys: &[],
434            name: "a",
435            is_distinct: false,
436            is_reversed: false,
437            exprs: &[expr],
438        };
439
440        let mut accum1 = agg1.accumulator(args1)?;
441        let mut accum2 = agg2.accumulator(args2)?;
442
443        let value1 = vec![
444            col("a", schema)?
445                .evaluate(batch1)
446                .and_then(|v| v.into_array(batch1.num_rows()))?,
447        ];
448        let value2 = vec![
449            col("a", schema)?
450                .evaluate(batch2)
451                .and_then(|v| v.into_array(batch2.num_rows()))?,
452        ];
453
454        accum1.update_batch(&value1)?;
455        accum2.update_batch(&value2)?;
456        let state2 = get_accum_scalar_values_as_arrays(accum2.as_mut())?;
457        accum1.merge_batch(&state2)?;
458        let result = accum1.evaluate()?;
459        Ok(result)
460    }
461}
datafusion_functions_aggregate/stddev.rs

datafusion_functions_aggregate/
stddev.rs
