chore(cubestore): Upgrade DF 46: Let DATE_ADD and DATE_SUB tolerate time zones

Author: srh

rust/cubestore/cubestore-sql-tests/src/tests.rs

@@ -6735,6 +6735,21 @@ async fn date_add(service: Box<dyn SqlClient>) {
             None,
         ]),
     );
+
+    // Check we tolerate NOW(), perhaps with +00:00 time zone.
+    let r = service
+        .exec_query("SELECT NOW(), date_add(NOW(), INTERVAL '1 day')")
+        .await
+        .unwrap();
+    let rows = to_rows(&r);
+    assert_eq!(1, rows.len());
+    assert_eq!(2, rows[0].len());
+    match (&rows[0][0], &rows[0][1]) {
+        (TableValue::Timestamp(tv), TableValue::Timestamp(day_later)) => {
+            assert_eq!(day_later.get_time_stamp(), tv.get_time_stamp() + 86400i64 * 1_000_000_000);
+        },
+        _ => panic!("row has wrong types: {:?}", rows[0]),
+    }
 }
 
 async fn date_bin(service: Box<dyn SqlClient>) {

rust/cubestore/cubestore/src/queryplanner/udfs.rs

@@ -6,12 +6,13 @@ use datafusion::arrow::array::{
 };
 use datafusion::arrow::buffer::ScalarBuffer;
 use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
+use datafusion::common::internal_err;
 use datafusion::error::DataFusionError;
 use datafusion::logical_expr::function::AccumulatorArgs;
 use datafusion::logical_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 use datafusion::logical_expr::{
     AggregateUDF, AggregateUDFImpl, Expr, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
-    Volatility,
+    Volatility, TIMEZONE_WILDCARD,
 };
 use datafusion::physical_plan::{Accumulator, ColumnarValue};
 use datafusion::scalar::ScalarValue;
@@ -457,6 +458,7 @@ struct DateAddSub {
 
 impl DateAddSub {
     pub fn new(is_add: bool) -> DateAddSub {
+        let tz_wildcard: Arc<str> = Arc::from(TIMEZONE_WILDCARD);
         DateAddSub {
             is_add,
             signature: Signature {
@@ -473,6 +475,22 @@ impl DateAddSub {
                         DataType::Timestamp(TimeUnit::Nanosecond, None),
                         DataType::Interval(IntervalUnit::MonthDayNano),
                     ]),
+                    // We wanted this for NOW(), which has "+00:00" time zone.  Using
+                    // TIMEZONE_WILDCARD to favor DST-related questions over "UTC" == "+00:00"
+                    // questions.  MySQL doesn't have a timezone as this function is applied, and we
+                    // simply invoke DF's date + interval behavior.
+                    TypeSignature::Exact(vec![
+                        DataType::Timestamp(TimeUnit::Nanosecond, Some(tz_wildcard.clone())),
+                        DataType::Interval(IntervalUnit::YearMonth),
+                    ]),
+                    TypeSignature::Exact(vec![
+                        DataType::Timestamp(TimeUnit::Nanosecond, Some(tz_wildcard.clone())),
+                        DataType::Interval(IntervalUnit::DayTime),
+                    ]),
+                    TypeSignature::Exact(vec![
+                        DataType::Timestamp(TimeUnit::Nanosecond, Some(tz_wildcard)),
+                        DataType::Interval(IntervalUnit::MonthDayNano),
+                    ]),
                 ]),
                 volatility: Volatility::Immutable,
             },
@@ -505,18 +523,22 @@ impl ScalarUDFImpl for DateAddSub {
     fn signature(&self) -> &Signature {
         &self.signature
     }
-    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType, DataFusionError> {
-        Ok(DataType::Timestamp(TimeUnit::Nanosecond, None))
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType, DataFusionError> {
+        if arg_types.len() != 2 {
+            return Err(DataFusionError::Internal(format!("DateAddSub return_type expects 2 arguments, got {:?}", arg_types)));
+        }
+        match (&arg_types[0], &arg_types[1]) {
+            (ts@DataType::Timestamp(_, _), DataType::Interval(_)) => Ok(ts.clone()),
+            _ => Err(DataFusionError::Internal(format!("DateAddSub return_type expects Timestamp and Interval arguments, got {:?}", arg_types))),
+        }
     }
     fn invoke(&self, inputs: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
         use datafusion::arrow::compute::kernels::numeric::add;
         use datafusion::arrow::compute::kernels::numeric::sub;
         assert_eq!(inputs.len(), 2);
         // DF 42.2.0 already has date + interval or date - interval.  Note that `add` and `sub` are
         // public (defined in arrow_arith), while timestamp-specific functions they invoke,
-        // `arithmetic_op` and then `timestamp_op::<TimestampNanosecondType>`, are not.
-        //
-        // TODO upgrade DF: Double-check that the TypeSignature is actually enforced.
+        // Arrow's `arithmetic_op` and then `timestamp_op::<TimestampNanosecondType>`, are not.
         datafusion::physical_expr_common::datum::apply(
             &inputs[0],
             &inputs[1],