chore: address comments

apache · Sep 10, 2024 · c73baf2 · c73baf2
1 parent db7cc31
commit c73baf2
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 15 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,24 @@ 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())]),
+            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 {:?}", arg_types[0])
         }
     }
 }

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\("The to_local_time function can only accept timestamp as the arg, got Utf8"\)
 select to_local_time('2024-04-01T00:00:20Z');
 
 # invalid timezone