fix: avoid internal errors for OneOf signature mismatches

Signed-off-by: yaommen <myanstu@163.com>

Author: myandpr

datafusion/expr/src/type_coercion/functions.rs

@@ -26,7 +26,9 @@ use std::vec;
 
 use arrow::datatypes::{DataType, Field, FieldRef};
 
-use datafusion_common::{Result, ScalarValue, Statistics, exec_err, not_impl_err};
+use datafusion_common::{
+    DataFusionError, Result, ScalarValue, Statistics, exec_err, not_impl_err,
+};
 use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
 use datafusion_expr_common::operator::Operator;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
@@ -517,6 +519,15 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     /// types are accepted and the function's Volatility.
     fn signature(&self) -> &Signature;
 
+    /// Allows a function to provide a more actionable signature failure
+    /// message than the generic `OneOf` fallback.
+    fn diagnose_failed_signature(
+        &self,
+        _arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        None
+    }
+
     /// What [`DataType`] will be returned by this function, given the types of
     /// the arguments
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;
@@ -1243,6 +1254,13 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl {
         self.inner.signature()
     }
 
+    fn diagnose_failed_signature(
+        &self,
+        arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        self.inner.diagnose_failed_signature(arg_types)
+    }
+
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         self.inner.return_type(arg_types)
     }

datafusion/expr/src/udaf.rs

@@ -28,7 +28,7 @@ use arrow::datatypes::{DataType, Field, FieldRef};
 #[cfg(debug_assertions)]
 use datafusion_common::assert_or_internal_err;
 use datafusion_common::config::ConfigOptions;
-use datafusion_common::{ExprSchema, Result, ScalarValue, not_impl_err};
+use datafusion_common::{DataFusionError, ExprSchema, Result, ScalarValue, not_impl_err};
 use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
 use datafusion_expr_common::interval_arithmetic::Interval;
 use datafusion_expr_common::placement::ExpressionPlacement;
@@ -572,6 +572,15 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     /// [`DataFusionError::Internal`]: datafusion_common::DataFusionError::Internal
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;
 
+    /// Allows a function to provide a more actionable signature failure
+    /// message than the generic `OneOf` fallback.
+    fn diagnose_failed_signature(
+        &self,
+        _arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        None
+    }
+
     /// Create a new instance of this function with updated configuration.
     ///
     /// This method is called when configuration options change at runtime
@@ -1033,6 +1042,13 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
         self.inner.signature()
     }
 
+    fn diagnose_failed_signature(
+        &self,
+        arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        self.inner.diagnose_failed_signature(arg_types)
+    }
+
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         self.inner.return_type(arg_types)
     }

datafusion/expr/src/udf.rs

@@ -33,7 +33,7 @@ use crate::udf_eq::UdfEq;
 use crate::{
     Expr, PartitionEvaluator, Signature, function::WindowFunctionSimplification,
 };
-use datafusion_common::{Result, not_impl_err};
+use datafusion_common::{DataFusionError, Result, not_impl_err};
 use datafusion_doc::Documentation;
 use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
 use datafusion_functions_window_common::expr::ExpressionArgs;
@@ -333,6 +333,15 @@ pub trait WindowUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     /// types are accepted and the function's Volatility.
     fn signature(&self) -> &Signature;
 
