datafusion_expr/udaf.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//! [`AggregateUDF`]: User Defined Aggregate Functions
19
20use std::any::Any;
21use std::cmp::Ordering;
22use std::fmt::{self, Debug, Formatter, Write};
23use std::hash::{Hash, Hasher};
24use std::sync::Arc;
25use std::vec;
26
27use arrow::datatypes::{DataType, Field, FieldRef};
28
29use datafusion_common::{exec_err, not_impl_err, Result, ScalarValue, Statistics};
30use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
31use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
32
33use crate::expr::{
34 schema_name_from_exprs, schema_name_from_exprs_comma_separated_without_space,
35 schema_name_from_sorts, AggregateFunction, AggregateFunctionParams, ExprListDisplay,
36 WindowFunctionParams,
37};
38use crate::function::{
39 AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs,
40};
41use crate::groups_accumulator::GroupsAccumulator;
42use crate::udf_eq::UdfEq;
43use crate::utils::format_state_name;
44use crate::utils::AggregateOrderSensitivity;
45use crate::{expr_vec_fmt, Accumulator, Expr};
46use crate::{Documentation, Signature};
47
48/// Logical representation of a user-defined [aggregate function] (UDAF).
49///
50/// An aggregate function combines the values from multiple input rows
51/// into a single output "aggregate" (summary) row. It is different
52/// from a scalar function because it is stateful across batches. User
53/// defined aggregate functions can be used as normal SQL aggregate
54/// functions (`GROUP BY` clause) as well as window functions (`OVER`
55/// clause).
56///
57/// `AggregateUDF` provides DataFusion the information needed to plan and call
58/// aggregate functions, including name, type information, and a factory
59/// function to create an [`Accumulator`] instance, to perform the actual
60/// aggregation.
61///
62/// For more information, please see [the examples]:
63///
64/// 1. For simple use cases, use [`create_udaf`] (examples in [`simple_udaf.rs`]).
65///
66/// 2. For advanced use cases, use [`AggregateUDFImpl`] which provides full API
67/// access (examples in [`advanced_udaf.rs`]).
68///
69/// # API Note
70/// This is a separate struct from `AggregateUDFImpl` to maintain backwards
71/// compatibility with the older API.
72///
73/// [the examples]: https://github.com/apache/datafusion/tree/main/datafusion-examples#single-process
74/// [aggregate function]: https://en.wikipedia.org/wiki/Aggregate_function
75/// [`Accumulator`]: Accumulator
76/// [`create_udaf`]: crate::expr_fn::create_udaf
77/// [`simple_udaf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/simple_udaf.rs
78/// [`advanced_udaf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udaf.rs
79#[derive(Debug, Clone, PartialOrd)]
80pub struct AggregateUDF {
81 inner: Arc<dyn AggregateUDFImpl>,
82}
83
84impl PartialEq for AggregateUDF {
85 fn eq(&self, other: &Self) -> bool {
86 self.inner.dyn_eq(other.inner.as_any())
87 }
88}
89
90impl Eq for AggregateUDF {}
91
92impl Hash for AggregateUDF {
93 fn hash<H: Hasher>(&self, state: &mut H) {
94 self.inner.dyn_hash(state)
95 }
96}
97
98impl fmt::Display for AggregateUDF {
99 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
100 write!(f, "{}", self.name())
101 }
102}
103
104/// Arguments passed to [`AggregateUDFImpl::value_from_stats`]
105#[derive(Debug)]
106pub struct StatisticsArgs<'a> {
107 /// The statistics of the aggregate input
108 pub statistics: &'a Statistics,
109 /// The resolved return type of the aggregate function
110 pub return_type: &'a DataType,
111 /// Whether the aggregate function is distinct.
112 ///
113 /// ```sql
114 /// SELECT COUNT(DISTINCT column1) FROM t;
115 /// ```
116 pub is_distinct: bool,
117 /// The physical expression of arguments the aggregate function takes.
118 pub exprs: &'a [Arc<dyn PhysicalExpr>],
119}
120
121impl AggregateUDF {
122 /// Create a new `AggregateUDF` from a `[AggregateUDFImpl]` trait object
123 ///
124 /// Note this is the same as using the `From` impl (`AggregateUDF::from`)
125 pub fn new_from_impl<F>(fun: F) -> AggregateUDF
126 where
127 F: AggregateUDFImpl + 'static,
128 {
129 Self::new_from_shared_impl(Arc::new(fun))
130 }
131
132 /// Create a new `AggregateUDF` from a `[AggregateUDFImpl]` trait object
133 pub fn new_from_shared_impl(fun: Arc<dyn AggregateUDFImpl>) -> AggregateUDF {
134 Self { inner: fun }
135 }
136
137 /// Return the underlying [`AggregateUDFImpl`] trait object for this function
138 pub fn inner(&self) -> &Arc<dyn AggregateUDFImpl> {
139 &self.inner
140 }
141
142 /// Adds additional names that can be used to invoke this function, in
143 /// addition to `name`
144 ///
145 /// If you implement [`AggregateUDFImpl`] directly you should return aliases directly.
146 pub fn with_aliases(self, aliases: impl IntoIterator<Item = &'static str>) -> Self {
147 Self::new_from_impl(AliasedAggregateUDFImpl::new(
148 Arc::clone(&self.inner),
149 aliases,
150 ))
151 }
152
153 /// Creates an [`Expr`] that calls the aggregate function.
154 ///
155 /// This utility allows using the UDAF without requiring access to
156 /// the registry, such as with the DataFrame API.
157 pub fn call(&self, args: Vec<Expr>) -> Expr {
158 Expr::AggregateFunction(AggregateFunction::new_udf(
159 Arc::new(self.clone()),
160 args,
161 false,
162 None,
163 vec![],
164 None,
165 ))
166 }
167
168 /// Returns this function's name
169 ///
170 /// See [`AggregateUDFImpl::name`] for more details.
171 pub fn name(&self) -> &str {
172 self.inner.name()
173 }
174
175 /// Returns the aliases for this function.
176 pub fn aliases(&self) -> &[String] {
177 self.inner.aliases()
178 }
179
180 /// See [`AggregateUDFImpl::schema_name`] for more details.
181 pub fn schema_name(&self, params: &AggregateFunctionParams) -> Result<String> {
182 self.inner.schema_name(params)
183 }
184
185 /// Returns a human readable expression.
186 ///
187 /// See [`Expr::human_display`] for details.
188 pub fn human_display(&self, params: &AggregateFunctionParams) -> Result<String> {
189 self.inner.human_display(params)
190 }
191
192 pub fn window_function_schema_name(
193 &self,
194 params: &WindowFunctionParams,
195 ) -> Result<String> {
196 self.inner.window_function_schema_name(params)
197 }
198
199 /// See [`AggregateUDFImpl::display_name`] for more details.
200 pub fn display_name(&self, params: &AggregateFunctionParams) -> Result<String> {
201 self.inner.display_name(params)
202 }
203
204 pub fn window_function_display_name(
205 &self,
206 params: &WindowFunctionParams,
207 ) -> Result<String> {
208 self.inner.window_function_display_name(params)
209 }
210
211 pub fn is_nullable(&self) -> bool {
212 self.inner.is_nullable()
213 }
214
215 /// Returns this function's signature (what input types are accepted)
216 ///
217 /// See [`AggregateUDFImpl::signature`] for more details.
218 pub fn signature(&self) -> &Signature {
219 self.inner.signature()
220 }
221
222 /// Return the type of the function given its input types
223 ///
224 /// See [`AggregateUDFImpl::return_type`] for more details.
225 pub fn return_type(&self, args: &[DataType]) -> Result<DataType> {
226 self.inner.return_type(args)
227 }
228
229 /// Return the field of the function given its input fields
230 ///
231 /// See [`AggregateUDFImpl::return_field`] for more details.
232 pub fn return_field(&self, args: &[FieldRef]) -> Result<FieldRef> {
233 self.inner.return_field(args)
234 }
235
236 /// Return an accumulator the given aggregate, given its return datatype
237 pub fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
238 self.inner.accumulator(acc_args)
239 }
240
241 /// Return the fields used to store the intermediate state for this aggregator, given
242 /// the name of the aggregate, value type and ordering fields. See [`AggregateUDFImpl::state_fields`]
243 /// for more details.
244 ///
245 /// This is used to support multi-phase aggregations
246 pub fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
247 self.inner.state_fields(args)
248 }
249
250 /// See [`AggregateUDFImpl::groups_accumulator_supported`] for more details.
251 pub fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
252 self.inner.groups_accumulator_supported(args)
253 }
254
255 /// See [`AggregateUDFImpl::create_groups_accumulator`] for more details.
256 pub fn create_groups_accumulator(
257 &self,
258 args: AccumulatorArgs,
259 ) -> Result<Box<dyn GroupsAccumulator>> {
260 self.inner.create_groups_accumulator(args)
261 }
262
263 pub fn create_sliding_accumulator(
264 &self,
265 args: AccumulatorArgs,
266 ) -> Result<Box<dyn Accumulator>> {
267 self.inner.create_sliding_accumulator(args)
268 }
269
270 pub fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
271 self.inner.coerce_types(arg_types)
272 }
273
274 /// See [`AggregateUDFImpl::with_beneficial_ordering`] for more details.
275 pub fn with_beneficial_ordering(
276 self,
277 beneficial_ordering: bool,
278 ) -> Result<Option<AggregateUDF>> {
279 self.inner
280 .with_beneficial_ordering(beneficial_ordering)
281 .map(|updated_udf| updated_udf.map(|udf| Self { inner: udf }))
282 }
283
284 /// Gets the order sensitivity of the UDF. See [`AggregateOrderSensitivity`]
285 /// for possible options.
286 pub fn order_sensitivity(&self) -> AggregateOrderSensitivity {
287 self.inner.order_sensitivity()
288 }
289
290 /// Reserves the `AggregateUDF` (e.g. returns the `AggregateUDF` that will
291 /// generate same result with this `AggregateUDF` when iterated in reverse
292 /// order, and `None` if there is no such `AggregateUDF`).
293 pub fn reverse_udf(&self) -> ReversedUDAF {
294 self.inner.reverse_expr()
295 }
296
297 /// Do the function rewrite
298 ///
299 /// See [`AggregateUDFImpl::simplify`] for more details.
300 pub fn simplify(&self) -> Option<AggregateFunctionSimplification> {
301 self.inner.simplify()
302 }
303
304 /// Returns true if the function is max, false if the function is min
305 /// None in all other cases, used in certain optimizations for
306 /// or aggregate
307 pub fn is_descending(&self) -> Option<bool> {
308 self.inner.is_descending()
309 }
310
311 /// Return the value of this aggregate function if it can be determined
312 /// entirely from statistics and arguments.
313 ///
314 /// See [`AggregateUDFImpl::value_from_stats`] for more details.
315 pub fn value_from_stats(
316 &self,
317 statistics_args: &StatisticsArgs,
318 ) -> Option<ScalarValue> {
319 self.inner.value_from_stats(statistics_args)
320 }
321
322 /// See [`AggregateUDFImpl::default_value`] for more details.
323 pub fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
324 self.inner.default_value(data_type)
325 }
326
327 /// See [`AggregateUDFImpl::supports_null_handling_clause`] for more details.
328 pub fn supports_null_handling_clause(&self) -> bool {
329 self.inner.supports_null_handling_clause()
330 }
331
332 /// See [`AggregateUDFImpl::supports_within_group_clause`] for more details.
333 pub fn supports_within_group_clause(&self) -> bool {
334 self.inner.supports_within_group_clause()
335 }
336
337 /// Returns the documentation for this Aggregate UDF.
338 ///
339 /// Documentation can be accessed programmatically as well as
340 /// generating publicly facing documentation.
341 pub fn documentation(&self) -> Option<&Documentation> {
342 self.inner.documentation()
343 }
344}
345
346impl<F> From<F> for AggregateUDF
347where
348 F: AggregateUDFImpl + Send + Sync + 'static,
349{
350 fn from(fun: F) -> Self {
351 Self::new_from_impl(fun)
352 }
353}
354
355/// Trait for implementing [`AggregateUDF`].
356///
357/// This trait exposes the full API for implementing user defined aggregate functions and
358/// can be used to implement any function.
359///
360/// See [`advanced_udaf.rs`] for a full example with complete implementation and
361/// [`AggregateUDF`] for other available options.
362///
363/// [`advanced_udaf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udaf.rs
364///
365/// # Basic Example
366/// ```
367/// # use std::any::Any;
368/// # use std::sync::{Arc, LazyLock};
369/// # use arrow::datatypes::{DataType, FieldRef};
370/// # use datafusion_common::{DataFusionError, plan_err, Result};
371/// # use datafusion_expr::{col, ColumnarValue, Signature, Volatility, Expr, Documentation};
372/// # use datafusion_expr::{AggregateUDFImpl, AggregateUDF, Accumulator, function::{AccumulatorArgs, StateFieldsArgs}};
373/// # use datafusion_expr::window_doc_sections::DOC_SECTION_AGGREGATE;
374/// # use arrow::datatypes::Schema;
375/// # use arrow::datatypes::Field;
376///
377/// #[derive(Debug, Clone, PartialEq, Eq, Hash)]
378/// struct GeoMeanUdf {
379/// signature: Signature,
380/// }
381///
382/// impl GeoMeanUdf {
383/// fn new() -> Self {
384/// Self {
385/// signature: Signature::uniform(1, vec![DataType::Float64], Volatility::Immutable),
386/// }
387/// }
388/// }
389///
390/// static DOCUMENTATION: LazyLock<Documentation> = LazyLock::new(|| {
391/// Documentation::builder(DOC_SECTION_AGGREGATE, "calculates a geometric mean", "geo_mean(2.0)")
392/// .with_argument("arg1", "The Float64 number for the geometric mean")
393/// .build()
394/// });
395///
396/// fn get_doc() -> &'static Documentation {
397/// &DOCUMENTATION
398/// }
399///
400/// /// Implement the AggregateUDFImpl trait for GeoMeanUdf
401/// impl AggregateUDFImpl for GeoMeanUdf {
402/// fn as_any(&self) -> &dyn Any { self }
403/// fn name(&self) -> &str { "geo_mean" }
404/// fn signature(&self) -> &Signature { &self.signature }
405/// fn return_type(&self, args: &[DataType]) -> Result<DataType> {
406/// if !matches!(args.get(0), Some(&DataType::Float64)) {
407/// return plan_err!("geo_mean only accepts Float64 arguments");
408/// }
409/// Ok(DataType::Float64)
410/// }
411/// // This is the accumulator factory; DataFusion uses it to create new accumulators.
412/// fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { unimplemented!() }
413/// fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
414/// Ok(vec![
415/// Arc::new(args.return_field.as_ref().clone().with_name("value")),
416/// Arc::new(Field::new("ordering", DataType::UInt32, true))
417/// ])
418/// }
419/// fn documentation(&self) -> Option<&Documentation> {
420/// Some(get_doc())
421/// }
422/// }
423///
424/// // Create a new AggregateUDF from the implementation
425/// let geometric_mean = AggregateUDF::from(GeoMeanUdf::new());
426///
427/// // Call the function `geo_mean(col)`
428/// let expr = geometric_mean.call(vec![col("a")]);
429/// ```
430pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
431 /// Returns this object as an [`Any`] trait object
432 fn as_any(&self) -> &dyn Any;
433
434 /// Returns this function's name
435 fn name(&self) -> &str;
436
437 /// Returns any aliases (alternate names) for this function.
438 ///
439 /// Note: `aliases` should only include names other than [`Self::name`].
440 /// Defaults to `[]` (no aliases)
441 fn aliases(&self) -> &[String] {
442 &[]
443 }
444
445 /// Returns the name of the column this expression would create
446 ///
447 /// See [`Expr::schema_name`] for details
448 ///
449 /// Example of schema_name: count(DISTINCT column1) FILTER (WHERE column2 > 10) ORDER BY [..]
450 fn schema_name(&self, params: &AggregateFunctionParams) -> Result<String> {
451 udaf_default_schema_name(self, params)
452 }
453
454 /// Returns a human readable expression.
455 ///
456 /// See [`Expr::human_display`] for details.
457 fn human_display(&self, params: &AggregateFunctionParams) -> Result<String> {
458 udaf_default_human_display(self, params)
459 }
460
461 /// Returns the name of the column this expression would create
462 ///
463 /// See [`Expr::schema_name`] for details
464 ///
465 /// Different from `schema_name` in that it is used for window aggregate function
466 ///
467 /// Example of schema_name: count(DISTINCT column1) FILTER (WHERE column2 > 10) [PARTITION BY [..]] [ORDER BY [..]]
468 fn window_function_schema_name(
469 &self,
470 params: &WindowFunctionParams,
471 ) -> Result<String> {
472 udaf_default_window_function_schema_name(self, params)
473 }
474
475 /// Returns the user-defined display name of function, given the arguments
476 ///
477 /// This can be used to customize the output column name generated by this
478 /// function.
479 ///
480 /// Defaults to `function_name([DISTINCT] column1, column2, ..) [null_treatment] [filter] [order_by [..]]`
481 fn display_name(&self, params: &AggregateFunctionParams) -> Result<String> {
482 udaf_default_display_name(self, params)
483 }
484
485 /// Returns the user-defined display name of function, given the arguments
486 ///
487 /// This can be used to customize the output column name generated by this
488 /// function.
489 ///
490 /// Different from `display_name` in that it is used for window aggregate function
491 ///
492 /// Defaults to `function_name([DISTINCT] column1, column2, ..) [null_treatment] [partition by [..]] [order_by [..]]`
493 fn window_function_display_name(
494 &self,
495 params: &WindowFunctionParams,
496 ) -> Result<String> {
497 udaf_default_window_function_display_name(self, params)
498 }
499
500 /// Returns the function's [`Signature`] for information about what input
501 /// types are accepted and the function's Volatility.
502 fn signature(&self) -> &Signature;
503
504 /// What [`DataType`] will be returned by this function, given the types of
505 /// the arguments
506 fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;
507
508 /// What type will be returned by this function, given the arguments?
509 ///
510 /// By default, this function calls [`Self::return_type`] with the
511 /// types of each argument.
512 ///
513 /// # Notes
514 ///
515 /// Most UDFs should implement [`Self::return_type`] and not this
516 /// function as the output type for most functions only depends on the types
517 /// of their inputs (e.g. `sum(f64)` is always `f64`).
518 ///
519 /// This function can be used for more advanced cases such as:
520 ///
521 /// 1. specifying nullability
522 /// 2. return types based on the **values** of the arguments (rather than
523 /// their **types**.
524 /// 3. return types based on metadata within the fields of the inputs
525 fn return_field(&self, arg_fields: &[FieldRef]) -> Result<FieldRef> {
526 udaf_default_return_field(self, arg_fields)
527 }
528
529 /// Whether the aggregate function is nullable.
530 ///
531 /// Nullable means that the function could return `null` for any inputs.
532 /// For example, aggregate functions like `COUNT` always return a non null value
533 /// but others like `MIN` will return `NULL` if there is nullable input.
534 /// Note that if the function is declared as *not* nullable, make sure the [`AggregateUDFImpl::default_value`] is `non-null`
535 fn is_nullable(&self) -> bool {
536 true
537 }
538
539 /// Return a new [`Accumulator`] that aggregates values for a specific
540 /// group during query execution.
541 ///
542 /// acc_args: [`AccumulatorArgs`] contains information about how the
543 /// aggregate function was called.
544 fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>>;
545
546 /// Return the fields used to store the intermediate state of this accumulator.
547 ///
548 /// See [`Accumulator::state`] for background information.
549 ///
550 /// args: [`StateFieldsArgs`] contains arguments passed to the
551 /// aggregate function's accumulator.
552 ///
553 /// # Notes:
554 ///
555 /// The default implementation returns a single state field named `name`
556 /// with the same type as `value_type`. This is suitable for aggregates such
557 /// as `SUM` or `MIN` where partial state can be combined by applying the
558 /// same aggregate.
559 ///
560 /// For aggregates such as `AVG` where the partial state is more complex
561 /// (e.g. a COUNT and a SUM), this method is used to define the additional
562 /// fields.
563 ///
564 /// The name of the fields must be unique within the query and thus should
565 /// be derived from `name`. See [`format_state_name`] for a utility function
566 /// to generate a unique name.
567 fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
568 let fields = vec![args
569 .return_field
570 .as_ref()
571 .clone()
572 .with_name(format_state_name(args.name, "value"))];
573
574 Ok(fields
575 .into_iter()
576 .map(Arc::new)
577 .chain(args.ordering_fields.to_vec())
578 .collect())
579 }
580
581 /// If the aggregate expression has a specialized
582 /// [`GroupsAccumulator`] implementation. If this returns true,
583 /// `[Self::create_groups_accumulator]` will be called.
584 ///
585 /// # Notes
586 ///
587 /// Even if this function returns true, DataFusion will still use
588 /// [`Self::accumulator`] for certain queries, such as when this aggregate is
589 /// used as a window function or when there no GROUP BY columns in the
590 /// query.
591 fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
592 false
593 }
594
595 /// Return a specialized [`GroupsAccumulator`] that manages state
596 /// for all groups.
597 ///
598 /// For maximum performance, a [`GroupsAccumulator`] should be
599 /// implemented in addition to [`Accumulator`].
600 fn create_groups_accumulator(
601 &self,
602 _args: AccumulatorArgs,
603 ) -> Result<Box<dyn GroupsAccumulator>> {
604 not_impl_err!("GroupsAccumulator hasn't been implemented for {self:?} yet")
605 }
606
607 /// Sliding accumulator is an alternative accumulator that can be used for
608 /// window functions. It has retract method to revert the previous update.
609 ///
610 /// See [retract_batch] for more details.
611 ///
612 /// [retract_batch]: Accumulator::retract_batch
613 fn create_sliding_accumulator(
614 &self,
615 args: AccumulatorArgs,
616 ) -> Result<Box<dyn Accumulator>> {
617 self.accumulator(args)
618 }
619
620 /// Sets the indicator whether ordering requirements of the AggregateUDFImpl is
621 /// satisfied by its input. If this is not the case, UDFs with order
622 /// sensitivity `AggregateOrderSensitivity::Beneficial` can still produce
623 /// the correct result with possibly more work internally.
624 ///
625 /// # Returns
626 ///
627 /// Returns `Ok(Some(updated_udf))` if the process completes successfully.
628 /// If the expression can benefit from existing input ordering, but does
629 /// not implement the method, returns an error. Order insensitive and hard
630 /// requirement aggregators return `Ok(None)`.
631 fn with_beneficial_ordering(
632 self: Arc<Self>,
633 _beneficial_ordering: bool,
634 ) -> Result<Option<Arc<dyn AggregateUDFImpl>>> {
635 if self.order_sensitivity().is_beneficial() {
636 return exec_err!(
637 "Should implement with satisfied for aggregator :{:?}",
638 self.name()
639 );
640 }
641 Ok(None)
642 }
643
644 /// Gets the order sensitivity of the UDF. See [`AggregateOrderSensitivity`]
645 /// for possible options.
646 fn order_sensitivity(&self) -> AggregateOrderSensitivity {
647 // We have hard ordering requirements by default, meaning that order
648 // sensitive UDFs need their input orderings to satisfy their ordering
649 // requirements to generate correct results.
650 AggregateOrderSensitivity::HardRequirement
651 }
652
653 /// Optionally apply per-UDaF simplification / rewrite rules.
654 ///
655 /// This can be used to apply function specific simplification rules during
656 /// optimization (e.g. `arrow_cast` --> `Expr::Cast`). The default
657 /// implementation does nothing.
658 ///
659 /// Note that DataFusion handles simplifying arguments and "constant
660 /// folding" (replacing a function call with constant arguments such as
661 /// `my_add(1,2) --> 3` ). Thus, there is no need to implement such
662 /// optimizations manually for specific UDFs.
663 ///
664 /// # Returns
665 ///
666 /// [None] if simplify is not defined or,
667 ///
668 /// Or, a closure with two arguments:
669 /// * 'aggregate_function': [AggregateFunction] for which simplified has been invoked
670 /// * 'info': [crate::simplify::SimplifyInfo]
671 ///
672 /// closure returns simplified [Expr] or an error.
673 ///
674 /// # Notes
675 ///
676 /// The returned expression must have the same schema as the original
677 /// expression, including both the data type and nullability. For example,
678 /// if the original expression is nullable, the returned expression must
679 /// also be nullable, otherwise it may lead to schema verification errors
680 /// later in query planning.
681 fn simplify(&self) -> Option<AggregateFunctionSimplification> {
682 None
683 }
684
685 /// Returns the reverse expression of the aggregate function.
686 fn reverse_expr(&self) -> ReversedUDAF {
687 ReversedUDAF::NotSupported
688 }
689
690 /// Coerce arguments of a function call to types that the function can evaluate.
691 ///
692 /// This function is only called if [`AggregateUDFImpl::signature`] returns [`crate::TypeSignature::UserDefined`]. Most
693 /// UDAFs should return one of the other variants of `TypeSignature` which handle common
694 /// cases
695 ///
696 /// See the [type coercion module](crate::type_coercion)
697 /// documentation for more details on type coercion
698 ///
699 /// For example, if your function requires a floating point arguments, but the user calls
700 /// it like `my_func(1::int)` (aka with `1` as an integer), coerce_types could return `[DataType::Float64]`
701 /// to ensure the argument was cast to `1::double`
702 ///
703 /// # Parameters
704 /// * `arg_types`: The argument types of the arguments this function with
705 ///
706 /// # Return value
707 /// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call
708 /// arguments to these specific types.
709 fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
710 not_impl_err!("Function {} does not implement coerce_types", self.name())
711 }
712
713 /// If this function is max, return true
714 /// If the function is min, return false
715 /// Otherwise return None (the default)
716 ///
717 ///
718 /// Note: this is used to use special aggregate implementations in certain conditions
719 fn is_descending(&self) -> Option<bool> {
720 None
721 }
722
723 /// Return the value of this aggregate function if it can be determined
724 /// entirely from statistics and arguments.
725 ///
726 /// Using a [`ScalarValue`] rather than a runtime computation can significantly
727 /// improving query performance.
728 ///
729 /// For example, if the minimum value of column `x` is known to be `42` from
730 /// statistics, then the aggregate `MIN(x)` should return `Some(ScalarValue(42))`
731 fn value_from_stats(&self, _statistics_args: &StatisticsArgs) -> Option<ScalarValue> {
732 None
733 }
734
735 /// Returns default value of the function given the input is all `null`.
736 ///
737 /// Most of the aggregate function return Null if input is Null,
738 /// while `count` returns 0 if input is Null
739 fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
740 ScalarValue::try_from(data_type)
741 }
742
743 /// If this function supports `[IGNORE NULLS | RESPECT NULLS]` clause, return true
744 /// If the function does not, return false
745 fn supports_null_handling_clause(&self) -> bool {
746 true
747 }
748
749 /// If this function supports the `WITHIN GROUP (ORDER BY column [ASC|DESC])`
750 /// SQL syntax, return `true`. Otherwise, return `false` (default) which will
751 /// cause an error when parsing SQL where this syntax is detected for this
752 /// function.
753 ///
754 /// This function should return `true` for ordered-set aggregate functions
755 /// only.
756 ///
757 /// # Ordered-set aggregate functions
758 ///
759 /// Ordered-set aggregate functions allow specifying a sort order that affects
760 /// how the function calculates its result, unlike other aggregate functions
761 /// like `sum` or `count`. For example, `percentile_cont` is an ordered-set
762 /// aggregate function that calculates the exact percentile value from a list
763 /// of values; the output of calculating the `0.75` percentile depends on if
764 /// you're calculating on an ascending or descending list of values.
765 ///
766 /// An example of how an ordered-set aggregate function is called with the
767 /// `WITHIN GROUP` SQL syntax:
768 ///
769 /// ```sql
770 /// -- Ascending
771 /// SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY c1 ASC) FROM table;
772 /// -- Default ordering is ascending if not explicitly specified
773 /// SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY c1) FROM table;
774 /// -- Descending
775 /// SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY c1 DESC) FROM table;
776 /// ```
777 ///
778 /// This calculates the `0.75` percentile of the column `c1` from `table`,
779 /// according to the specific ordering. The column specified in the `WITHIN GROUP`
780 /// ordering clause is taken as the column to calculate values on; specifying
781 /// the `WITHIN GROUP` clause is optional so these queries are equivalent:
782 ///
783 /// ```sql
784 /// -- If no WITHIN GROUP is specified then default ordering is implementation
785 /// -- dependent; in this case ascending for percentile_cont
786 /// SELECT percentile_cont(c1, 0.75) FROM table;
787 /// SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY c1 ASC) FROM table;
788 /// ```
789 ///
790 /// Aggregate UDFs can define their default ordering if the function is called
791 /// without the `WITHIN GROUP` clause, though a default of ascending is the
792 /// standard practice.
793 ///
794 /// Ordered-set aggregate function implementations are responsible for handling
795 /// the input sort order themselves (e.g. `percentile_cont` must buffer and
796 /// sort the values internally). That is, DataFusion does not introduce any
797 /// kind of sort into the plan for these functions with this syntax.
798 fn supports_within_group_clause(&self) -> bool {
799 false
800 }
801
802 /// Returns the documentation for this Aggregate UDF.
803 ///
804 /// Documentation can be accessed programmatically as well as
805 /// generating publicly facing documentation.
806 fn documentation(&self) -> Option<&Documentation> {
807 None
808 }
809
810 /// Indicates whether the aggregation function is monotonic as a set
811 /// function. See [`SetMonotonicity`] for details.
812 fn set_monotonicity(&self, _data_type: &DataType) -> SetMonotonicity {
813 SetMonotonicity::NotMonotonic
814 }
815}
816
817impl PartialEq for dyn AggregateUDFImpl {
818 fn eq(&self, other: &Self) -> bool {
819 self.dyn_eq(other.as_any())
820 }
821}
822
823impl PartialOrd for dyn AggregateUDFImpl {
824 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
825 match self.name().partial_cmp(other.name()) {
826 Some(Ordering::Equal) => self.signature().partial_cmp(other.signature()),
827 cmp => cmp,
828 }
829 // TODO (https://github.com/apache/datafusion/issues/17477) avoid recomparing all fields
830 .filter(|cmp| *cmp != Ordering::Equal || self == other)
831 }
832}
833
834/// Encapsulates default implementation of [`AggregateUDFImpl::schema_name`].
835pub fn udaf_default_schema_name<F: AggregateUDFImpl + ?Sized>(
836 func: &F,
837 params: &AggregateFunctionParams,
838) -> Result<String> {
839 let AggregateFunctionParams {
840 args,
841 distinct,
842 filter,
843 order_by,
844 null_treatment,
845 } = params;
846
847 // exclude the first function argument(= column) in ordered set aggregate function,
848 // because it is duplicated with the WITHIN GROUP clause in schema name.
849 let args = if func.supports_within_group_clause() && !order_by.is_empty() {
850 &args[1..]
851 } else {
852 &args[..]
853 };
854
855 let mut schema_name = String::new();
856
857 schema_name.write_fmt(format_args!(
858 "{}({}{})",
859 func.name(),
860 if *distinct { "DISTINCT " } else { "" },
861 schema_name_from_exprs_comma_separated_without_space(args)?
862 ))?;
863
864 if let Some(null_treatment) = null_treatment {
865 schema_name.write_fmt(format_args!(" {null_treatment}"))?;
866 }
867
868 if let Some(filter) = filter {
869 schema_name.write_fmt(format_args!(" FILTER (WHERE {filter})"))?;
870 };
871
872 if !order_by.is_empty() {
873 let clause = match func.supports_within_group_clause() {
874 true => "WITHIN GROUP",
875 false => "ORDER BY",
876 };
877
878 schema_name.write_fmt(format_args!(
879 " {} [{}]",
880 clause,
881 schema_name_from_sorts(order_by)?
882 ))?;
883 };
884
885 Ok(schema_name)
886}
887
888/// Encapsulates default implementation of [`AggregateUDFImpl::human_display`].
889pub fn udaf_default_human_display<F: AggregateUDFImpl + ?Sized>(
890 func: &F,
891 params: &AggregateFunctionParams,
892) -> Result<String> {
893 let AggregateFunctionParams {
894 args,
895 distinct,
896 filter,
897 order_by,
898 null_treatment,
899 } = params;
900
901 let mut schema_name = String::new();
902
903 schema_name.write_fmt(format_args!(
904 "{}({}{})",
905 func.name(),
906 if *distinct { "DISTINCT " } else { "" },
907 ExprListDisplay::comma_separated(args.as_slice())
908 ))?;
909
910 if let Some(null_treatment) = null_treatment {
911 schema_name.write_fmt(format_args!(" {null_treatment}"))?;
912 }
913
914 if let Some(filter) = filter {
915 schema_name.write_fmt(format_args!(" FILTER (WHERE {filter})"))?;
916 };
917
918 if !order_by.is_empty() {
919 schema_name.write_fmt(format_args!(
920 " ORDER BY [{}]",
921 schema_name_from_sorts(order_by)?
922 ))?;
923 };
924
925 Ok(schema_name)
926}
927
928/// Encapsulates default implementation of [`AggregateUDFImpl::window_function_schema_name`].
929pub fn udaf_default_window_function_schema_name<F: AggregateUDFImpl + ?Sized>(
930 func: &F,
931 params: &WindowFunctionParams,
932) -> Result<String> {
933 let WindowFunctionParams {
934 args,
935 partition_by,
936 order_by,
937 window_frame,
938 filter,
939 null_treatment,
940 distinct,
941 } = params;
942
943 let mut schema_name = String::new();
944
945 // Inject DISTINCT into the schema name when requested
946 if *distinct {
947 schema_name.write_fmt(format_args!(
948 "{}(DISTINCT {})",
949 func.name(),
950 schema_name_from_exprs(args)?
951 ))?;
952 } else {
953 schema_name.write_fmt(format_args!(
954 "{}({})",
955 func.name(),
956 schema_name_from_exprs(args)?
957 ))?;
958 }
959
960 if let Some(null_treatment) = null_treatment {
961 schema_name.write_fmt(format_args!(" {null_treatment}"))?;
962 }
963
964 if let Some(filter) = filter {
965 schema_name.write_fmt(format_args!(" FILTER (WHERE {filter})"))?;
966 }
967
968 if !partition_by.is_empty() {
969 schema_name.write_fmt(format_args!(
970 " PARTITION BY [{}]",
971 schema_name_from_exprs(partition_by)?
972 ))?;
973 }
974
975 if !order_by.is_empty() {
976 schema_name.write_fmt(format_args!(
977 " ORDER BY [{}]",
978 schema_name_from_sorts(order_by)?
979 ))?;
980 }
981
982 schema_name.write_fmt(format_args!(" {window_frame}"))?;
983
984 Ok(schema_name)
985}
986
987/// Encapsulates default implementation of [`AggregateUDFImpl::display_name`].
988pub fn udaf_default_display_name<F: AggregateUDFImpl + ?Sized>(
989 func: &F,
990 params: &AggregateFunctionParams,
991) -> Result<String> {
992 let AggregateFunctionParams {
993 args,
994 distinct,
995 filter,
996 order_by,
997 null_treatment,
998 } = params;
999
1000 let mut display_name = String::new();
1001
1002 display_name.write_fmt(format_args!(
1003 "{}({}{})",
1004 func.name(),
1005 if *distinct { "DISTINCT " } else { "" },
1006 expr_vec_fmt!(args)
1007 ))?;
1008
1009 if let Some(nt) = null_treatment {
1010 display_name.write_fmt(format_args!(" {nt}"))?;
1011 }
1012 if let Some(fe) = filter {
1013 display_name.write_fmt(format_args!(" FILTER (WHERE {fe})"))?;
1014 }
1015 if !order_by.is_empty() {
1016 display_name.write_fmt(format_args!(
1017 " ORDER BY [{}]",
1018 order_by
1019 .iter()
1020 .map(|o| format!("{o}"))
1021 .collect::<Vec<String>>()
1022 .join(", ")
1023 ))?;
1024 }
1025
1026 Ok(display_name)
1027}
1028
1029/// Encapsulates default implementation of [`AggregateUDFImpl::window_function_display_name`].
1030pub fn udaf_default_window_function_display_name<F: AggregateUDFImpl + ?Sized>(
1031 func: &F,
1032 params: &WindowFunctionParams,
1033) -> Result<String> {
1034 let WindowFunctionParams {
1035 args,
1036 partition_by,
1037 order_by,
1038 window_frame,
1039 filter,
1040 null_treatment,
1041 distinct,
1042 } = params;
1043
1044 let mut display_name = String::new();
1045
1046 if *distinct {
1047 display_name.write_fmt(format_args!(
1048 "{}(DISTINCT {})",
1049 func.name(),
1050 expr_vec_fmt!(args)
1051 ))?;
1052 } else {
1053 display_name.write_fmt(format_args!(
1054 "{}({})",
1055 func.name(),
1056 expr_vec_fmt!(args)
1057 ))?;
1058 }
1059
1060 if let Some(null_treatment) = null_treatment {
1061 display_name.write_fmt(format_args!(" {null_treatment}"))?;
1062 }
1063
1064 if let Some(fe) = filter {
1065 display_name.write_fmt(format_args!(" FILTER (WHERE {fe})"))?;
1066 }
1067
1068 if !partition_by.is_empty() {
1069 display_name.write_fmt(format_args!(
1070 " PARTITION BY [{}]",
1071 expr_vec_fmt!(partition_by)
1072 ))?;
1073 }
1074
1075 if !order_by.is_empty() {
1076 display_name
1077 .write_fmt(format_args!(" ORDER BY [{}]", expr_vec_fmt!(order_by)))?;
1078 };
1079
1080 display_name.write_fmt(format_args!(
1081 " {} BETWEEN {} AND {}",
1082 window_frame.units, window_frame.start_bound, window_frame.end_bound
1083 ))?;
1084
1085 Ok(display_name)
1086}
1087
1088/// Encapsulates default implementation of [`AggregateUDFImpl::return_field`].
1089pub fn udaf_default_return_field<F: AggregateUDFImpl + ?Sized>(
1090 func: &F,
1091 arg_fields: &[FieldRef],
1092) -> Result<FieldRef> {
1093 let arg_types: Vec<_> = arg_fields.iter().map(|f| f.data_type()).cloned().collect();
1094 let data_type = func.return_type(&arg_types)?;
1095
1096 Ok(Arc::new(Field::new(
1097 func.name(),
1098 data_type,
1099 func.is_nullable(),
1100 )))
1101}
1102
1103pub enum ReversedUDAF {
1104 /// The expression is the same as the original expression, like SUM, COUNT
1105 Identical,
1106 /// The expression does not support reverse calculation
1107 NotSupported,
1108 /// The expression is different from the original expression
1109 Reversed(Arc<AggregateUDF>),
1110}
1111
1112/// AggregateUDF that adds an alias to the underlying function. It is better to
1113/// implement [`AggregateUDFImpl`], which supports aliases, directly if possible.
1114#[derive(Debug, PartialEq, Eq, Hash)]
1115struct AliasedAggregateUDFImpl {
1116 inner: UdfEq<Arc<dyn AggregateUDFImpl>>,
1117 aliases: Vec<String>,
1118}
1119
1120impl AliasedAggregateUDFImpl {
1121 pub fn new(
1122 inner: Arc<dyn AggregateUDFImpl>,
1123 new_aliases: impl IntoIterator<Item = &'static str>,
1124 ) -> Self {
1125 let mut aliases = inner.aliases().to_vec();
1126 aliases.extend(new_aliases.into_iter().map(|s| s.to_string()));
1127
1128 Self {
1129 inner: inner.into(),
1130 aliases,
1131 }
1132 }
1133}
1134
1135#[warn(clippy::missing_trait_methods)] // Delegates, so it should implement every single trait method
1136impl AggregateUDFImpl for AliasedAggregateUDFImpl {
1137 fn as_any(&self) -> &dyn Any {
1138 self
1139 }
1140
1141 fn name(&self) -> &str {
1142 self.inner.name()
1143 }
1144
1145 fn signature(&self) -> &Signature {
1146 self.inner.signature()
1147 }
1148
1149 fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
1150 self.inner.return_type(arg_types)
1151 }
1152
1153 fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
1154 self.inner.accumulator(acc_args)
1155 }
1156
1157 fn aliases(&self) -> &[String] {
1158 &self.aliases
1159 }
1160
1161 fn schema_name(&self, params: &AggregateFunctionParams) -> Result<String> {
1162 self.inner.schema_name(params)
1163 }
1164
1165 fn human_display(&self, params: &AggregateFunctionParams) -> Result<String> {
1166 self.inner.human_display(params)
1167 }
1168
1169 fn window_function_schema_name(
1170 &self,
1171 params: &WindowFunctionParams,
1172 ) -> Result<String> {
1173 self.inner.window_function_schema_name(params)
1174 }
1175
1176 fn display_name(&self, params: &AggregateFunctionParams) -> Result<String> {
1177 self.inner.display_name(params)
1178 }
1179
1180 fn window_function_display_name(
1181 &self,
1182 params: &WindowFunctionParams,
1183 ) -> Result<String> {
1184 self.inner.window_function_display_name(params)
1185 }
1186
1187 fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
1188 self.inner.state_fields(args)
1189 }
1190
1191 fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
1192 self.inner.groups_accumulator_supported(args)
1193 }
1194
1195 fn create_groups_accumulator(
1196 &self,
1197 args: AccumulatorArgs,
1198 ) -> Result<Box<dyn GroupsAccumulator>> {
1199 self.inner.create_groups_accumulator(args)
1200 }
1201
1202 fn create_sliding_accumulator(
1203 &self,
1204 args: AccumulatorArgs,
1205 ) -> Result<Box<dyn Accumulator>> {
1206 self.inner.accumulator(args)
1207 }
1208
1209 fn with_beneficial_ordering(
1210 self: Arc<Self>,
1211 beneficial_ordering: bool,
1212 ) -> Result<Option<Arc<dyn AggregateUDFImpl>>> {
1213 Arc::clone(&self.inner)
1214 .with_beneficial_ordering(beneficial_ordering)
1215 .map(|udf| {
1216 udf.map(|udf| {
1217 Arc::new(AliasedAggregateUDFImpl {
1218 inner: udf.into(),
1219 aliases: self.aliases.clone(),
1220 }) as Arc<dyn AggregateUDFImpl>
1221 })
1222 })
1223 }
1224
1225 fn order_sensitivity(&self) -> AggregateOrderSensitivity {
1226 self.inner.order_sensitivity()
1227 }
1228
1229 fn simplify(&self) -> Option<AggregateFunctionSimplification> {
1230 self.inner.simplify()
1231 }
1232
1233 fn reverse_expr(&self) -> ReversedUDAF {
1234 self.inner.reverse_expr()
1235 }
1236
1237 fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
1238 self.inner.coerce_types(arg_types)
1239 }
1240
1241 fn return_field(&self, arg_fields: &[FieldRef]) -> Result<FieldRef> {
1242 self.inner.return_field(arg_fields)
1243 }
1244
1245 fn is_nullable(&self) -> bool {
1246 self.inner.is_nullable()
1247 }
1248
1249 fn is_descending(&self) -> Option<bool> {
1250 self.inner.is_descending()
1251 }
1252
1253 fn value_from_stats(&self, statistics_args: &StatisticsArgs) -> Option<ScalarValue> {
1254 self.inner.value_from_stats(statistics_args)
1255 }
1256
1257 fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
1258 self.inner.default_value(data_type)
1259 }
1260
1261 fn supports_null_handling_clause(&self) -> bool {
1262 self.inner.supports_null_handling_clause()
1263 }
1264
1265 fn supports_within_group_clause(&self) -> bool {
1266 self.inner.supports_within_group_clause()
1267 }
1268
1269 fn set_monotonicity(&self, data_type: &DataType) -> SetMonotonicity {
1270 self.inner.set_monotonicity(data_type)
1271 }
1272
1273 fn documentation(&self) -> Option<&Documentation> {
1274 self.inner.documentation()
1275 }
1276}
1277
1278/// Indicates whether an aggregation function is monotonic as a set
1279/// function. A set function is monotonically increasing if its value
1280/// increases as its argument grows (as a set). Formally, `f` is a
1281/// monotonically increasing set function if `f(S) >= f(T)` whenever `S`
1282/// is a superset of `T`.
1283///
1284/// For example `COUNT` and `MAX` are monotonically increasing as their
1285/// values always increase (or stay the same) as new values are seen. On
1286/// the other hand, `MIN` is monotonically decreasing as its value always
1287/// decreases or stays the same as new values are seen.
1288#[derive(Debug, Clone, PartialEq)]
1289pub enum SetMonotonicity {
1290 /// Aggregate value increases or stays the same as the input set grows.
1291 Increasing,
1292 /// Aggregate value decreases or stays the same as the input set grows.
1293 Decreasing,
1294 /// Aggregate value may increase, decrease, or stay the same as the input
1295 /// set grows.
1296 NotMonotonic,
1297}
1298
1299#[cfg(test)]
1300mod test {
1301 use crate::{AggregateUDF, AggregateUDFImpl};
1302 use arrow::datatypes::{DataType, FieldRef};
1303 use datafusion_common::Result;
1304 use datafusion_expr_common::accumulator::Accumulator;
1305 use datafusion_expr_common::signature::{Signature, Volatility};
1306 use datafusion_functions_aggregate_common::accumulator::{
1307 AccumulatorArgs, StateFieldsArgs,
1308 };
1309 use std::any::Any;
1310 use std::cmp::Ordering;
1311 use std::hash::{DefaultHasher, Hash, Hasher};
1312
1313 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
1314 struct AMeanUdf {
1315 signature: Signature,
1316 }
1317
1318 impl AMeanUdf {
1319 fn new() -> Self {
1320 Self {
1321 signature: Signature::uniform(
1322 1,
1323 vec![DataType::Float64],
1324 Volatility::Immutable,
1325 ),
1326 }
1327 }
1328 }
1329
1330 impl AggregateUDFImpl for AMeanUdf {
1331 fn as_any(&self) -> &dyn Any {
1332 self
1333 }
1334 fn name(&self) -> &str {
1335 "a"
1336 }
1337 fn signature(&self) -> &Signature {
1338 &self.signature
1339 }
1340 fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
1341 unimplemented!()
1342 }
1343 fn accumulator(
1344 &self,
1345 _acc_args: AccumulatorArgs,
1346 ) -> Result<Box<dyn Accumulator>> {
1347 unimplemented!()
1348 }
1349 fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
1350 unimplemented!()
1351 }
1352 }
1353
1354 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
1355 struct BMeanUdf {
1356 signature: Signature,
1357 }
1358 impl BMeanUdf {
1359 fn new() -> Self {
1360 Self {
1361 signature: Signature::uniform(
1362 1,
1363 vec![DataType::Float64],
1364 Volatility::Immutable,
1365 ),
1366 }
1367 }
1368 }
1369
1370 impl AggregateUDFImpl for BMeanUdf {
1371 fn as_any(&self) -> &dyn Any {
1372 self
1373 }
1374 fn name(&self) -> &str {
1375 "b"
1376 }
1377 fn signature(&self) -> &Signature {
1378 &self.signature
1379 }
1380 fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
1381 unimplemented!()
1382 }
1383 fn accumulator(
1384 &self,
1385 _acc_args: AccumulatorArgs,
1386 ) -> Result<Box<dyn Accumulator>> {
1387 unimplemented!()
1388 }
1389 fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
1390 unimplemented!()
1391 }
1392 }
1393
1394 #[test]
1395 fn test_partial_eq() {
1396 let a1 = AggregateUDF::from(AMeanUdf::new());
1397 let a2 = AggregateUDF::from(AMeanUdf::new());
1398 let eq = a1 == a2;
1399 assert!(eq);
1400 assert_eq!(a1, a2);
1401 assert_eq!(hash(a1), hash(a2));
1402 }
1403
1404 #[test]
1405 fn test_partial_ord() {
1406 // Test validates that partial ord is defined for AggregateUDF using the name and signature,
1407 // not intended to exhaustively test all possibilities
1408 let a1 = AggregateUDF::from(AMeanUdf::new());
1409 let a2 = AggregateUDF::from(AMeanUdf::new());
1410 assert_eq!(a1.partial_cmp(&a2), Some(Ordering::Equal));
1411
1412 let b1 = AggregateUDF::from(BMeanUdf::new());
1413 assert!(a1 < b1);
1414 assert!(!(a1 == b1));
1415 }
1416
1417 fn hash<T: Hash>(value: T) -> u64 {
1418 let hasher = &mut DefaultHasher::new();
1419 value.hash(hasher);
1420 hasher.finish()
1421 }
1422}