feat(cubestore): implement date_bin()

Author: KSDaemon

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -409,6 +409,7 @@ impl ContextProvider for MetaStoreSchemaProvider {
             "unix_timestamp" | "UNIX_TIMESTAMP" => CubeScalarUDFKind::UnixTimestamp,
             "date_add" | "DATE_ADD" => CubeScalarUDFKind::DateAdd,
             "date_sub" | "DATE_SUB" => CubeScalarUDFKind::DateSub,
+            "date_bin" | "DATE_BIN" => CubeScalarUDFKind::DateBin,
             _ => return None,
         };
         return Some(Arc::new(scalar_udf_by_kind(kind).descriptor()));

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

@@ -2,7 +2,9 @@ use crate::queryplanner::coalesce::{coalesce, SUPPORTED_COALESCE_TYPES};
 use crate::queryplanner::hll::{Hll, HllUnion};
 use crate::CubeError;
 use chrono::{TimeZone, Utc};
-use datafusion::arrow::array::{Array, BinaryArray, TimestampNanosecondArray, UInt64Builder};
+use datafusion::arrow::array::{
+    Array, ArrayRef, BinaryArray, TimestampNanosecondArray, UInt64Builder,
+};
 use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
 use datafusion::cube_ext::datetime::{date_addsub_array, date_addsub_scalar};
 use datafusion::error::DataFusionError;
@@ -16,6 +18,12 @@ use smallvec::smallvec;
 use smallvec::SmallVec;
 use std::sync::Arc;
 
+const NANOS_IN_DAY: i64 = 86_400_000_000_000;
+const DAYS_IN_MONTH: i64 = 30; // Approximation for months
+const DAYS_IN_YEAR: i64 = 365; // Approximation for years
+const NANOS_IN_MONTH: i64 = DAYS_IN_MONTH * NANOS_IN_DAY;
+const NANOS_IN_YEAR: i64 = DAYS_IN_YEAR * NANOS_IN_DAY;
+
 #[derive(Copy, Clone, Debug, Serialize, Deserialize)]
 pub enum CubeScalarUDFKind {
     HllCardinality, // cardinality(), accepting the HyperLogLog sketches.
@@ -24,6 +32,7 @@ pub enum CubeScalarUDFKind {
     UnixTimestamp,
     DateAdd,
     DateSub,
+    DateBin,
 }
 
 pub trait CubeScalarUDF {
@@ -40,6 +49,7 @@ pub fn scalar_udf_by_kind(k: CubeScalarUDFKind) -> Box<dyn CubeScalarUDF> {
         CubeScalarUDFKind::UnixTimestamp => Box::new(UnixTimestamp {}),
         CubeScalarUDFKind::DateAdd => Box::new(DateAddSub { is_add: true }),
         CubeScalarUDFKind::DateSub => Box::new(DateAddSub { is_add: false }),
+        CubeScalarUDFKind::DateBin => Box::new(DateBin {}),
     }
 }
 
@@ -63,6 +73,9 @@ pub fn scalar_kind_by_name(n: &str) -> Option<CubeScalarUDFKind> {
     if n == "DATE_SUB" {
         return Some(CubeScalarUDFKind::DateSub);
     }
+    if n == "DATE_BIN" {
+        return Some(CubeScalarUDFKind::DateBin);
+    }
     return None;
 }
 
@@ -192,6 +205,157 @@ impl CubeScalarUDF for UnixTimestamp {
     }
 }
 
