Jefffrey commented on code in PR #18317:
URL: https://github.com/apache/datafusion/pull/18317#discussion_r2468267625


##########
datafusion/functions-nested/src/range.rs:
##########
@@ -378,44 +420,45 @@ impl Range {
 
     fn gen_range_timestamp(&self, args: &[ArrayRef]) -> Result<ArrayRef> {
         let [start, stop, step] = take_function_args(self.name(), args)?;
+        let step = as_interval_mdn_array(step)?;
+
+        // Signature can only guarantee we get a timestamp type, not 
specifically
+        // timestamp(ns) so handle potential cast from other timestamps here.
+        fn cast_to_ns(arr: &ArrayRef) -> Result<ArrayRef> {
+            match arr.data_type() {
+                DataType::Timestamp(TimeUnit::Nanosecond, _) => 
Ok(Arc::clone(arr)),
+                DataType::Timestamp(_, tz) => Ok(cast(
+                    arr,
+                    &DataType::Timestamp(TimeUnit::Nanosecond, tz.to_owned()),
+                )?),
+                _ => unreachable!(),
+            }
+        }
+        let start = cast_to_ns(start)?;
+        let start = as_timestamp_nanosecond_array(&start)?;
+        let stop = cast_to_ns(stop)?;
+        let stop = as_timestamp_nanosecond_array(&stop)?;

Review Comment:
   I wonder if is better to instead try uplift the signature coercible API to 
allow specifying we only want timestamp(ns) (of any timezone) and let that 
handle coercion for us 🤔 



##########
datafusion/sqllogictest/test_files/array.slt:
##########
@@ -7131,9 +7142,18 @@ 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: Internal error: Unexpected argument type for 
generate_series : Date32
+query error DataFusion error: Error during planning: Internal error: 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
+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
+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
+select generate_series('2024-01-01'::timestamp, '2025-01-02', interval '1 
day');

Review Comment:
   So these are the original examples from the issue; **technically** we still 
return an internal error, so it feels a bit wrong to claim we close #15881 
(well its an internal error nested in a plan error apparently)
   
   However this internal error comes from the function signature related code, 
so we can fix it there and it should propagate to here (aka its not a specific 
issue with generate_series)



##########
datafusion/functions-nested/src/range.rs:
##########
@@ -320,29 +363,28 @@ impl Range {
 
     fn gen_range_date(&self, args: &[ArrayRef]) -> Result<ArrayRef> {
         let [start, stop, step] = take_function_args(self.name(), args)?;
+        let step = as_interval_mdn_array(step)?;
 
-        let (start_array, stop_array, step_array) = (
-            as_date32_array(start)?,
-            as_date32_array(stop)?,
-            as_interval_mdn_array(step)?,
-        );
+        // Signature can only guarantee we get a date type, not specifically
+        // date32 so handle potential cast from date64 here.
+        let start = cast(start, &DataType::Date32)?;
+        let start = as_date32_array(&start)?;
+        let stop = cast(stop, &DataType::Date32)?;
+        let stop = as_date32_array(&stop)?;

Review Comment:
   I don't see much benefit in handling date64 natively anyway (given it's a 
bit of a weird type) so I think this is fine



##########
datafusion/expr-common/src/signature.rs:
##########
@@ -360,6 +357,7 @@ impl TypeSignatureClass {
             TypeSignatureClass::Binary if native_type.is_binary() => {
                 Ok(origin_type.to_owned())
             }
+            _ if native_type.is_null() => Ok(origin_type.to_owned()),

Review Comment:
   We were missing this even though we check for it in `matches_native_type`; I 
think this will allow passing null types to coercible signatures without 
needing to explicitly specify null type handling



##########
datafusion/sqllogictest/test_files/array.slt:
##########
@@ -6938,13 +6938,13 @@ select range(5),
        range(2, 5),
        range(2, 10, 3),
        range(10, 2, -3),
-       range(1, 5, -1),
-       range(1, -5, 1),
-       range(1, -5, -1),
-       range(DATE '1992-09-01', DATE '1993-03-01', INTERVAL '1' MONTH),
-       range(DATE '1993-02-01', DATE '1993-01-01', INTERVAL '-1' DAY),
-       range(DATE '1989-04-01', DATE '1993-03-01', INTERVAL '1' YEAR),
-       range(DATE '1993-03-01', DATE '1989-04-01', INTERVAL '1' YEAR)
+       range(arrow_cast(1, 'Int8'), 5, -1),
+       range(arrow_cast(1, 'Int16'), arrow_cast(-5, 'Int8'), 1),
+       range(arrow_cast(1, 'Int32'), arrow_cast(-5, 'Int16'), arrow_cast(-1, 
'Int8')),
+       range(DATE '1992-09-01', DATE '1993-03-01', arrow_cast('1 MONTH', 
'Interval(YearMonth)')),
+       range(DATE '1993-02-01', arrow_cast(DATE '1993-01-01', 'Date64'), 
INTERVAL '-1' DAY),
+       range(arrow_cast(DATE '1989-04-01', 'Date64'), DATE '1993-03-01', 
INTERVAL '1' YEAR),
+       range(arrow_cast(DATE '1993-03-01', 'Date64'), arrow_cast(DATE 
'1989-04-01', 'Date64'), INTERVAL '1' YEAR)

Review Comment:
   Uplift tests to ensure coercion is handled properly for different types



##########
datafusion/functions-nested/src/range.rs:
##########
@@ -146,18 +152,60 @@ impl Default for Range {
 }
 
 impl Range {
+    fn defined_signature() -> Signature {

Review Comment:
   This is the main change



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to