Trait AggregateUDFImpl

Source

pub trait AggregateUDFImpl:
    Debug
    + Send
    + Sync {
Show 28 methods    // Required methods
    fn as_any(&self) -> &(dyn Any + 'static);
    fn name(&self) -> &str;
    fn signature(&self) -> &Signature;
    fn return_type(
        &self,
        arg_types: &[DataType],
    ) -> Result<DataType, DataFusionError>;
    fn accumulator(
        &self,
        acc_args: AccumulatorArgs<'_>,
    ) -> Result<Box<dyn Accumulator>, DataFusionError>;

    // Provided methods
    fn schema_name(
        &self,
        params: &AggregateFunctionParams,
    ) -> Result<String, DataFusionError> { ... }
    fn human_display(
        &self,
        params: &AggregateFunctionParams,
    ) -> Result<String, DataFusionError> { ... }
    fn window_function_schema_name(
        &self,
        params: &WindowFunctionParams,
    ) -> Result<String, DataFusionError> { ... }
    fn display_name(
        &self,
        params: &AggregateFunctionParams,
    ) -> Result<String, DataFusionError> { ... }
    fn window_function_display_name(
        &self,
        params: &WindowFunctionParams,
    ) -> Result<String, DataFusionError> { ... }
    fn is_nullable(&self) -> bool { ... }
    fn state_fields(
        &self,
        args: StateFieldsArgs<'_>,
    ) -> Result<Vec<Field>, DataFusionError> { ... }
    fn groups_accumulator_supported(&self, _args: AccumulatorArgs<'_>) -> bool { ... }
    fn create_groups_accumulator(
        &self,
        _args: AccumulatorArgs<'_>,
    ) -> Result<Box<dyn GroupsAccumulator>, DataFusionError> { ... }
    fn aliases(&self) -> &[String] { ... }
    fn create_sliding_accumulator(
        &self,
        args: AccumulatorArgs<'_>,
    ) -> Result<Box<dyn Accumulator>, DataFusionError> { ... }
    fn with_beneficial_ordering(
        self: Arc<Self>,
        _beneficial_ordering: bool,
    ) -> Result<Option<Arc<dyn AggregateUDFImpl>>, DataFusionError> { ... }
    fn order_sensitivity(&self) -> AggregateOrderSensitivity { ... }
    fn simplify(
        &self,
    ) -> Option<Box<dyn Fn(AggregateFunction, &dyn SimplifyInfo) -> Result<Expr, DataFusionError>>> { ... }
    fn reverse_expr(&self) -> ReversedUDAF { ... }
    fn coerce_types(
        &self,
        _arg_types: &[DataType],
    ) -> Result<Vec<DataType>, DataFusionError> { ... }
    fn equals(&self, other: &dyn AggregateUDFImpl) -> bool { ... }
    fn hash_value(&self) -> u64 { ... }
    fn is_descending(&self) -> Option<bool> { ... }
    fn value_from_stats(
        &self,
        _statistics_args: &StatisticsArgs<'_>,
    ) -> Option<ScalarValue> { ... }
    fn default_value(
        &self,
        data_type: &DataType,
    ) -> Result<ScalarValue, DataFusionError> { ... }
    fn documentation(&self) -> Option<&Documentation> { ... }
    fn set_monotonicity(&self, _data_type: &DataType) -> SetMonotonicity { ... }
}

Expand description

Trait for implementing AggregateUDF.

This trait exposes the full API for implementing user defined aggregate functions and can be used to implement any function.

See advanced_udaf.rs for a full example with complete implementation and AggregateUDF for other available options.

§Basic Example


#[derive(Debug, Clone)]
struct GeoMeanUdf {
  signature: Signature,
}

impl GeoMeanUdf {
  fn new() -> Self {
    Self {
      signature: Signature::uniform(1, vec![DataType::Float64], Volatility::Immutable),
     }
  }
}

static DOCUMENTATION: LazyLock<Documentation> = LazyLock::new(|| {
        Documentation::builder(DOC_SECTION_AGGREGATE, "calculates a geometric mean", "geo_mean(2.0)")
            .with_argument("arg1", "The Float64 number for the geometric mean")
            .build()
    });

fn get_doc() -> &'static Documentation {
    &DOCUMENTATION
}
    
/// Implement the AggregateUDFImpl trait for GeoMeanUdf
impl AggregateUDFImpl for GeoMeanUdf {
   fn as_any(&self) -> &dyn Any { self }
   fn name(&self) -> &str { "geo_mean" }
   fn signature(&self) -> &Signature { &self.signature }
   fn return_type(&self, args: &[DataType]) -> Result<DataType> {
     if !matches!(args.get(0), Some(&DataType::Float64)) {
       return plan_err!("geo_mean only accepts Float64 arguments");
     }
     Ok(DataType::Float64)
   }
   // This is the accumulator factory; DataFusion uses it to create new accumulators.
   fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { unimplemented!() }
   fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
       Ok(vec![
            Field::new("value", args.return_type.clone(), true),
            Field::new("ordering", DataType::UInt32, true)
       ])
   }
   fn documentation(&self) -> Option<&Documentation> {
       Some(get_doc())  
   }
}

