-
Notifications
You must be signed in to change notification settings - Fork 1.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Make fields of ScalarUDF
, AggregateUDF
and WindowUDF
non pub
#8079
Changes from 2 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -15,12 +15,14 @@ | |
// specific language governing permissions and limitations | ||
// under the License. | ||
|
||
//! Udaf module contains functions and structs supporting user-defined aggregate functions. | ||
//! [`AggregrateUDF`]: User Defined Aggregate Functions | ||
|
||
use crate::Expr; | ||
use crate::{Accumulator, Expr}; | ||
use crate::{ | ||
AccumulatorFactoryFunction, ReturnTypeFunction, Signature, StateTypeFunction, | ||
}; | ||
use arrow::datatypes::DataType; | ||
use datafusion_common::Result; | ||
use std::fmt::{self, Debug, Formatter}; | ||
use std::sync::Arc; | ||
|
||
|
@@ -46,15 +48,15 @@ use std::sync::Arc; | |
#[derive(Clone)] | ||
pub struct AggregateUDF { | ||
/// name | ||
pub name: String, | ||
name: String, | ||
/// Signature (input arguments) | ||
pub signature: Signature, | ||
signature: Signature, | ||
/// Return type | ||
pub return_type: ReturnTypeFunction, | ||
return_type: ReturnTypeFunction, | ||
/// actual implementation | ||
pub accumulator: AccumulatorFactoryFunction, | ||
accumulator: AccumulatorFactoryFunction, | ||
/// the accumulator's state's description as a function of the return type | ||
pub state_type: StateTypeFunction, | ||
state_type: StateTypeFunction, | ||
} | ||
|
||
impl Debug for AggregateUDF { | ||
|
@@ -112,4 +114,35 @@ impl AggregateUDF { | |
order_by: None, | ||
}) | ||
} | ||
|
||
/// Returns this function's name | ||
pub fn name(&self) -> &str { | ||
&self.name | ||
} | ||
|
||
/// Returns this function's signature (what input types are accepted) | ||
pub fn signature(&self) -> &Signature { | ||
&self.signature | ||
} | ||
|
||
/// Return the type of the function given its input types | ||
pub fn return_type(&self, args: &[DataType]) -> Result<DataType> { | ||
// Old API returns an Arc of the datatype for some reason | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Another change I did was to move the handling of function pointers into the implementation of This is in preparation for potentially changing the implementation of |
||
let res = (self.return_type)(args)?; | ||
Ok(res.as_ref().clone()) | ||
} | ||
|
||
/// Return an accumualator the given aggregate, given | ||
/// its return datatype. | ||
pub fn accumulator(&self, return_type: &DataType) -> Result<Box<dyn Accumulator>> { | ||
(self.accumulator)(return_type) | ||
} | ||
|
||
/// Return the type of the intermediate state used by this aggregator, given | ||
/// its return datatype. Supports multi-phase aggregations | ||
pub fn state_type(&self, return_type: &DataType) -> Result<Vec<DataType>> { | ||
// old API returns an Arc for some reason, try and unwrap it here | ||
let res = (self.state_type)(return_type)?; | ||
Ok(Arc::try_unwrap(res).unwrap_or_else(|res| res.as_ref().clone())) | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -171,12 +171,8 @@ impl WindowFunction { | |
WindowFunction::BuiltInWindowFunction(fun) => { | ||
fun.return_type(input_expr_types) | ||
} | ||
WindowFunction::AggregateUDF(fun) => { | ||
Ok((*(fun.return_type)(input_expr_types)?).clone()) | ||
} | ||
WindowFunction::WindowUDF(fun) => { | ||
Ok((*(fun.return_type)(input_expr_types)?).clone()) | ||
} | ||
WindowFunction::AggregateUDF(fun) => fun.return_type(input_expr_types), | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think this now reads more nicely and idomatically |
||
WindowFunction::WindowUDF(fun) => fun.return_type(input_expr_types), | ||
} | ||
} | ||
} | ||
|
@@ -234,8 +230,8 @@ impl WindowFunction { | |
match self { | ||
WindowFunction::AggregateFunction(fun) => fun.signature(), | ||
WindowFunction::BuiltInWindowFunction(fun) => fun.signature(), | ||
WindowFunction::AggregateUDF(fun) => fun.signature.clone(), | ||
WindowFunction::WindowUDF(fun) => fun.signature.clone(), | ||
WindowFunction::AggregateUDF(fun) => fun.signature().clone(), | ||
WindowFunction::WindowUDF(fun) => fun.signature().clone(), | ||
} | ||
} | ||
} | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This piece of code looks really clean.