+fn interval_to_nanoseconds(interval: &ScalarValue) -> Result<i64, DataFusionError> {
+    match interval {
+        ScalarValue::IntervalYearMonth(Some(i)) => {
+            // since years and months do not have a fixed number of nanoseconds,
+            // months can have different lengths, and leap years affect the length of a year,
+            // converting an IntervalYearMonth to nanoseconds requires certain assumptions:
+            // 1 month = 30 days
+            // 1 year = 12 months = 365 days
+
+            let years = i / 12;
+            let months = i % 12;
+            let total_nanoseconds =
+                (years as i64) * NANOS_IN_YEAR + (months as i64) * NANOS_IN_MONTH;
+
+            Ok(total_nanoseconds)
+        }
+        ScalarValue::IntervalDayTime(Some(i)) => {
+            let days: i64 = i.signum() * (i.abs() >> 32);
+            let millis: i64 = i.signum() * ((i.abs() << 32) >> 32);
+            let total_nanoseconds = days * NANOS_IN_DAY + millis * 1_000_000;
+
+            Ok(total_nanoseconds)
+        }
+        _ => Err(DataFusionError::Execution(format!(
+            "Unsupported interval type"
+        ))),
+    }
+}
+
+struct DateBin {}
+impl DateBin {
+    fn signature() -> Signature {
+        Signature::OneOf(vec![
+            Signature::Exact(vec![
+                DataType::Interval(IntervalUnit::YearMonth),
+                DataType::Timestamp(TimeUnit::Nanosecond, None),
+                DataType::Timestamp(TimeUnit::Nanosecond, None),
+            ]),
+            Signature::Exact(vec![
+                DataType::Interval(IntervalUnit::DayTime),
+                DataType::Timestamp(TimeUnit::Nanosecond, None),
+                DataType::Timestamp(TimeUnit::Nanosecond, None),
+            ]),
+        ])
+    }
+}
+impl CubeScalarUDF for DateBin {
+    fn kind(&self) -> CubeScalarUDFKind {
+        CubeScalarUDFKind::DateBin
+    }
+
+    fn name(&self) -> &str {
+        "DATE_BIN"
+    }
+
+    fn descriptor(&self) -> ScalarUDF {
+        return ScalarUDF {
+            name: self.name().to_string(),
+            signature: Self::signature(),
+            return_type: Arc::new(|_| {
+                Ok(Arc::new(DataType::Timestamp(TimeUnit::Nanosecond, None)))
+            }),
+            fun: Arc::new(move |inputs| {
+                assert_eq!(inputs.len(), 3);
+                let interval = match &inputs[0] {
+                    ColumnarValue::Scalar(i) => i.clone(),
+                    _ => {
+                        // We leave this case out for simplicity.
+                        // CubeStore does not allow intervals inside tables, so this is super rare.
+                        return Err(DataFusionError::Execution(format!(
+                            "Only scalar intervals are supported in DATE_BIN"
+                        )));
+                    }
+                };
+                let interval_ns = interval_to_nanoseconds(&interval)?;
+
+                let origin_ns = match &inputs[2] {
+                    ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(o))) => *o,
+                    ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(None)) => {
+                        return Err(DataFusionError::Execution(format!(
+                            "Third argument (origin) of DATE_BIN must be a non-null timestamp"
+                        )));
+                    }
+                    _ => {
+                        // Leaving out other rare cases.
+                        // The initial need for the date_bin comes from custom granularities support
+                        // and there will always be a scalar origin point
+                        return Err(DataFusionError::Execution(format!(
+                            "Only scalar origins are supported in DATE_BIN"
+                        )));
+                    }
+                };
+
+                match &inputs[1] {
+                    ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(None)) => Ok(
+                        ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(None)),
+                    ),
+                    ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(t))) => {
+                        let timestamp_ns = *t;
+                        let diff = timestamp_ns - origin_ns;
+                        let num_intervals = diff / interval_ns;
+                        let mut nearest_timestamp = origin_ns + num_intervals * interval_ns;
+
+                        // We need the closest date before, not after
+                        if diff < 0 {
+                            nearest_timestamp -= interval_ns
+                        }
+
+                        Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+                            Some(nearest_timestamp),
+                        )))
+                    }
+                    ColumnarValue::Array(t) if t.as_any().is::<TimestampNanosecondArray>() => {
+                        let t = t
+                            .as_any()
+                            .downcast_ref::<TimestampNanosecondArray>()
+                            .unwrap();
+
+                        let mut builder = TimestampNanosecondArray::builder(t.len());
+
+                        for i in 0..t.len() {
+                            if t.is_null(i) {
+                                builder.append_null()?;
+                            } else {
+                                let ts = t.value(i);
+                                let diff = ts - origin_ns;
+                                let num_intervals = diff / interval_ns;
+                                let mut nearest_timestamp = origin_ns + num_intervals * interval_ns;
+
+                                // We need the closest date before, not after
+                                if diff < 0 {
+                                    nearest_timestamp -= interval_ns
+                                }
+
+                                builder.append_value(nearest_timestamp)?;
+                            }
+                        }
+
+                        Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef))
+                    }
+                    _ => {
+                        return Err(DataFusionError::Execution(format!(
+                            "Second argument of DATE_BIN must be a non-null timestamp"
+                        )));
+                    }
+                }
+            }),
+        };
+    }
+}
+
 struct DateAddSub {
     is_add: bool,
 }