// Create a new AggregateUDF from the implementation
let geometric_mean = AggregateUDF::from(GeoMeanUdf::new());

// Call the function `geo_mean(col)`
let expr = geometric_mean.call(vec![col("a")]);

Required Methods§

Source

fn as_any(&self) -> &(dyn Any + 'static)

Returns this object as an Any trait object

Source

fn name(&self) -> &str

Returns this function’s name

Source

fn signature(&self) -> &Signature

Returns the function’s Signature for information about what input types are accepted and the function’s Volatility.

Source

fn return_type( &self, arg_types: &[DataType], ) -> Result<DataType, DataFusionError>

What DataType will be returned by this function, given the types of the arguments

Source

fn accumulator( &self, acc_args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

Return a new Accumulator that aggregates values for a specific group during query execution.

acc_args: AccumulatorArgs contains information about how the aggregate function was called.

Provided Methods§

Source

fn schema_name( &self, params: &AggregateFunctionParams, ) -> Result<String, DataFusionError>

Returns the name of the column this expression would create

See Expr::schema_name for details

Example of schema_name: count(DISTINCT column1) FILTER (WHERE column2 > 10) ORDER BY [..]

Source

fn human_display( &self, params: &AggregateFunctionParams, ) -> Result<String, DataFusionError>

Returns a human readable expression.

See Expr::human_display for details.

Source

fn window_function_schema_name( &self, params: &WindowFunctionParams, ) -> Result<String, DataFusionError>

Returns the name of the column this expression would create

See Expr::schema_name for details

Different from schema_name in that it is used for window aggregate function

Example of schema_name: count(DISTINCT column1) FILTER (WHERE column2 > 10) [PARTITION BY [..]] [ORDER BY [..]]

Source

fn display_name( &self, params: &AggregateFunctionParams, ) -> Result<String, DataFusionError>

Returns the user-defined display name of function, given the arguments

This can be used to customize the output column name generated by this function.

Defaults to function_name([DISTINCT] column1, column2, ..) [null_treatment] [filter] [order_by [..]]

Source

fn window_function_display_name( &self, params: &WindowFunctionParams, ) -> Result<String, DataFusionError>

Returns the user-defined display name of function, given the arguments

This can be used to customize the output column name generated by this function.

Different from display_name in that it is used for window aggregate function

Defaults to function_name([DISTINCT] column1, column2, ..) [null_treatment] [partition by [..]] [order_by [..]]

Source

fn is_nullable(&self) -> bool

Whether the aggregate function is nullable.

Nullable means that the function could return null for any inputs. For example, aggregate functions like COUNT always return a non null value but others like MIN will return NULL if there is nullable input. Note that if the function is declared as not nullable, make sure the AggregateUDFImpl::default_value is non-null

Source

fn state_fields( &self, args: StateFieldsArgs<'_>, ) -> Result<Vec<Field>, DataFusionError>

Return the fields used to store the intermediate state of this accumulator.

See Accumulator::state for background information.

args: StateFieldsArgs contains arguments passed to the aggregate function’s accumulator.

§Notes:

The default implementation returns a single state field named name with the same type as value_type. This is suitable for aggregates such as SUM or MIN where partial state can be combined by applying the same aggregate.

For aggregates such as AVG where the partial state is more complex (e.g. a COUNT and a SUM), this method is used to define the additional fields.

The name of the fields must be unique within the query and thus should be derived from name. See format_state_name for a utility function to generate a unique name.

Source

fn groups_accumulator_supported(&self, _args: AccumulatorArgs<'_>) -> bool

If the aggregate expression has a specialized GroupsAccumulator implementation. If this returns true, [Self::create_groups_accumulator] will be called.

§Notes

Even if this function returns true, DataFusion will still use Self::accumulator for certain queries, such as when this aggregate is used as a window function or when there no GROUP BY columns in the query.

Source

fn create_groups_accumulator( &self, _args: AccumulatorArgs<'_>, ) -> Result<Box<dyn GroupsAccumulator>, DataFusionError>

Return a specialized GroupsAccumulator that manages state for all groups.

For maximum performance, a GroupsAccumulator should be implemented in addition to Accumulator.

Source

fn aliases(&self) -> &[String]

Returns any aliases (alternate names) for this function.

Note: aliases should only include names other than Self::name. Defaults to [] (no aliases)

Source

fn create_sliding_accumulator( &self, args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

Sliding accumulator is an alternative accumulator that can be used for window functions. It has retract method to revert the previous update.

See retract_batch for more details.

Source

fn with_beneficial_ordering( self: Arc<Self>, _beneficial_ordering: bool, ) -> Result<Option<Arc<dyn AggregateUDFImpl>>, DataFusionError>

Sets the indicator whether ordering requirements of the AggregateUDFImpl is satisfied by its input. If this is not the case, UDFs with order sensitivity AggregateOrderSensitivity::Beneficial can still produce the correct result with possibly more work internally.

§Returns

Returns Ok(Some(updated_udf)) if the process completes successfully. If the expression can benefit from existing input ordering, but does not implement the method, returns an error. Order insensitive and hard requirement aggregators return Ok(None).

Source

fn order_sensitivity(&self) -> AggregateOrderSensitivity

Gets the order sensitivity of the UDF. See AggregateOrderSensitivity for possible options.

Source

fn simplify( &self, ) -> Option<Box<dyn Fn(AggregateFunction, &dyn SimplifyInfo) -> Result<Expr, DataFusionError>>>

Optionally apply per-UDaF simplification / rewrite rules.

This can be used to apply function specific simplification rules during optimization (e.g. arrow_cast –> Expr::Cast). The default implementation does nothing.

Note that DataFusion handles simplifying arguments and “constant folding” (replacing a function call with constant arguments such as my_add(1,2) --> 3 ). Thus, there is no need to implement such optimizations manually for specific UDFs.

§Returns

None if simplify is not defined or,

Or, a closure with two arguments:

‘aggregate_function’: crate::expr::AggregateFunction for which simplified has been invoked
‘info’: crate::simplify::SimplifyInfo

closure returns simplified Expr or an error.

Source

fn reverse_expr(&self) -> ReversedUDAF

Returns the reverse expression of the aggregate function.

Source

fn coerce_types( &self, _arg_types: &[DataType], ) -> Result<Vec<DataType>, DataFusionError>

Coerce arguments of a function call to types that the function can evaluate.

This function is only called if AggregateUDFImpl::signature returns crate::TypeSignature::UserDefined. Most UDAFs should return one of the other variants of TypeSignature which handle common cases

See the type coercion module documentation for more details on type coercion

For example, if your function requires a floating point arguments, but the user calls it like my_func(1::int) (aka with 1 as an integer), coerce_types could return [DataType::Float64] to ensure the argument was cast to 1::double

§Parameters

arg_types: The argument types of the arguments this function with

§Return value

A Vec the same length as arg_types. DataFusion will CAST the function call arguments to these specific types.

Source

fn equals(&self, other: &dyn AggregateUDFImpl) -> bool

Return true if this aggregate UDF is equal to the other.

Allows customizing the equality of aggregate UDFs. Must be consistent with Self::hash_value and follow the same rules as Eq:

reflexive: a.equals(a);
symmetric: a.equals(b) implies b.equals(a);
transitive: a.equals(b) and b.equals(c) implies a.equals(c).

By default, compares Self::name and Self::signature.

Source

fn hash_value(&self) -> u64

Returns a hash value for this aggregate UDF.

Allows customizing the hash code of aggregate UDFs. Similarly to Hash and Eq, if Self::equals returns true for two UDFs, their hash_values must be the same.

By default, hashes Self::name and Self::signature.

Source

fn is_descending(&self) -> Option<bool>

If this function is max, return true If the function is min, return false Otherwise return None (the default)

Note: this is used to use special aggregate implementations in certain conditions

Source

fn value_from_stats( &self, _statistics_args: &StatisticsArgs<'_>, ) -> Option<ScalarValue>

Return the value of this aggregate function if it can be determined entirely from statistics and arguments.

Using a ScalarValue rather than a runtime computation can significantly improving query performance.

For example, if the minimum value of column x is known to be 42 from statistics, then the aggregate MIN(x) should return Some(ScalarValue(42))

Source

fn default_value( &self, data_type: &DataType, ) -> Result<ScalarValue, DataFusionError>

Returns default value of the function given the input is all null.

Most of the aggregate function return Null if input is Null, while count returns 0 if input is Null

Source

fn documentation(&self) -> Option<&Documentation>

Returns the documentation for this Aggregate UDF.

Documentation can be accessed programmatically as well as generating publicly facing documentation.

Source

fn set_monotonicity(&self, _data_type: &DataType) -> SetMonotonicity

Indicates whether the aggregation function is monotonic as a set function. See SetMonotonicity for details.

Trait AggregateUDFImplCopy item path

§Basic Example

Required Methods§

fn as_any(&self) -> &(dyn Any + 'static)

fn name(&self) -> &str

fn signature(&self) -> &Signature

fn return_type( &self, arg_types: &[DataType], ) -> Result<DataType, DataFusionError>

fn accumulator( &self, acc_args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

Provided Methods§

fn schema_name( &self, params: &AggregateFunctionParams, ) -> Result<String, DataFusionError>

fn human_display( &self, params: &AggregateFunctionParams, ) -> Result<String, DataFusionError>

fn window_function_schema_name( &self, params: &WindowFunctionParams, ) -> Result<String, DataFusionError>

fn display_name( &self, params: &AggregateFunctionParams, ) -> Result<String, DataFusionError>

fn window_function_display_name( &self, params: &WindowFunctionParams, ) -> Result<String, DataFusionError>

fn is_nullable(&self) -> bool

fn state_fields( &self, args: StateFieldsArgs<'_>, ) -> Result<Vec<Field>, DataFusionError>

§Notes:

fn groups_accumulator_supported(&self, _args: AccumulatorArgs<'_>) -> bool

§Notes

fn create_groups_accumulator( &self, _args: AccumulatorArgs<'_>, ) -> Result<Box<dyn GroupsAccumulator>, DataFusionError>

fn aliases(&self) -> &[String]

fn create_sliding_accumulator( &self, args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

fn with_beneficial_ordering( self: Arc<Self>, _beneficial_ordering: bool, ) -> Result<Option<Arc<dyn AggregateUDFImpl>>, DataFusionError>

§Returns

fn order_sensitivity(&self) -> AggregateOrderSensitivity

fn simplify( &self, ) -> Option<Box<dyn Fn(AggregateFunction, &dyn SimplifyInfo) -> Result<Expr, DataFusionError>>>

§Returns

fn reverse_expr(&self) -> ReversedUDAF

fn coerce_types( &self, _arg_types: &[DataType], ) -> Result<Vec<DataType>, DataFusionError>

§Parameters

§Return value

fn equals(&self, other: &dyn AggregateUDFImpl) -> bool

fn hash_value(&self) -> u64

fn is_descending(&self) -> Option<bool>

fn value_from_stats( &self, _statistics_args: &StatisticsArgs<'_>, ) -> Option<ScalarValue>

fn default_value( &self, data_type: &DataType, ) -> Result<ScalarValue, DataFusionError>

fn documentation(&self) -> Option<&Documentation>

fn set_monotonicity(&self, _data_type: &DataType) -> SetMonotonicity

Trait Implementations§

impl PartialEq for dyn AggregateUDFImpl

fn eq(&self, other: &(dyn AggregateUDFImpl + 'static)) -> bool

fn ne(&self, other: &Rhs) -> bool

impl PartialOrd for dyn AggregateUDFImpl

fn partial_cmp( &self, other: &(dyn AggregateUDFImpl + 'static), ) -> Option<Ordering>

fn lt(&self, other: &Rhs) -> bool

fn le(&self, other: &Rhs) -> bool

fn gt(&self, other: &Rhs) -> bool

fn ge(&self, other: &Rhs) -> bool

Implementations on Foreign Types§

impl AggregateUDFImpl for ApproxDistinct

fn as_any(&self) -> &(dyn Any + 'static)

fn name(&self) -> &str

fn signature(&self) -> &Signature

fn return_type(&self, _: &[DataType]) -> Result<DataType, DataFusionError>

fn state_fields( &self, args: StateFieldsArgs<'_>, ) -> Result<Vec<Field>, DataFusionError>

fn accumulator( &self, acc_args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

fn documentation(&self) -> Option<&Documentation>

impl AggregateUDFImpl for ApproxMedian

fn as_any(&self) -> &(dyn Any + 'static)

fn state_fields( &self, args: StateFieldsArgs<'_>, ) -> Result<Vec<Field>, DataFusionError>

fn name(&self) -> &str

fn signature(&self) -> &Signature

fn return_type( &self, arg_types: &[DataType], ) -> Result<DataType, DataFusionError>

fn accumulator( &self, acc_args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

fn documentation(&self) -> Option<&Documentation>

impl AggregateUDFImpl for ApproxPercentileCont

fn state_fields( &self, args: StateFieldsArgs<'_>, ) -> Result<Vec<Field>, DataFusionError>

fn as_any(&self) -> &(dyn Any + 'static)

fn name(&self) -> &str

fn signature(&self) -> &Signature

fn accumulator( &self, acc_args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

fn return_type( &self, arg_types: &[DataType], ) -> Result<DataType, DataFusionError>

fn documentation(&self) -> Option<&Documentation>

impl AggregateUDFImpl for ApproxPercentileContWithWeight

fn state_fields( &self, args: StateFieldsArgs<'_>, ) -> Result<Vec<Field>, DataFusionError>

fn as_any(&self) -> &(dyn Any + 'static)

fn name(&self) -> &str

fn signature(&self) -> &Signature

fn return_type( &self, arg_types: &[DataType], ) -> Result<DataType, DataFusionError>

fn accumulator( &self, acc_args: AccumulatorArgs<'_>, ) -> Result<Box<dyn Accumulator>, DataFusionError>

Trait AggregateUDFImpl