chore: better messaging

apache · Sep 9, 2024 · 01fab57 · 01fab57
1 parent 8322811
commit 01fab57
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 30 deletions.
diff --git a/datafusion/functions/src/unicode/substr.rs b/datafusion/functions/src/unicode/substr.rs
@@ -81,35 +81,42 @@ impl ScalarUDFImpl for SubstrFunc {
     }
 
     fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
-        match arg_types {
-            [first, second] => match (first, second) {
-                (
-                    DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8,
-                    DataType::Int64 | DataType::Int32,
-                ) => Ok(vec![first.to_owned(), DataType::Int64]),
-                _ => Err(DataFusionError::Execution(format!(
-                    "The {} function can only accept strings, but got {:?}.",
-                    self.name(),
-                    arg_types
-                ))),
-            },
-            [first, second, third] => match (first, second, third) {
-                (
-                    DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8,
-                    DataType::Int64 | DataType::Int32,
-                    DataType::Int64 | DataType::Int32,
-                ) => Ok(vec![first.to_owned(), DataType::Int64, DataType::Int64]),
-                _ => Err(DataFusionError::Execution(format!(
-                    "The {} function can only accept strings, but got {:?}.",
-                    self.name(),
-                    arg_types
-                ))),
-            },
-            _ => Err(DataFusionError::Execution(format!(
-                "The {} function can only accept strings, but got {:?}",
+        if ![DataType::LargeUtf8, DataType::Utf8View, DataType::Utf8]
+            .contains(&arg_types[0])
+        {
+            return Err(DataFusionError::Execution(format!(
+                "The first argument of the {} function can only be a string, but got {:?}.",
                 self.name(),
-                arg_types
-            ))),
+                arg_types[0]
+            )));
+        }
+
+        if ![DataType::Int64, DataType::Int32].contains(&arg_types[1]) {
+            return Err(DataFusionError::Execution(format!(
+                "The second argument of the {} function can only be an integer, but got {:?}.",
+                self.name(),
+                arg_types[1]
+            )));
+        }
+
+        if arg_types.len() == 3
+            && ![DataType::Int64, DataType::Int32].contains(&arg_types[2])
+        {
+            return Err(DataFusionError::Execution(format!(
+                "The third argument of the {} function can only be an integer, but got {:?}.",
+                self.name(),
+                arg_types[2]
+            )));
+        }
+
+        if arg_types.len() == 2 {
+            Ok(vec![arg_types[0].to_owned(), DataType::Int64])
+        } else {
+            Ok(vec![
+                arg_types[0].to_owned(),
+                DataType::Int64,
+                DataType::Int64,
+            ])
         }
     }
 }

diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt
@@ -497,10 +497,10 @@ SELECT substr('alphabet', 3, CAST(NULL AS int))
 ----
 NULL
 
-statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Execution\("The substr function can only accept strings, but got \[Int64, Int64]\."\)
+statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Execution\("The first argument of the substr function can only be a string, but got Int64\."\)
 SELECT substr(1, 3)
 
-statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Execution\("The substr function can only accept strings, but got \[Int64, Int64, Int64\]\."\)
+statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Execution\("The first argument of the substr function can only be a string, but got Int64\."\)
 SELECT substr(1, 3, 4)
 
 query T