improve OneOf signature diagnostics

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;
@@ -272,6 +274,16 @@ impl AggregateUDF {
         self.inner.coerce_types(arg_types)
     }
 
+    /// Returns a function-specific signature failure for argument combinations
+    /// that are better described semantically than by generic signature
+    /// diagnostics.
+    pub fn diagnose_failed_signature(
+        &self,
+        arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        self.inner.diagnose_failed_signature(arg_types)
+    }
+
     /// See [`AggregateUDFImpl::with_beneficial_ordering`] for more details.
     pub fn with_beneficial_ordering(
         self,
@@ -799,6 +811,16 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         not_impl_err!("Function {} does not implement coerce_types", self.name())
     }
 
+    /// Returns a semantic signature failure for argument combinations that
+    /// should be surfaced directly instead of going through generic `OneOf`
+    /// diagnostics.
+    fn diagnose_failed_signature(
+        &self,
+        _arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        None
+    }
+
     /// If this function is max, return true
     /// If the function is min, return false
     /// Otherwise return None (the default)
@@ -1339,6 +1361,13 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl {
         self.inner.coerce_types(arg_types)
     }
 
+    fn diagnose_failed_signature(
+        &self,
+        arg_types: &[DataType],
+    ) -> Option<DataFusionError> {
+        self.inner.diagnose_failed_signature(arg_types)
+    }
+
     fn return_field(&self, arg_fields: &[FieldRef]) -> Result<FieldRef> {
         self.inner.return_field(arg_fields)
     }

datafusion/expr/src/udaf.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::{
@@ -134,6 +136,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};
@@ -208,6 +211,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

@@ -28,7 +28,7 @@ use datafusion_doc::window_doc_sections::DOC_SECTION_ANALYTICAL;
 use datafusion_expr::window_state::WindowAggState;
 use datafusion_expr::{
     Documentation, LimitEffect, Literal, PartitionEvaluator, ReversedUDWF, Signature,
-    TypeSignature, Volatility, WindowUDFImpl,
+    Volatility, WindowUDFImpl,
 };
 use datafusion_functions_window_common::field;
 use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
@@ -96,17 +96,14 @@ pub struct NthValue {
 impl NthValue {
     /// Create a new `nth_value` function
     pub fn new(kind: NthValueKind) -> Self {
-        Self {
-            signature: Signature::one_of(
-                vec![
-                    TypeSignature::Nullary,
-                    TypeSignature::Any(1),
-                    TypeSignature::Any(2),
-                ],
-                Volatility::Immutable,
-            ),
-            kind,
-        }
+        let signature = match kind {
+            NthValueKind::First | NthValueKind::Last => {
+                Signature::any(1, Volatility::Immutable)
+            }
+            NthValueKind::Nth => Signature::any(2, Volatility::Immutable),
+        };
+
+        Self { signature, kind }
     }
 
     pub fn first() -> Self {

datafusion/functions-window/src/nth_value.rs

@@ -8399,17 +8399,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,39 @@ 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: The function 'nth_value' expected 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
 
+statement error DataFusion error: Error during planning: The function 'nth_value' expected 2 arguments but received 1
+select
+c9,
+nth_value(c5) over (order by c9) as nv2
+from aggregate_test_100
+order by c9
+
+statement error DataFusion error: Error during planning: The function 'first_value' expected 1 arguments but received 2
+select
+c9,
+first_value(c5, 2) over (order by c9) as fv1
+from aggregate_test_100
+order by c9
+
+statement error DataFusion error: Error during planning: 'nth_value' does not support zero arguments
+select nth_value() over (order by x) from (values (1), (2)) as t(x)
+
+statement error DataFusion error: Error during planning: 'first_value' does not support zero arguments
+select first_value() over (order by x) from (values (1), (2)) as t(x)
+
+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');
 
 #############

datafusion/sqllogictest/test_files/named_arguments.slt

@@ -56,7 +56,7 @@ SELECT hex(column1) FROM VALUES (arrow_cast('hello', 'LargeBinary')), (NULL), (a
 NULL
 776F726C64
 
-statement error Function 'hex' expects 1 arguments but received 2
+statement error DataFusion error: Error during planning: Function 'hex' expects 1 arguments but received 2
 SELECT hex(1, 2);
 
 query T