+    /// Allows a function to provide a more actionable signature failure
+    /// message than the generic `OneOf` fallback.
+    fn diagnose_failed_signature(
+        &self,
+        _arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        None
+    }
+
     /// Returns the expressions that are passed to the [`PartitionEvaluator`].
     fn expressions(&self, expr_args: ExpressionArgs) -> Vec<Arc<dyn PhysicalExpr>> {
         expr_args.input_exprs().into()
@@ -514,6 +523,13 @@ impl WindowUDFImpl for AliasedWindowUDFImpl {
         self.inner.signature()
     }
 
+    fn diagnose_failed_signature(
+        &self,
+        arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        self.inner.diagnose_failed_signature(arg_types)
+    }
+
     fn expressions(&self, expr_args: ExpressionArgs) -> Vec<Arc<dyn PhysicalExpr>> {
         expr_args
             .input_exprs()

datafusion/expr/src/udwf.rs

@@ -32,7 +32,9 @@ use arrow::datatypes::{
     DurationSecondType, Field, FieldRef, Float64Type, TimeUnit, UInt64Type, i256,
 };
 use datafusion_common::types::{NativeType, logical_float64};
-use datafusion_common::{Result, ScalarValue, exec_err, not_impl_err};
+use datafusion_common::{
+    DataFusionError, Result, ScalarValue, exec_err, not_impl_err, plan_datafusion_err,
+};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
 use datafusion_expr::{
@@ -139,6 +141,18 @@ impl AggregateUDFImpl for Avg {
         &self.signature
     }
 
+    fn diagnose_failed_signature(
+        &self,
+        arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        match arg_types {
+            [DataType::Boolean] => {
+                Some(plan_datafusion_err!("Avg not supported for Boolean"))
+            }
+            _ => None,
+        }
+    }
+
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         match &arg_types[0] {
             DataType::Decimal32(precision, scale) => {

datafusion/functions-aggregate/src/average.rs

@@ -32,7 +32,10 @@ use datafusion_common::types::{
     NativeType, logical_float64, logical_int8, logical_int16, logical_int32,
     logical_int64, logical_uint8, logical_uint16, logical_uint32, logical_uint64,
 };
-use datafusion_common::{HashMap, Result, ScalarValue, exec_err, not_impl_err};
+use datafusion_common::{
+    DataFusionError, HashMap, Result, ScalarValue, exec_err, not_impl_err,
+    plan_datafusion_err,
+};
 use datafusion_expr::expr::AggregateFunction;
 use datafusion_expr::expr_fn::cast;
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
@@ -213,6 +216,18 @@ impl AggregateUDFImpl for Sum {
         &self.signature
     }
 
+    fn diagnose_failed_signature(
+        &self,
+        arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        match arg_types {
+            [DataType::Boolean] => {
+                Some(plan_datafusion_err!("Sum not supported for Boolean"))
+            }
+            _ => None,
+        }
+    }
+
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         match &arg_types[0] {
             DataType::Int64 => Ok(DataType::Int64),

datafusion/functions-aggregate/src/sum.rs

@@ -7896,17 +7896,17 @@ select generate_series(arrow_cast('2021-01-01T00:00:00', 'Timestamp(Nanosecond,
 [2021-01-01T00:00:00-05:00, 2021-01-01T01:29:54.500-05:00, 2021-01-01T02:59:49-05:00, 2021-01-01T04:29:43.500-05:00, 2021-01-01T05:59:38-05:00]
 
 ## mixing types for timestamps is not supported
-query error DataFusion error: Error during planning: Internal error: Function 'generate_series' failed to match any signature
+query error DataFusion error: Error during planning: Function 'generate_series' failed to match any signature
 select generate_series(arrow_cast('2021-01-01T00:00:00', 'Timestamp(Nanosecond, Some("-05:00"))'), DATE '2021-01-02', INTERVAL '1' HOUR);
 
 ## mixing types not allowed even if an argument is null
-query error DataFusion error: Error during planning: Internal error: Function 'generate_series' failed to match any signature
+query error DataFusion error: Error during planning: Function 'generate_series' failed to match any signature
 select generate_series(TIMESTAMP '1992-09-01', DATE '1993-03-01', NULL);
 
-query error DataFusion error: Error during planning: Internal error: Function 'generate_series' failed to match any signature
+query error DataFusion error: Error during planning: Function 'generate_series' failed to match any signature
 select generate_series(1, '2024-01-01', '2025-01-02');
 
-query error DataFusion error: Error during planning: Internal error: Function 'generate_series' failed to match any signature
+query error DataFusion error: Error during planning: Function 'generate_series' failed to match any signature
 select generate_series('2024-01-01'::timestamp, '2025-01-02', interval '1 day');
 
 ## should return NULL

datafusion/sqllogictest/test_files/array.slt

@@ -125,13 +125,19 @@ from aggregate_test_100
 order by c9
 
 # WindowFunction wrong signature
-statement error DataFusion error: Error during planning: Internal error: Function 'nth_value' failed to match any signature
+statement error DataFusion error: Error during planning: Function 'nth_value' expects 0 to 2 arguments but received 3
 select
 c9,
 nth_value(c5, 2, 3) over (order by c9) as nv1
 from aggregate_test_100
 order by c9
 
+query error DataFusion error: Error during planning: Sum not supported for Boolean
+select sum(bool_col) from (values (true), (false), (null)) as t(bool_col);
+
+query error DataFusion error: Error during planning: Avg not supported for Boolean
+select avg(bool_col) from (values (true), (false), (null)) as t(bool_col);
+
 
 # nth_value with wrong name
 statement error DataFusion error: Error during planning: Invalid function 'nth_vlue'.\nDid you mean 'nth_value'?

datafusion/sqllogictest/test_files/errors.slt

@@ -208,10 +208,10 @@ SELECT substr('alphabet', NULL, 2)
 ----
 NULL
 
-statement error Function 'substr' failed to match any signature
+statement error DataFusion error: Error during planning: Function 'substr' requires String, but received Int64
 SELECT substr(1, 3)
 
-statement error Function 'substr' failed to match any signature
+statement error DataFusion error: Error during planning: Function 'substr' requires String, but received Int64
 SELECT substr(1, 3, 4)
 
 query T

datafusion/sqllogictest/test_files/functions.slt

@@ -86,7 +86,7 @@ SELECT substr("STR" => 'hello world', "start_pos" => 7);
 
 # Error: wrong number of arguments
 # This query provides only 1 argument but substr requires 2 or 3
-query error Function 'substr' failed to match any signature
+query error DataFusion error: Error during planning: Function 'substr' expects 2 to 3 arguments but received 1
 SELECT substr(str => 'hello world');
 
 #############