From e5fe7bcdc1043e3d8d4b4e380324a91f30ba68e9 Mon Sep 17 00:00:00 2001
From: Jon Mease <jonmmease@gmail.com>
Date: Mon, 6 Nov 2023 09:42:36 -0500
Subject: [PATCH] Fix cast timestamp to integer millisecond, add stack
 timestamp tests

---
 vegafusion-common/src/datatypes.rs            |  9 ++-
 .../tests/test_transform_stack.rs             | 71 +++++++++++++++++++
 vegafusion-sql/src/dataframe/mod.rs           |  5 ++
 3 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/vegafusion-common/src/datatypes.rs b/vegafusion-common/src/datatypes.rs
index 089aa7629..debeb7846 100644
--- a/vegafusion-common/src/datatypes.rs
+++ b/vegafusion-common/src/datatypes.rs
@@ -85,9 +85,12 @@ pub fn to_numeric(value: Expr, schema: &DFSchema) -> Result<Expr> {
         value
     } else if matches!(dtype, DataType::Timestamp(_, _)) {
         // Convert to milliseconds
-        Expr::ScalarFunction(expr::ScalarFunction {
-            fun: BuiltinScalarFunction::ToTimestampMillis,
-            args: vec![value],
+        Expr::TryCast(TryCast {
+            expr: Box::new(Expr::ScalarFunction(expr::ScalarFunction {
+                fun: BuiltinScalarFunction::ToTimestampMillis,
+                args: vec![value],
+            })),
+            data_type: DataType::Int64,
         })
     } else {
         // Cast non-numeric types (like UTF-8) to Float64
diff --git a/vegafusion-runtime/tests/test_transform_stack.rs b/vegafusion-runtime/tests/test_transform_stack.rs
index 611b9d846..bc447fee8 100644
--- a/vegafusion-runtime/tests/test_transform_stack.rs
+++ b/vegafusion-runtime/tests/test_transform_stack.rs
@@ -311,3 +311,74 @@ mod test_stack_no_divide_by_zero {
     #[test]
     fn test_marker() {} // Help IDE detect test module
 }
+
+#[cfg(test)]
+mod test_stack_timestamp_group {
+    use super::*;
+    use serde_json::json;
+    use vegafusion_common::data::table::VegaFusionTable;
+    use vegafusion_core::spec::transform::stack::StackOffsetSpec;
+
+    #[rstest(
+        offset,
+        case(None),
+        case(Some(StackOffsetSpec::Zero)),
+        case(Some(StackOffsetSpec::Normalize)),
+        case(Some(StackOffsetSpec::Center))
+    )]
+    fn test(offset: Option<StackOffsetSpec>) {
+        let dataset = VegaFusionTable::from_json(&json!([
+            {
+              "RELEASE_DATE": "Jun 12 1998"
+            },
+            {
+              "RELEASE_DATE": "Aug 07 1998"
+            },
+            {
+              "RELEASE_DATE": "Aug 28 1998"
+            },
+            {
+              "RELEASE_DATE": "Sep 11 1998"
+            },
+            {
+              "RELEASE_DATE": "Oct 09 1998"
+            }
+        ]))
+        .unwrap();
+
+        let transform_specs: Vec<TransformSpec> = serde_json::from_value(json!([
+            {
+              "type": "formula",
+              "expr": "toDate(datum[\"RELEASE_DATE\"])",
+              "as": "RELEASE_DATE"
+            },
+            {
+              "type": "stack",
+              "groupby": [
+                "RELEASE_DATE"
+              ],
+              "field": "RELEASE_DATE",
+              "sort": {
+                "field": [],
+                "order": []
+              },
+              "as": [
+                "__RELEASE_DATE_start__",
+                "__RELEASE_DATE_end__"
+              ],
+              "offset": offset
+            }
+        ]))
+        .unwrap();
+
+        check_transform_evaluation(
+            &dataset,
+            transform_specs.as_slice(),
+            &Default::default(),
+            &Default::default(),
+        );
+    }
+
+    #[test]
+    fn test_marker() {} // Help IDE detect test module
+}
diff --git a/vegafusion-sql/src/dataframe/mod.rs b/vegafusion-sql/src/dataframe/mod.rs
index 5ee92268d..98a30db29 100644
--- a/vegafusion-sql/src/dataframe/mod.rs
+++ b/vegafusion-sql/src/dataframe/mod.rs
@@ -1401,6 +1401,11 @@ fn make_new_schema_from_exprs(
         } else {
             // Add field for expression
             let schema_df = DFSchema::try_from(schema.clone())?;
+            println!(
+                "Expr: {:?}\n{:?}",
+                expr.to_sql_select(dialect, &schema_df).unwrap().to_string(),
+                schema_df
+            );
             let dtype = expr.get_type(&schema_df)?;
             let name = expr.display_name()?;
             fields.push(Field::new(name, dtype, true));