feat(cubesql): Support CURRENT_DATE scalar function

cube-js · Jun 2, 2023 · ec928a6 · ec928a6
1 parent 5c298ba
commit ec928a6
Show file tree

Hide file tree

Showing 7 changed files with 95 additions and 12 deletions.
diff --git a/rust/cubesql/Cargo.lock b/rust/cubesql/Cargo.lock
diff --git a/rust/cubesql/cubesql/Cargo.toml b/rust/cubesql/cubesql/Cargo.toml
@@ -9,7 +9,7 @@ documentation = "https://cube.dev/docs"
 homepage = "https://cube.dev"
 
 [dependencies]
-datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "b4f61f57339e6dce162903975c78a2b0a7dd5295", default-features = false, features = ["regex_expressions", "unicode_expressions"] }
+datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "104887f467aaa7172bcfc8b96200231e115bc177", default-features = false, features = ["regex_expressions", "unicode_expressions"] }
 anyhow = "1.0"
 thiserror = "1.0"
 cubeclient = { path = "../cubeclient" }

diff --git a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs
@@ -9,7 +9,9 @@ use async_trait::async_trait;
 use cubeclient::models::{V1LoadRequestQuery, V1LoadResult, V1LoadResultAnnotation};
 pub use datafusion::{
     arrow::{
-        array::{ArrayRef, BooleanBuilder, Float64Builder, Int64Builder, StringBuilder},
+        array::{
+            ArrayRef, BooleanBuilder, Date32Builder, Float64Builder, Int64Builder, StringBuilder,
+        },
         datatypes::{DataType, SchemaRef},
         error::{ArrowError, Result as ArrowResult},
         record_batch::RecordBatch,
@@ -30,7 +32,7 @@ use crate::{
     transport::{CubeStreamReceiver, LoadRequestMeta, TransportService},
     CubeError,
 };
-use chrono::NaiveDateTime;
+use chrono::{Datelike, NaiveDate, NaiveDateTime};
 use datafusion::{
     arrow::{array::TimestampNanosecondBuilder, datatypes::TimeUnit},
     execution::context::TaskContext,
@@ -727,8 +729,7 @@ pub fn transform_response<V: ValueObject>(
                     {
                         (FieldValue::String(s), builder) => {
                             let timestamp = NaiveDateTime::parse_from_str(s.as_str(), "%Y-%m-%dT%H:%M:%S.%f")
-                        .or_else(|_| NaiveDateTime::parse_from_str(s.as_str(), "%Y-%m-%d %H:%M:%S.%f")
-                        )
+                                .or_else(|_| NaiveDateTime::parse_from_str(s.as_str(), "%Y-%m-%d %H:%M:%S.%f"))
                                 .map_err(|e| {
                                     DataFusionError::Execution(format!(
                                         "Can't parse timestamp: '{}': {}",
@@ -748,6 +749,43 @@ pub fn transform_response<V: ValueObject>(
                     }
                 )
             }
+            DataType::Date32 => {
+                build_column!(
+                    DataType::Date32,
+                    Date32Builder,
+                    response,
+                    field_name,
+                    {
+                        (FieldValue::String(s), builder) => {
+                            let date = NaiveDate::parse_from_str(s.as_str(), "%Y-%m-%d")
+                                .map_err(|e| {
+                                    DataFusionError::Execution(format!(
+                                        "Can't parse date: '{}': {}",
+                                        s, e
+                                    ))
+                                });
+                            match date {
+                                Ok(date) => {
+                                    let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
+                                    let days_since_epoch = date.num_days_from_ce()  - epoch.num_days_from_ce();
+                                    builder.append_value(days_since_epoch)?;
+                                }
+                                Err(error) => {
+                                    log::error!(
+                                        "Unable to parse value as Date32: {}",
+                                        error.to_string()
+                                    );
+
+                                    builder.append_null()?
+                                }
+                            }
+                        }
+                    },
+                    {
+                        (ScalarValue::Date32(v), builder) => builder.append_option(v.clone())?,
+                    }
+                )
+            }
             t => {
                 return Err(CubeError::user(format!(
                     "Type {} is not supported in response transformation from Cube",

diff --git a/rust/cubesql/cubesql/src/compile/engine/udf.rs b/rust/cubesql/cubesql/src/compile/engine/udf.rs
@@ -3567,7 +3567,6 @@ pub fn register_fun_stubs(mut ctx: SessionContext) -> SessionContext {
         rettyp = Utf8,
         vol = Stable
     );
-    register_fun_stub!(udf, "current_date", argc = 0, rettyp = Date32, vol = Stable);
     register_fun_stub!(udf, "current_query", argc = 0, rettyp = Utf8, vol = Stable);
     register_fun_stub!(udf, "current_role", argc = 0, rettyp = Utf8, vol = Stable);
     register_fun_stub!(

diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -11233,6 +11233,42 @@ ORDER BY \"COUNT(count)\" DESC"
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_current_date() -> Result<(), CubeError> {
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            "SELECT CURRENT_DATE AS \"COL\"".to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+        .await;
+
+        let logical_plan = &query_plan.print(true).unwrap();
+
+        let re = Regex::new(r#"Date32\("\d+"\)"#).unwrap();
+        let logical_plan = re
+            .replace_all(logical_plan, "Date32(\"0\")")
+            .as_ref()
+            .to_string();
+
+        assert_eq!(
+            logical_plan,
+            "Projection: Date32(\"0\") AS COL\
+            \n  EmptyRelation",
+        );
+
+        insta::assert_snapshot!(
+            "current_date",
+            execute_query(
+                "SELECT current_timestamp::date = current_date".to_string(),
+                DatabaseProtocol::PostgreSQL
+            )
+            .await?
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_union_ctes() -> Result<(), CubeError> {
         insta::assert_snapshot!(

diff --git a/rust/cubesql/cubesql/src/compile/rewrite/language.rs b/rust/cubesql/cubesql/src/compile/rewrite/language.rs
@@ -361,6 +361,7 @@ macro_rules! variant_field_struct {
                 BuiltinScalarFunction::ToDayInterval => "ToDayInterval",
                 BuiltinScalarFunction::Now => "Now",
                 BuiltinScalarFunction::UtcTimestamp => "UtcTimestamp",
+                BuiltinScalarFunction::CurrentDate => "CurrentDate",
                 BuiltinScalarFunction::Translate => "Translate",
                 BuiltinScalarFunction::Trim => "Trim",
                 BuiltinScalarFunction::Upper => "Upper",

diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__current_date.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__current_date.snap
@@ -0,0 +1,9 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"SELECT current_timestamp::date = current_date\".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
++---------------+
+| Boolean(true) |
++---------------+
+| true          |
++---------------+