datafusion_functions_aggregate/
approx_median.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 for APPROX_MEDIAN that can be evaluated MEDIAN at runtime during query execution
19
20use arrow::datatypes::DataType::{Float64, UInt64};
21use arrow::datatypes::{DataType, Field, FieldRef};
22use datafusion_common::types::NativeType;
23use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator;
24use std::fmt::Debug;
25use std::sync::Arc;
26
27use datafusion_common::{Result, not_impl_err};
28use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
29use datafusion_expr::utils::format_state_name;
30use datafusion_expr::{
31    Accumulator, AggregateUDFImpl, Coercion, Documentation, Signature, TypeSignature,
32    TypeSignatureClass, Volatility,
33};
34use datafusion_macros::user_doc;
35
36use crate::approx_percentile_cont::ApproxPercentileAccumulator;
37
38make_udaf_expr_and_func!(
39    ApproxMedian,
40    approx_median,
41    expression,
42    "Computes the approximate median of a set of numbers",
43    approx_median_udaf
44);
45
46/// APPROX_MEDIAN aggregate expression
47#[user_doc(
48    doc_section(label = "Approximate Functions"),
49    description = "Returns the approximate median (50th percentile) of input values. It is an alias of `approx_percentile_cont(0.5) WITHIN GROUP (ORDER BY x)`.",
50    syntax_example = "approx_median(expression)",
51    sql_example = r#"```sql
52> SELECT approx_median(column_name) FROM table_name;
53+-----------------------------------+
54| approx_median(column_name)        |
55+-----------------------------------+
56| 23.5                              |
57+-----------------------------------+
58```"#,
59    standard_argument(name = "expression",)
60)]
61#[derive(Debug, PartialEq, Eq, Hash)]
62pub struct ApproxMedian {
63    signature: Signature,
64}
65
66impl Default for ApproxMedian {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl ApproxMedian {
73    /// Create a new APPROX_MEDIAN aggregate function
74    pub fn new() -> Self {
75        Self {
76            signature: Signature::one_of(
77                vec![TypeSignature::Coercible(vec![Coercion::new_implicit(
78                    TypeSignatureClass::Float,
79                    vec![TypeSignatureClass::Numeric],
80                    NativeType::Float64,
81                )])],
82                Volatility::Immutable,
83            ),
84        }
85    }
86}
87
88impl AggregateUDFImpl for ApproxMedian {
89    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
90        if args.input_fields[0].data_type().is_null() {
91            Ok(vec![
92                Field::new(
93                    format_state_name(args.name, self.name()),
94                    DataType::Null,
95                    true,
96                )
97                .into(),
98            ])
99        } else {
100            Ok(vec![
101                Field::new(format_state_name(args.name, "max_size"), UInt64, false),
102                Field::new(format_state_name(args.name, "sum"), Float64, false),
103                Field::new(format_state_name(args.name, "count"), Float64, false),
104                Field::new(format_state_name(args.name, "max"), Float64, false),
105                Field::new(format_state_name(args.name, "min"), Float64, false),
106                Field::new_list(
107                    format_state_name(args.name, "centroids"),
108                    Field::new_list_field(Float64, true),
109                    false,
110                ),
111            ]
112            .into_iter()
113            .map(Arc::new)
114            .collect())
115        }
116    }
117
118    fn name(&self) -> &str {
119        "approx_median"
120    }
121
122    fn signature(&self) -> &Signature {
123        &self.signature
124    }
125
126    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
127        Ok(arg_types[0].clone())
128    }
129
130    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
131        if acc_args.is_distinct {
132            return not_impl_err!(
133                "APPROX_MEDIAN(DISTINCT) aggregations are not available"
134            );
135        }
136
137        if acc_args.expr_fields[0].data_type().is_null() {
138            Ok(Box::new(NoopAccumulator::default()))
139        } else {
140            Ok(Box::new(ApproxPercentileAccumulator::new(
141                0.5_f64,
142                acc_args.expr_fields[0].data_type().clone(),
143            )))
144        }
145    }
146
147    fn documentation(&self) -> Option<&Documentation> {
148        self.doc()
149    }
150}
datafusion_functions_aggregate/approx_median.rs

datafusion_functions_aggregate/
approx_median.rs
