Fix(clickhouse): support non-datetime time column partitioning (#3357)

Co-authored-by: Jo <46752250+georgesittas@users.noreply.github.com>

Author: treysp

sqlmesh/core/model/definition.py

@@ -982,6 +982,20 @@ def full_depends_on(self) -> t.Set[str]:
 
         return self._full_depends_on
 
+    @property
+    def partitioned_by(self) -> t.List[exp.Expression]:
+        """Columns to partition the model by, including the time column if it is not already included."""
+        if self.time_column and self.time_column.column not in {
+            col for expr in self.partitioned_by_ for col in expr.find_all(exp.Column)
+        }:
+            return [
+                TIME_COL_PARTITION_FUNC.get(self.dialect, lambda x, y: x)(
+                    self.time_column.column, self.columns_to_types
+                ),
+                *self.partitioned_by_,
+            ]
+        return self.partitioned_by_
+
 
 class _SqlBasedModel(_Model):
     inline_audits_: t.Dict[str, t.Any] = Field(default={}, alias="inline_audits")
@@ -2456,3 +2470,45 @@ def _meta_renderer(
 def get_model_name(path: Path) -> str:
     path_parts = list(path.parts[path.parts.index("models") + 1 : -1]) + [path.stem]
     return ".".join(path_parts[-3:])
+
+
+# function applied to time column when automatically used for partitioning in INCREMENTAL_BY_TIME_RANGE models
+def clickhouse_partition_func(
+    column: exp.Expression, columns_to_types: t.Optional[t.Dict[str, exp.DataType]]
+) -> exp.Expression:
+    # `toMonday()` function accepts a Date or DateTime type column
+
+    col_type = (columns_to_types and columns_to_types.get(column.name)) or exp.DataType.build(
+        "UNKNOWN"
+    )
+    col_type_is_conformable = col_type.is_type(
+        exp.DataType.Type.DATE,
+        exp.DataType.Type.DATE32,
+        exp.DataType.Type.DATETIME,
+        exp.DataType.Type.DATETIME64,
+    )
+
+    #  if input column is already a conformable type, just pass the column
+    if col_type_is_conformable:
+        return exp.func("toMonday", column, dialect="clickhouse")
+
+    # if input column type is not known, cast input to DateTime64
+    if col_type.is_type(exp.DataType.Type.UNKNOWN):
+        return exp.func(
+            "toMonday",
+            exp.cast(column, exp.DataType.build("DateTime64('UTC')", dialect="clickhouse")),
+            dialect="clickhouse",
+        )
+
+    # if input column type is known but not conformable, cast input to DateTime64 and cast output back to original type
+    return exp.cast(
+        exp.func(
+            "toMonday",
+            exp.cast(column, exp.DataType.build("DateTime64('UTC')", dialect="clickhouse")),
+            dialect="clickhouse",
+        ),
+        col_type,
+    )
+
+
+TIME_COL_PARTITION_FUNC = {"clickhouse": clickhouse_partition_func}

sqlmesh/core/model/meta.py

@@ -367,18 +367,6 @@ def unique_key(self) -> t.List[exp.Expression]:
             return self.kind.unique_key
         return []
 
-    @property
-    def partitioned_by(self) -> t.List[exp.Expression]:
-        """Columns to partition the model by, including the time column if it is not already included."""
-        if self.time_column and self.time_column.column not in [
-            col for col in self._partition_by_columns
-        ]:
-            return [
-                TIME_COL_PARTITION_FUNC.get(self.dialect, lambda x: x)(self.time_column.column),
-                *self.partitioned_by_,
-            ]
-        return self.partitioned_by_
-
     @property
     def clustered_by(self) -> t.List[exp.Expression]:
         return self.clustered_by_ or []
@@ -446,10 +434,6 @@ def all_references(self) -> t.List[Reference]:
             Reference(model_name=self.name, expression=e, unique=True) for e in self.references
         ]
 
-    @property
-    def _partition_by_columns(self) -> t.List[exp.Column]:
-        return [col for expr in self.partitioned_by_ for col in expr.find_all(exp.Column)]
-
     @property
     def managed_columns(self) -> t.Dict[str, exp.DataType]:
         return getattr(self.kind, "managed_columns", {})
@@ -478,12 +462,3 @@ def fqn(self) -> str:
     @property
     def on_destructive_change(self) -> OnDestructiveChange:
         return getattr(self.kind, "on_destructive_change", OnDestructiveChange.ALLOW)
-
-
-# function applied to time column when automatically used for partitioning in
-#   INCREMENTAL_BY_TIME_RANGE models
-TIME_COL_PARTITION_FUNC = {
-    "clickhouse": lambda x: exp.func(
-        "toMonday", exp.cast(x, exp.DataType.build("DateTime64", dialect="clickhouse"))
-    )
-}

tests/core/engine_adapter/test_clickhouse.py

@@ -350,6 +350,7 @@ def build_properties_sql(storage_format="", order_by="", primary_key="", propert
 
 
 def test_partitioned_by_expr(make_mocked_engine_adapter: t.Callable):
+    # user doesn't specify, unknown time column type
     model = load_sql_based_model(
         parse(
             """
@@ -369,8 +370,11 @@ def test_partitioned_by_expr(make_mocked_engine_adapter: t.Callable):
         )
     )
 
-    assert model.partitioned_by[0].sql("clickhouse") == 'toMonday(CAST("ds" AS DateTime64))'
+    assert (
+        model.partitioned_by[0].sql("clickhouse") == """toMonday(CAST("ds" AS DateTime64('UTC')))"""
+    )
 
+    # user specifies without time column, unknown time column type
     model = load_sql_based_model(
         parse(
             """
@@ -392,10 +396,58 @@ def test_partitioned_by_expr(make_mocked_engine_adapter: t.Callable):
     )
 
     assert [p.sql("clickhouse") for p in model.partitioned_by] == [
-        'toMonday(CAST("ds" AS DateTime64))',
+        """toMonday(CAST("ds" AS DateTime64('UTC')))""",
         '"x"',
     ]
 
+    # user doesn't specify, conformable date/datetime time column type
+    model = load_sql_based_model(
+        parse(
+            """
+        MODEL (
+            name foo,
+            dialect clickhouse,
+            kind INCREMENTAL_BY_TIME_RANGE(
+              time_column ds
+            )
+        );
+
+        select
+            ds::DATE as ds
+        from bar;
+        """,
+            default_dialect="clickhouse",
+        )
+    )
+
+    assert model.partitioned_by[0].sql("clickhouse") == 'toMonday("ds")'
+
+    # user doesn't specify, non-conformable time column type
+    model = load_sql_based_model(
+        parse(
+            """
+        MODEL (
+            name foo,
+            dialect clickhouse,
+            kind INCREMENTAL_BY_TIME_RANGE(
+              time_column ds
+            )
+        );
+
+        select
+            ds::String as ds
+        from bar;
+        """,
+            default_dialect="clickhouse",
+        )
+    )
+
+    assert (
+        model.partitioned_by[0].sql("clickhouse")
+        == """CAST(toMonday(CAST("ds" AS DateTime64('UTC'))) AS String)"""
+    )
+
+    # user specifies partitioned_by with time column
     model = load_sql_based_model(
         parse(
             """