Feat: Support partitioned_by expressions for bigquery (#1041)

Author: crericha

sqlmesh/core/dialect.py

@@ -14,6 +14,7 @@
 from sqlglot.tokens import Token
 
 from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE
+from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pandas import columns_to_types_from_df
 
 
@@ -463,6 +464,15 @@ class ChunkType(Enum):
     SQL = auto()
 
 
+def parse_one(sql: str, dialect: t.Optional[str] = None) -> exp.Expression:
+    expressions = parse(sql, default_dialect=dialect)
+    if not expressions:
+        raise SQLMeshError(f"No expressions found in '{sql}'")
+    elif len(expressions) > 1:
+        raise SQLMeshError(f"Multiple expressions found in '{sql}'")
+    return expressions[0]
+
+
 def parse(sql: str, default_dialect: t.Optional[str] = None) -> t.List[exp.Expression]:
     """Parse a sql string.
 

sqlmesh/core/engine_adapter/base.py

@@ -855,7 +855,7 @@ def temp_table(self, query_or_df: QueryOrDF, name: str = "diff") -> t.Iterator[e
     def _create_table_properties(
         self,
         storage_format: t.Optional[str] = None,
-        partitioned_by: t.Optional[t.List[str]] = None,
+        partitioned_by: t.Optional[t.List[exp.Expression]] = None,
         partition_interval_unit: t.Optional[IntervalUnit] = None,
     ) -> t.Optional[exp.Properties]:
         return None

sqlmesh/core/engine_adapter/bigquery.py

@@ -335,28 +335,38 @@ def _fetch_native_df(self, query: t.Union[exp.Expression, str]) -> DF:
     def _create_table_properties(
         self,
         storage_format: t.Optional[str] = None,
-        partitioned_by: t.Optional[t.List[str]] = None,
+        partitioned_by: t.Optional[t.List[exp.Expression]] = None,
         partition_interval_unit: t.Optional[IntervalUnit] = None,
     ) -> t.Optional[exp.Properties]:
         if not partitioned_by:
             return None
         if partition_interval_unit is None:
             raise SQLMeshError("partition_interval_unit is required when partitioning a table")
-        if partition_interval_unit == IntervalUnit.MINUTE:
-            raise SQLMeshError("BigQuery does not support partitioning by minute")
         if len(partitioned_by) > 1:
             raise SQLMeshError("BigQuery only supports partitioning by a single column")
-        partition_col = exp.to_column(partitioned_by[0])
-        this: t.Union[exp.Func, exp.Column]
-        if partition_interval_unit == IntervalUnit.HOUR:
-            this = exp.func(
-                "TIMESTAMP_TRUNC",
-                partition_col,
-                exp.var(IntervalUnit.HOUR.value.upper()),
-                dialect=self.dialect,
-            )
+
+        this: exp.Expression
+        if isinstance(partitioned_by[0], exp.Column):
+            if partition_interval_unit == IntervalUnit.MINUTE:
+                raise SQLMeshError("BigQuery does not support partitioning by minute")
+
+            trunc_func: t.Optional[str] = None
+            if partition_interval_unit == IntervalUnit.HOUR:
+                trunc_func = "TIMESTAMP_TRUNC"
+            elif partition_interval_unit in (IntervalUnit.MONTH, IntervalUnit.YEAR):
+                trunc_func = "DATE_TRUNC"
+
+            if trunc_func:
+                this = exp.func(
+                    trunc_func,
+                    partitioned_by[0],
+                    exp.var(partition_interval_unit.value.upper()),
+                    dialect=self.dialect,
+                )
+            else:
+                this = partitioned_by[0]
         else:
-            this = partition_col
+            this = partitioned_by[0]
 
         partition_columns_property = exp.PartitionedByProperty(this=this)
         return exp.Properties(expressions=[partition_columns_property])

sqlmesh/core/engine_adapter/spark.py

@@ -237,18 +237,21 @@ def create_view(
     def _create_table_properties(
         self,
         storage_format: t.Optional[str] = None,
-        partitioned_by: t.Optional[t.List[str]] = None,
+        partitioned_by: t.Optional[t.List[exp.Expression]] = None,
         partition_interval_unit: t.Optional[IntervalUnit] = None,
     ) -> t.Optional[exp.Properties]:
         format_property = None
         partition_columns_property = None
         if storage_format:
             format_property = exp.FileFormatProperty(this=exp.Var(this=storage_format))
         if partitioned_by:
+            for expr in partitioned_by:
+                if not isinstance(expr, exp.Column):
+                    raise SQLMeshError(
+                        f"PARTITIONED BY contains non-column value '{expr.sql(dialect='spark')}'."
+                    )
             partition_columns_property = exp.PartitionedByProperty(
-                this=exp.Schema(
-                    expressions=[exp.to_identifier(column) for column in partitioned_by]
-                ),
+                this=exp.Schema(expressions=partitioned_by),
             )
         return exp.Properties(
             expressions=[

sqlmesh/core/model/definition.py

@@ -100,7 +100,7 @@ class _Model(ModelMeta, frozen=True):
         lookback: The number of previous incremental intervals in the lookback window.
         storage_format: The storage format used to store the physical table, only applicable in certain engines.
             (eg. 'parquet')
-        partitioned_by: The partition columns, only applicable in certain engines. (eg. (ds, hour))
+        partitioned_by: The partition columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour))
         python_env: Dictionary containing all global variables needed to render the model's macros.
         mapping_schema: The schema of table names to column and types.
     """
@@ -557,7 +557,11 @@ def validate_definition(self) -> None:
             ConfigError
         """
         if self.partitioned_by:
-            unique_partition_keys = {k.strip().lower() for k in self.partitioned_by}
+            unique_partition_keys = {
+                col.name.strip().lower()
+                for expr in self.partitioned_by
+                for col in expr.find_all(exp.Column)
+            }
             if len(self.partitioned_by) != len(unique_partition_keys):
                 raise_config_error(
                     "All partition keys must be unique in the model definition",
@@ -1644,12 +1648,16 @@ def _single_value_or_tuple(values: t.Sequence) -> exp.Identifier | exp.Tuple:
     )
 
 
+def _single_expr_or_tuple(values: t.Sequence[exp.Expression]) -> exp.Expression | exp.Tuple:
+    return values[0] if len(values) == 1 else exp.Tuple(expressions=values)
+
+
 META_FIELD_CONVERTER: t.Dict[str, t.Callable] = {
     "name": lambda value: exp.to_table(value),
     "start": lambda value: exp.Literal.string(value),
     "cron": lambda value: exp.Literal.string(value),
     "batch_size": lambda value: exp.Literal.number(value),
-    "partitioned_by_": _single_value_or_tuple,
+    "partitioned_by_": _single_expr_or_tuple,
     "depends_on_": lambda value: exp.Tuple(expressions=value),
     "pre": _list_of_calls_to_exp,
     "post": _list_of_calls_to_exp,

sqlmesh/core/model/meta.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 import typing as t
-from datetime import timedelta
 from enum import Enum
 
 from pydantic import Field, root_validator, validator
 from sqlglot import exp
+from sqlglot.helper import ensure_list
 
 from sqlmesh.core import dialect as d
 from sqlmesh.core.model.kind import (
@@ -15,7 +15,6 @@
     ViewKind,
     _Incremental,
 )
-from sqlmesh.utils import unique
 from sqlmesh.utils.cron import CroniterCache
 from sqlmesh.utils.date import TimeLike, to_datetime
 from sqlmesh.utils.errors import ConfigError
@@ -25,15 +24,21 @@
 class IntervalUnit(str, Enum):
     """IntervalUnit is the inferred granularity of an incremental model.
 
-    IntervalUnit can be one of 4 types, DAY, HOUR, MINUTE. The unit is inferred
+    IntervalUnit can be one of 5 types, YEAR, MONTH, DAY, HOUR, MINUTE. The unit is inferred
     based on the cron schedule of a model. The minimum time delta between a sample set of dates
     is used to determine which unit a model's schedule is.
     """
 
+    YEAR = "year"
+    MONTH = "month"
     DAY = "day"
     HOUR = "hour"
     MINUTE = "minute"
 
+    @property
+    def is_date_granularity(self) -> bool:
+        return self in (IntervalUnit.YEAR, IntervalUnit.MONTH, IntervalUnit.DAY)
+
 
 AuditReference = t.Tuple[str, t.Dict[str, exp.Expression]]
 
@@ -51,7 +56,7 @@ class ModelMeta(PydanticModel):
     start: t.Optional[TimeLike]
     retention: t.Optional[int]  # not implemented yet
     storage_format: t.Optional[str]
-    partitioned_by_: t.List[str] = Field(default=[], alias="partitioned_by")
+    partitioned_by_: t.List[exp.Expression] = Field(default=[], alias="partitioned_by")
     depends_on_: t.Optional[t.Set[str]] = Field(default=None, alias="depends_on")
     columns_to_types_: t.Optional[t.Dict[str, exp.DataType]] = Field(default=None, alias="columns")
     column_descriptions_: t.Optional[t.Dict[str, str]]
@@ -110,7 +115,7 @@ def extract(v: exp.Expression) -> t.Tuple[str, t.Dict[str, str]]:
             ]
         return v
 
-    @validator("partitioned_by_", "tags", "grain", pre=True)
+    @validator("tags", "grain", pre=True)
     def _value_or_tuple_validator(cls, v: t.Any) -> t.Any:
         if isinstance(v, (exp.Tuple, exp.Array)):
             return [e.name for e in v.expressions]
@@ -136,6 +141,39 @@ def _cron_validator(cls, v: t.Any) -> t.Optional[str]:
                 raise ConfigError(f"Invalid cron expression '{cron}'")
         return cron
 
+    @validator("partitioned_by_", pre=True)
+    def _partition_by_validator(
+        cls, v: t.Any, values: t.Dict[str, t.Any]
+    ) -> t.List[exp.Expression]:
+        partitions: t.List[exp.Expression]
+        if isinstance(v, (exp.Tuple, exp.Array)):
+            partitions = v.expressions
+        elif isinstance(v, exp.Expression):
+            partitions = [v]
+        else:
+            dialect = values.get("dialect")
+            partitions = [
+                d.parse_one(entry, dialect=dialect) if isinstance(entry, str) else entry
+                for entry in ensure_list(v)
+            ]
+        partitions = [
+            exp.to_column(expr.name) if isinstance(expr, exp.Identifier) else expr
+            for expr in partitions
+        ]
+
+        for partition in partitions:
+            num_cols = len(list(partition.find_all(exp.Column)))
+            error_msg: t.Optional[str] = None
+            if num_cols == 0:
+                error_msg = "does not contain a column"
+            elif num_cols > 1:
+                error_msg = "contains multiple columns"
+
+            if error_msg:
+                raise ConfigError(f"partitioned_by field '{partition}' {error_msg}")
+
+        return partitions
+
     @validator("columns_to_types_", pre=True)
     def _columns_validator(
         cls, v: t.Any, values: t.Dict[str, t.Any]
@@ -194,9 +232,12 @@ def unique_key(self) -> t.List[str]:
         return []
 
     @property
-    def partitioned_by(self) -> t.List[str]:
-        time_column = [self.time_column.column] if self.time_column else []
-        return unique([*time_column, *self.partitioned_by_])
+    def partitioned_by(self) -> t.List[exp.Expression]:
+        if self.time_column and self.time_column.column not in [
+            col.name for col in self._partition_by_columns
+        ]:
+            return [*[exp.to_column(self.time_column.column)], *self.partitioned_by_]
+        return self.partitioned_by_
 
     @property
     def column_descriptions(self) -> t.Dict[str, str]:
@@ -208,18 +249,13 @@ def lookback(self) -> int:
         """The incremental lookback window."""
         return (self.kind.lookback if isinstance(self.kind, _Incremental) else 0) or 0
 
-    @property
-    def lookback_delta(self) -> timedelta:
-        """The incremental lookback time delta."""
-        if isinstance(self.kind, _Incremental):
-            interval_unit = self.interval_unit()
-            if interval_unit == IntervalUnit.DAY:
-                return timedelta(days=self.lookback)
-            if interval_unit == IntervalUnit.HOUR:
-                return timedelta(hours=self.lookback)
-            if interval_unit == IntervalUnit.MINUTE:
-                return timedelta(minutes=self.lookback)
-        return timedelta()
+    def lookback_start(self, start: TimeLike) -> TimeLike:
+        if self.lookback == 0:
+            return start
+
+        for _ in range(self.lookback):
+            start = self.cron_prev(start)
+        return start
 
     @property
     def batch_size(self) -> t.Optional[int]:
@@ -241,7 +277,11 @@ def interval_unit(self, sample_size: int = 10) -> IntervalUnit:
             croniter = CroniterCache(self.cron)
             samples = [croniter.get_next() for _ in range(sample_size)]
             min_interval = min(b - a for a, b in zip(samples, samples[1:]))
-            if min_interval >= 86400:
+            if min_interval >= 31536000:
+                self._interval_unit = IntervalUnit.YEAR
+            elif min_interval >= 2419200:
+                self._interval_unit = IntervalUnit.MONTH
+            elif min_interval >= 86400:
                 self._interval_unit = IntervalUnit.DAY
             elif min_interval >= 3600:
                 self._interval_unit = IntervalUnit.HOUR
@@ -252,8 +292,8 @@ def interval_unit(self, sample_size: int = 10) -> IntervalUnit:
     def normalized_cron(self) -> str:
         """Returns the UTC normalized cron based on sampling heuristics.
 
-        SQLMesh supports 3 interval units, daily, hourly, and minutes. If a job is scheduled
-        daily at 1PM, the actual intervals are shifted back to midnight UTC.
+        SQLMesh supports 5 interval units, yearly, monthly, daily, hourly, and minutes. If a
+        job is scheduled daily at 1PM, the actual intervals are shifted back to midnight UTC.
 
         Returns:
             The cron string representing either daily, hourly, or minutes.
@@ -265,6 +305,10 @@ def normalized_cron(self) -> str:
             return "0 * * * *"
         if unit == IntervalUnit.DAY:
             return "0 0 * * *"
+        if unit == IntervalUnit.MONTH:
+            return "0 0 1 * *"
+        if unit == IntervalUnit.YEAR:
+            return "0 0 1 1 *"
         return ""
 
     def croniter(self, value: TimeLike) -> CroniterCache:
@@ -309,3 +353,7 @@ def cron_floor(self, value: TimeLike) -> TimeLike:
             The timestamp floor.
         """
         return self.croniter(self.cron_next(value)).get_prev()
+
+    @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)]

sqlmesh/core/plan/definition.py

@@ -181,7 +181,7 @@ def _get_end_date(self, end_and_units: t.List[t.Tuple[int, IntervalUnit]]) -> Ti
         if end_and_units:
             end, unit = max(end_and_units)
 
-            if unit == IntervalUnit.DAY:
+            if unit.is_date_granularity:
                 return to_date(make_inclusive_end(end))
             return end
         return now()

sqlmesh/core/snapshot/definition.py

@@ -914,7 +914,7 @@ def _model_data_hash(model: Model) -> str:
         model.cron,
         model.storage_format,
         str(model.lookback),
-        *(model.partitioned_by or []),
+        *(expr.sql() for expr in (model.partitioned_by or [])),
         model.stamp,
     ]
 
@@ -1062,7 +1062,7 @@ def merge_intervals(intervals: Intervals) -> Intervals:
 
 
 def _format_date_time(time_like: TimeLike, unit: t.Optional[IntervalUnit]) -> str:
-    if unit is None or unit == IntervalUnit.DAY:
+    if unit is None or unit.is_date_granularity:
         return to_ds(time_like)
     return to_datetime(time_like).isoformat()[:19]