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

datafusion_extra_functions/
mode.rs
