Rolling window queries

datafusion/Cargo.toml

@@ -48,7 +48,7 @@ ahash = "0.7"
 hashbrown = "0.11"
 arrow = { git = "https://github.com/cube-js/arrow-rs.git", branch = "cube", features = ["prettyprint"] }
 parquet = { git = "https://github.com/cube-js/arrow-rs.git", branch = "cube", features = ["arrow"] }
-sqlparser = "0.9.0"
+sqlparser = { git = "https://github.com/cube-js/sqlparser-rs.git", rev = "22544370d1f1b608b9a5e75477b1979ed1f7eb72" }
 paste = "^1.0"
 num_cpus = "1.13.0"
 chrono = "0.4"

datafusion/src/cube_ext/datetime.rs

@@ -0,0 +1,155 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::DataFusionError;
+use crate::scalar::ScalarValue;
+use arrow::array::{Array, TimestampNanosecondArray, TimestampNanosecondBuilder};
+use chrono::{DateTime, Datelike, Duration, NaiveDate, TimeZone, Utc};
+
+pub fn date_addsub_array(
+    t: &TimestampNanosecondArray,
+    i: ScalarValue,
+    is_add: bool,
+) -> Result<TimestampNanosecondArray, DataFusionError> {
+    let mut result = TimestampNanosecondBuilder::new(t.len());
+    match i {
+        ScalarValue::IntervalYearMonth(Some(v)) => {
+            for i in 0..t.len() {
+                if t.is_null(i) {
+                    result.append_null()?;
+                } else {
+                    let t = Utc.timestamp_nanos(t.value(i));
+                    result.append_value(
+                        date_addsub_year_month(t, v, is_add)?.timestamp_nanos(),
+                    )?;
+                }
+            }
+        }
+        ScalarValue::IntervalDayTime(Some(v)) => {
+            for i in 0..t.len() {
+                if t.is_null(i) {
+                    result.append_null()?;
+                } else {
+                    let t = Utc.timestamp_nanos(t.value(i));
+                    result.append_value(
+                        date_addsub_day_time(t, v, is_add)?.timestamp_nanos(),
+                    )?;
+                }
+            }
+        }
+        _ => {
+            let name = match is_add {
+                true => "DATE_ADD",
+                false => "DATE_SUB",
+            };
+            return Err(DataFusionError::Plan(format!(
+                "Second argument of `{}` must be a non-null interval",
+                name
+            )));
+        }
+    }
+
+    Ok(result.finish())
+}
+
+pub fn date_addsub_scalar(
+    t: DateTime<Utc>,
+    i: ScalarValue,
+    is_add: bool,
+) -> Result<DateTime<Utc>, DataFusionError> {
+    match i {
+        ScalarValue::IntervalYearMonth(Some(v)) => date_addsub_year_month(t, v, is_add),
+        ScalarValue::IntervalDayTime(Some(v)) => date_addsub_day_time(t, v, is_add),
+        _ => {
+            let name = match is_add {
+                true => "DATE_ADD",
+                false => "DATE_SUB",
+            };
+            return Err(DataFusionError::Plan(format!(
+                "Second argument of `{}` must be a non-null interval",
+                name
+            )));
+        }
+    }
+}
+
+fn date_addsub_year_month(
+    t: DateTime<Utc>,
+    i: i32,
+    is_add: bool,
+) -> Result<DateTime<Utc>, DataFusionError> {
+    let i = match is_add {
+        true => i,
+        false => -i,
+    };
+
+    let mut year = t.year();
+    // Note month is numbered 0..11 in this function.
+    let mut month = t.month() as i32 - 1;
+
+    year += i / 12;
+    month += i % 12;
+
+    if month < 0 {
+        year -= 1;
+        month += 12;
+    }
+    debug_assert!(0 <= month);
+    year += month / 12;
+    month = month % 12;
+
+    match change_ym(t, year, 1 + month as u32) {
+        Some(t) => return Ok(t),
+        None => {
+            return Err(DataFusionError::Execution(format!(
+                "Failed to set date to ({}-{})",
+                year,
+                1 + month
+            )))
+        }
+    };
+}
+
+fn date_addsub_day_time(
+    t: DateTime<Utc>,
+    interval: i64,
+    is_add: bool,
+) -> Result<DateTime<Utc>, DataFusionError> {
+    let i = match is_add {
+        true => interval,
+        false => -interval,
+    };
+
+    let days: i64 = i.signum() * (i.abs() >> 32);
+    let millis: i64 = i.signum() * ((i.abs() << 32) >> 32);
+    return Ok(t + Duration::days(days) + Duration::milliseconds(millis));
+}
+
+fn change_ym(t: DateTime<Utc>, y: i32, m: u32) -> Option<DateTime<Utc>> {
+    debug_assert!(1 <= m && m <= 12);
+    let mut d = t.day();
+    d = d.min(last_day_of_month(y, m));
+    t.with_day(1)?.with_year(y)?.with_month(m)?.with_day(d)
+}
+
+fn last_day_of_month(y: i32, m: u32) -> u32 {
+    debug_assert!(1 <= m && m <= 12);
+    if m == 12 {
+        return 31;
+    }
+    NaiveDate::from_ymd(y, m + 1, 1).pred().day()
+}

datafusion/src/cube_ext/mod.rs

@@ -16,8 +16,10 @@
 // under the License.
 
 pub mod alias;
+pub mod datetime;
 pub mod join;
 pub mod joinagg;
+pub mod rolling;
 pub mod sequence;
 pub mod stream;
 pub mod util;