chore: address comments

apache · Sep 12, 2024 · 5477cf1 · 5477cf1
1 parent fa1627e
commit 5477cf1
Show file tree

Hide file tree

Showing 9 changed files with 50 additions and 41 deletions.
diff --git a/datafusion/functions/src/datetime/to_local_time.rs b/datafusion/functions/src/datetime/to_local_time.rs
@@ -30,7 +30,7 @@ use arrow::datatypes::{
 
 use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
 use datafusion_common::cast::as_primitive_array;
-use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_common::{exec_err, plan_err, DataFusionError, Result, ScalarValue};
 use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 
 /// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
@@ -330,20 +330,25 @@ impl ScalarUDFImpl for ToLocalTimeFunc {
 
     fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
         if arg_types.len() != 1 {
-            return Err(DataFusionError::Execution(format!(
+            return plan_err!(
                 "to_local_time function requires 1 argument, got {:?}",
                 arg_types.len()
-            )));
+            );
         }
 
-        match &arg_types[0] {
-            Timestamp(Nanosecond, timezone) => Ok(vec![Timestamp(Nanosecond, timezone.clone())]),
-            Timestamp(Microsecond, timezone) => Ok(vec![Timestamp(Microsecond, timezone.clone())]),
-            Timestamp(Millisecond, timezone) => Ok(vec![Timestamp(Millisecond, timezone.clone())]),
+        let first_arg = arg_types[0].clone();
+        match &first_arg {
+            Timestamp(Nanosecond, timezone) => {
+                Ok(vec![Timestamp(Nanosecond, timezone.clone())])
+            }
+            Timestamp(Microsecond, timezone) => {
+                Ok(vec![Timestamp(Microsecond, timezone.clone())])
+            }
+            Timestamp(Millisecond, timezone) => {
+                Ok(vec![Timestamp(Millisecond, timezone.clone())])
+            }
             Timestamp(Second, timezone) => Ok(vec![Timestamp(Second, timezone.clone())]),
-            _ => Err(DataFusionError::Execution(format!(
-                "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
-            )))
+            _ => plan_err!("The to_local_time function can only accept Timestamp as the arg got {first_arg}"),
         }
     }
 }

diff --git a/datafusion/functions/src/encoding/inner.rs b/datafusion/functions/src/encoding/inner.rs
@@ -76,11 +76,11 @@ impl ScalarUDFImpl for EncodeFunc {
 
     fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
         if arg_types.len() != 2 {
-            return Err(DataFusionError::Plan(format!(
+            return plan_err!(
                 "{} expects to get 2 arguments, but got {}",
                 self.name(),
                 arg_types.len()
-            )));
+            );
         }
 
         if arg_types[1] != DataType::Utf8 {
@@ -94,10 +94,10 @@ impl ScalarUDFImpl for EncodeFunc {
             DataType::LargeUtf8 | DataType::LargeBinary => {
                 Ok(vec![DataType::LargeUtf8, DataType::Utf8])
             }
-            _ => Err(DataFusionError::Plan(format!(
+            _ => plan_err!(
                 "1st argument should be Utf8 or Binary or Null, got {:?}",
                 arg_types[0]
-            ))),
+            ),
         }
     }
 }
@@ -143,15 +143,15 @@ impl ScalarUDFImpl for DecodeFunc {
 
     fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
         if arg_types.len() != 2 {
-            return Err(DataFusionError::Plan(format!(
+            return plan_err!(
                 "{} expects to get 2 arguments, but got {}",
                 self.name(),
                 arg_types.len()
-            )));
+            );
         }
 
         if arg_types[1] != DataType::Utf8 {
-            return Err(DataFusionError::Plan("2nd argument should be Utf8".into()));
+            return plan_err!("2nd argument should be Utf8");
         }
 
         match arg_types[0] {
@@ -161,10 +161,10 @@ impl ScalarUDFImpl for DecodeFunc {
             DataType::LargeUtf8 | DataType::LargeBinary => {
                 Ok(vec![DataType::LargeBinary, DataType::Utf8])
             }
-            _ => Err(DataFusionError::Plan(format!(
+            _ => plan_err!(
                 "1st argument should be Utf8 or Binary or Null, got {:?}",
                 arg_types[0]
-            ))),
+            ),
         }
     }
 }

diff --git a/datafusion/functions/src/unicode/strpos.rs b/datafusion/functions/src/unicode/strpos.rs
@@ -23,7 +23,7 @@ use arrow::array::{
 };
 use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type};
 
-use datafusion_common::{exec_err, DataFusionError, Result};
+use datafusion_common::{exec_err, plan_err, Result};
 use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 
 use crate::utils::{make_scalar_function, utf8_to_int_type};
