datafusion_extra_functions/
mode.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, common as df_common, error, logical_expr};
19use std::{any, fmt};
20
21use crate::common;
22
23make_udaf_expr_and_func!(
24    ModeFunction,
25    mode,
26    x,
27    "Calculates the most frequent value.",
28    mode_udaf
29);
30
31/// The `ModeFunction` calculates the mode (most frequent value) from a set of values.
32///
33/// - Null values are ignored during the calculation.
34/// - If multiple values have the same frequency, the MAX value with the highest frequency is returned.
35#[derive(Eq, Hash, PartialEq)]
36pub struct ModeFunction {
37    signature: logical_expr::Signature,
38}
39
40impl fmt::Debug for ModeFunction {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        f.debug_struct("ModeFunction")
43            .field("signature", &self.signature)
44            .finish()
45    }
46}
47
48impl Default for ModeFunction {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl ModeFunction {
55    pub fn new() -> Self {
56        Self {
57            signature: logical_expr::Signature::variadic_any(logical_expr::Volatility::Immutable),
58        }
59    }
60}
61
62impl logical_expr::AggregateUDFImpl for ModeFunction {
63    fn as_any(&self) -> &dyn any::Any {
64        self
65    }
66
67    fn name(&self) -> &str {
68        "mode"
69    }
70
71    fn signature(&self) -> &logical_expr::Signature {
72        &self.signature
73    }
74
75    fn return_type(
76        &self,
77        arg_types: &[arrow::datatypes::DataType],
78    ) -> error::Result<arrow::datatypes::DataType> {
79        Ok(arg_types[0].clone())
80    }
81
82    fn state_fields(
83        &self,
84        args: logical_expr::function::StateFieldsArgs,
85    ) -> error::Result<Vec<arrow::datatypes::FieldRef>> {
86        let value_type = args.input_fields[0].data_type().clone();
87
88        Ok(vec![
89            arrow::datatypes::Field::new("values", value_type, true).into(),
90            arrow::datatypes::Field::new("frequencies", arrow::datatypes::DataType::UInt64, true)
91                .into(),
92        ])
93    }
94
95    fn accumulator(
96        &self,
97        acc_args: logical_expr::function::AccumulatorArgs,
98    ) -> error::Result<Box<dyn logical_expr::Accumulator>> {
99        let data_type = &acc_args.exprs[0].data_type(acc_args.schema)?;
100
101        let accumulator: Box<dyn logical_expr::Accumulator> = match data_type {
102            arrow::datatypes::DataType::Int8 => Box::new(common::mode::PrimitiveModeAccumulator::<
103                arrow::datatypes::Int8Type,
104            >::new(data_type)),
105            arrow::datatypes::DataType::Int16 => {
106                Box::new(common::mode::PrimitiveModeAccumulator::<
107                    arrow::datatypes::Int16Type,
108                >::new(data_type))
109            }
110            arrow::datatypes::DataType::Int32 => {
111                Box::new(common::mode::PrimitiveModeAccumulator::<
112                    arrow::datatypes::Int32Type,
113                >::new(data_type))
114            }
115            arrow::datatypes::DataType::Int64 => {
116                Box::new(common::mode::PrimitiveModeAccumulator::<
117                    arrow::datatypes::Int64Type,
118                >::new(data_type))
119            }
120            arrow::datatypes::DataType::UInt8 => {
121                Box::new(common::mode::PrimitiveModeAccumulator::<
122                    arrow::datatypes::UInt8Type,
123                >::new(data_type))
124            }
125            arrow::datatypes::DataType::UInt16 => {
126                Box::new(common::mode::PrimitiveModeAccumulator::<
127                    arrow::datatypes::UInt16Type,
128                >::new(data_type))
129            }
130            arrow::datatypes::DataType::UInt32 => {
131                Box::new(common::mode::PrimitiveModeAccumulator::<
132                    arrow::datatypes::UInt32Type,
133                >::new(data_type))
134            }
135            arrow::datatypes::DataType::UInt64 => {
136                Box::new(common::mode::PrimitiveModeAccumulator::<
137                    arrow::datatypes::UInt64Type,
138                >::new(data_type))
139            }
140
141            arrow::datatypes::DataType::Date32 => {
142                Box::new(common::mode::PrimitiveModeAccumulator::<
143                    arrow::datatypes::Date32Type,
144                >::new(data_type))
145            }
146            arrow::datatypes::DataType::Date64 => {
147                Box::new(common::mode::PrimitiveModeAccumulator::<
148                    arrow::datatypes::Date64Type,
149                >::new(data_type))
150            }
151            arrow::datatypes::DataType::Time32(arrow::datatypes::TimeUnit::Millisecond) => {
152                Box::new(common::mode::PrimitiveModeAccumulator::<
153                    arrow::datatypes::Time32MillisecondType,
154                >::new(data_type))
155            }
156            arrow::datatypes::DataType::Time32(arrow::datatypes::TimeUnit::Second) => {
157                Box::new(common::mode::PrimitiveModeAccumulator::<
158                    arrow::datatypes::Time32SecondType,
159                >::new(data_type))
160            }
161            arrow::datatypes::DataType::Time64(arrow::datatypes::TimeUnit::Microsecond) => {
162                Box::new(common::mode::PrimitiveModeAccumulator::<
163                    arrow::datatypes::Time64MicrosecondType,
164                >::new(data_type))
165            }
166            arrow::datatypes::DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond) => {
167                Box::new(common::mode::PrimitiveModeAccumulator::<
168                    arrow::datatypes::Time64NanosecondType,
169                >::new(data_type))
170            }
171            arrow::datatypes::DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, _) => {
172                Box::new(common::mode::PrimitiveModeAccumulator::<
173                    arrow::datatypes::TimestampMicrosecondType,
174                >::new(data_type))
175            }
176            arrow::datatypes::DataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, _) => {
177                Box::new(common::mode::PrimitiveModeAccumulator::<
178                    arrow::datatypes::TimestampMillisecondType,
179                >::new(data_type))
180            }
181            arrow::datatypes::DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, _) => {
182                Box::new(common::mode::PrimitiveModeAccumulator::<
183                    arrow::datatypes::TimestampNanosecondType,
184                >::new(data_type))
185            }
186            arrow::datatypes::DataType::Timestamp(arrow::datatypes::TimeUnit::Second, _) => {
187                Box::new(common::mode::PrimitiveModeAccumulator::<
188                    arrow::datatypes::TimestampSecondType,
189                >::new(data_type))
190            }
191
192            arrow::datatypes::DataType::Float16 => Box::new(common::mode::FloatModeAccumulator::<
193                arrow::datatypes::Float16Type,
194            >::new(data_type)),
195            arrow::datatypes::DataType::Float32 => Box::new(common::mode::FloatModeAccumulator::<
196                arrow::datatypes::Float32Type,
197            >::new(data_type)),
198            arrow::datatypes::DataType::Float64 => Box::new(common::mode::FloatModeAccumulator::<
199                arrow::datatypes::Float64Type,
200            >::new(data_type)),
201
202            arrow::datatypes::DataType::Utf8
203            | arrow::datatypes::DataType::Utf8View
204            | arrow::datatypes::DataType::LargeUtf8 => {
205                Box::new(common::mode::BytesModeAccumulator::new(data_type))
206            }
207            _ => {
208                return df_common::not_impl_err!(
209                    "Unsupported data type: {:?} for mode function",
210                    data_type
211                );
212            }
213        };
214
215        Ok(accumulator)
216    }
217}
datafusion_extra_functions/mode.rs

datafusion_extra_functions/
mode.rs
