Restore support for Week and DayOfYear TimeUnits (#166)

* Re-enable week and year-week timeunit tests * Accept week and dayofyear time units * Implement fallback timeunit UDF for cases we don't handle with standard SQL
vega · Nov 2, 2022 · 2e00946 · 2e00946
1 parent 2e861d7
commit 2e00946
Show file tree

Hide file tree

Showing 5 changed files with 379 additions and 23 deletions.
diff --git a/vegafusion-core/src/spec/transform/timeunit.rs b/vegafusion-core/src/spec/transform/timeunit.rs
@@ -71,24 +71,12 @@ pub enum TimeUnitUnitSpec {
 
 impl TransformSpecTrait for TimeUnitTransformSpec {
     fn supported(&self) -> bool {
-        if self.units.is_none()
+        let unsupported = self.units.is_none()
             || self.step.is_some()
             || self.extent.is_some()
             || self.maxbins.is_some()
-            || self.signal.is_some()
-        {
-            false
-        } else {
-            if let Some(units) = &self.units {
-                for unit in units {
-                    if matches!(unit, TimeUnitUnitSpec::Week | TimeUnitUnitSpec::DayOfYear) {
-                        // week and dayofyear units are not yet supported
-                        return false;
-                    }
-                }
-            }
-            true
-        }
+            || self.signal.is_some();
+        !unsupported
     }
 
     fn output_signals(&self) -> Vec<String> {

diff --git a/vegafusion-rt-datafusion/src/expression/compiler/call.rs b/vegafusion-rt-datafusion/src/expression/compiler/call.rs
@@ -63,6 +63,7 @@ use crate::expression::compiler::builtin_functions::type_coercion::to_boolean::t
 use crate::expression::compiler::builtin_functions::type_coercion::to_number::to_number_transform;
 use crate::expression::compiler::builtin_functions::type_coercion::to_string::to_string_transform;
 use crate::task_graph::timezone::RuntimeTzConfig;
+use crate::transform::timeunit::TIMEUNIT_START_UDF;
 
 pub type MacroFn = Arc<dyn Fn(&[Expression]) -> Result<Expression> + Send + Sync>;
 pub type TransformFn = Arc<dyn Fn(&[Expr], &DFSchema) -> Result<Expr> + Send + Sync>;
@@ -454,6 +455,9 @@ pub fn make_session_context() -> SessionContext {
     ctx.register_udf((*MAKE_TIMESTAMPTZ).clone());
     ctx.register_udf((*TIMESTAMPTZ_TO_EPOCH_MS).clone());
 
+    // timeunit
+    ctx.register_udf((*TIMEUNIT_START_UDF).clone());
+
     // timeformat
     ctx.register_udf((*FORMAT_TIMESTAMP_UDF).clone());
 

diff --git a/vegafusion-rt-datafusion/src/sql/connection/datafusion_conn.rs b/vegafusion-rt-datafusion/src/sql/connection/datafusion_conn.rs
@@ -40,6 +40,9 @@ pub fn make_datafusion_dialect() -> Dialect {
     functions.insert("make_timestamptz".to_string());
     functions.insert("timestamptz_to_epoch_ms".to_string());
 
+    // timeunit
+    functions.insert("vega_timeunit".to_string());
+
     // timeformat
     functions.insert("format_timestamp".to_string());