alamb commented on code in PR #10439:
URL: https://github.com/apache/datafusion/pull/10439#discussion_r1597427699
##########
datafusion/expr/src/type_coercion/functions.rs:
##########
@@ -20,23 +20,124 @@ use std::sync::Arc;
use crate::signature::{
ArrayFunctionSignature, FIXED_SIZE_LIST_WILDCARD, TIMEZONE_WILDCARD,
};
-use crate::{Signature, TypeSignature};
+use crate::{AggregateUDF, ScalarUDF, Signature, TypeSignature};
use arrow::{
compute::can_cast_types,
datatypes::{DataType, TimeUnit},
};
use datafusion_common::utils::{coerced_fixed_size_list_to_list, list_ndims};
-use datafusion_common::{internal_datafusion_err, internal_err, plan_err,
Result};
+use datafusion_common::{
+ exec_err, internal_datafusion_err, internal_err, plan_err, Result,
+};
use super::binary::{comparison_binary_numeric_coercion, comparison_coercion};
+/// Performs type coercion for scalar function arguments.
+///
+/// Returns the data types to which each argument must be coerced to
+/// match `signature`.
+///
+/// For more details on coercion in general, please see the
+/// [`type_coercion`](crate::type_coercion) module.
+pub fn data_types_with_scalar_udf(
+ current_types: &[DataType],
+ func: &ScalarUDF,
+) -> Result<Vec<DataType>> {
+ let signature = func.signature();
+
+ if current_types.is_empty() {
+ if signature.type_signature.supports_zero_argument() {
+ return Ok(vec![]);
+ } else {
+ return plan_err!(
+ "[data_types_with_scalar_udf] signature {:?} does not support
zero arguments.",
+ &signature.type_signature
+ );
+ }
+ }
+
+ let valid_types =
+ get_valid_types_with_scalar_udf(&signature.type_signature,
current_types, func)?;
+
+ if valid_types
+ .iter()
+ .any(|data_type| data_type == current_types)
+ {
+ return Ok(current_types.to_vec());
+ }
+
+ // Try and coerce the argument types to match the signature, returning the
+ // coerced types from the first matching signature.
+ for valid_types in valid_types {
+ if let Some(types) = maybe_data_types(&valid_types, current_types) {
+ return Ok(types);
+ }
+ }
+
+ // none possible -> Error
+ plan_err!(
+ "[data_types_with_scalar_udf] Coercion from {:?} to the signature {:?}
failed.",
+ current_types,
+ &signature.type_signature
+ )
+}
+
+pub fn data_types_with_aggregate_udf(
+ current_types: &[DataType],
+ func: &AggregateUDF,
+) -> Result<Vec<DataType>> {
+ let signature = func.signature();
+
+ if current_types.is_empty() {
+ if signature.type_signature.supports_zero_argument() {
+ return Ok(vec![]);
+ } else {
+ return plan_err!(
+ "[data_types_with_aggregate_udf] Coercion from {:?} to the
signature {:?} failed.",
+ current_types,
+ &signature.type_signature
+ );
+ }
+ }
+
+ let valid_types = get_valid_types_with_aggregate_udf(
+ &signature.type_signature,
+ current_types,
+ func,
+ )?;
+ if valid_types
+ .iter()
+ .any(|data_type| data_type == current_types)
+ {
+ return Ok(current_types.to_vec());
+ }
+
+ // Try and coerce the argument types to match the signature, returning the
+ // coerced types from the first matching signature.
+ for valid_types in valid_types {
+ if let Some(types) = maybe_data_types(&valid_types, current_types) {
+ return Ok(types);
+ }
+ }
+
+ // none possible -> Error
+ plan_err!(
+ "[data_types_with_aggregate_udf] Coercion from {:?} to the signature
{:?} failed.",
+ current_types,
+ &signature.type_signature
+ )
+}
+
/// Performs type coercion for function arguments.
///
/// Returns the data types to which each argument must be coerced to
/// match `signature`.
///
/// For more details on coercion in general, please see the
/// [`type_coercion`](crate::type_coercion) module.
+///
+/// This function will be replaced with [data_types_with_scalar_udf],
Review Comment:
I am not sure we want to replace this function over time -- I think having
the basic simple Signatures that handle most common coercions makes sense to
have in DataFusion core (even if it could be done purely in a udf) as it will
make creating UDFs easier for users
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