@@ -84,12 +84,12 @@ impl ScalarUDFImpl for StrposFunc {
                     (_, DataType::Null) => Ok(vec![first.to_owned(), DataType::Utf8]),
                     (DataType::Dictionary(_, value_type), DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8)  => match **value_type {
                         DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8 | DataType::Null | DataType::Binary => Ok(vec![*value_type.clone(), second.to_owned()]),
-                        _ => Err(DataFusionError::Execution(format!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}.", **value_type))),
+                        _ => plan_err!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}.", **value_type),
                     },
-                    _ => Err(DataFusionError::Execution(format!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}.", arg_types)))
+                    _ => plan_err!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}.", arg_types)
                 }
             },
-            _ => Err(DataFusionError::Execution(format!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}", arg_types)))
+            _ => plan_err!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}", arg_types)
         }
     }
 }

diff --git a/datafusion/functions/src/unicode/substr.rs b/datafusion/functions/src/unicode/substr.rs
@@ -27,7 +27,7 @@ use arrow::array::{
 use arrow::datatypes::DataType;
 use arrow_buffer::{NullBufferBuilder, ScalarBuffer};
 use datafusion_common::cast::as_int64_array;
-use datafusion_common::{exec_datafusion_err, exec_err, DataFusionError, Result};
+use datafusion_common::{exec_datafusion_err, exec_err, plan_err, Result};
 use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 
 #[derive(Debug)]
@@ -84,29 +84,29 @@ impl ScalarUDFImpl for SubstrFunc {
         if ![DataType::LargeUtf8, DataType::Utf8View, DataType::Utf8]
             .contains(&arg_types[0])
         {
-            return Err(DataFusionError::Execution(format!(
+            return plan_err!(
                 "The first argument of the {} function can only be a string, but got {:?}.",
                 self.name(),
                 arg_types[0]
-            )));
+            );
         }
 
         if ![DataType::Int64, DataType::Int32].contains(&arg_types[1]) {
-            return Err(DataFusionError::Execution(format!(
+            return plan_err!(
                 "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!(
+            return plan_err!(
                 "The third argument of the {} function can only be an integer, but got {:?}.",
                 self.name(),
                 arg_types[2]
-            )));
+            );
         }
 
         if arg_types.len() == 2 {

diff --git a/datafusion/sqllogictest/test_files/encoding.slt b/datafusion/sqllogictest/test_files/encoding.slt
@@ -28,13 +28,13 @@ CREATE TABLE test(
 ;
 
 # errors
-query error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan\("1st argument should be Utf8 or Binary or Null, got Int64"\)
+query error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan
 select encode(12, 'hex');
 
 query error DataFusion error: Error during planning: There is no built\-in encoding named 'non_encoding', currently supported encodings are: base64, hex
 select encode(bin_field, 'non_encoding') from test;
 
-query error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan\("1st argument should be Utf8 or Binary or Null, got Int64"\)
+query error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan
 select decode(12, 'hex');
 
 query error DataFusion error: Error during planning: There is no built\-in encoding named 'non_encoding', currently supported encodings are: base64, hex

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 first argument of the substr function can only be a string, but got Int64\."\)
+statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan
 SELECT substr(1, 3)
 
-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\."\)
+statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan
 SELECT substr(1, 3, 4)
 
 query T

diff --git a/datafusion/sqllogictest/test_files/scalar.slt b/datafusion/sqllogictest/test_files/scalar.slt
@@ -1907,7 +1907,7 @@ select position('' in '')
 1
 
 
-query error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Execution\("The STRPOS/INSTR/POSITION function can only accept strings, but got \[Int64, Int64\]\."\)
+query error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan
 select position(1 in 1)
 
 

diff --git a/datafusion/sqllogictest/test_files/select.slt b/datafusion/sqllogictest/test_files/select.slt
@@ -483,15 +483,19 @@ CREATE TABLE foo2(c1 double, c2 double) AS VALUES
 (2, 5),
 (3, 6);
 
-statement ok
-SELECT COALESCE(column1, column2) FROM VALUES (null, 1.2);
+query T
+SELECT arrow_typeof(COALESCE(column1, column2)) FROM VALUES (null, 1.2);
+----
+Float64
 
 # multiple rows and columns with null need type coercion
-statement ok
-SELECT column1, column2, column3 FROM VALUES
+query TTT
+select arrow_typeof(column1), arrow_typeof(column2), arrow_typeof(column3) from (SELECT column1, column2, column3 FROM VALUES
 (null, 2, 'a'),
 (1.1, null, 'b'),
-(2, 5, null);
+(2, 5, null)) LIMIT 1;
+----
+Float64 Int64 Utf8
 
 
 # foo distinct

diff --git a/datafusion/sqllogictest/test_files/timestamps.slt b/datafusion/sqllogictest/test_files/timestamps.slt
@@ -2858,7 +2858,7 @@ statement error
 select to_local_time('2024-04-01T00:00:20Z'::timestamp, 'some string');
 
 # invalid argument data type
-statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Execution\("The to_local_time function can only accept timestamp as the arg, got Utf8"\)
+statement error DataFusion error: Error during planning: Execution error: User\-defined coercion failed with Plan
 select to_local_time('2024-04-01T00:00:20Z');
 
 # invalid timezone