jayzhan211 commented on code in PR #10439:
URL: https://github.com/apache/datafusion/pull/10439#discussion_r1597436309


##########
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:
   Agree with you



-- 
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

Reply via email to