Feat!: add more interval units and make the smallest 5 minutes (#1471)

additional performance boosts

Author: tobymao

sqlmesh/core/engine_adapter/bigquery.py

@@ -429,7 +429,7 @@ def _create_table_properties(
             if (
                 isinstance(this, exp.Column)
                 and partition_interval_unit is not None
-                and partition_interval_unit != IntervalUnit.MINUTE
+                and not partition_interval_unit.is_minute
             ):
                 column_type: t.Optional[exp.DataType] = (columns_to_types or {}).get(this.name)
 

sqlmesh/core/node.py

@@ -21,6 +21,9 @@
     from sqlmesh.core.snapshot import Node
 
 
+DEFAULT_SAMPLE_SIZE = 5
+
+
 class IntervalUnit(str, Enum):
     """IntervalUnit is the inferred granularity of an incremental node.
 
@@ -33,17 +36,19 @@ class IntervalUnit(str, Enum):
     MONTH = "month"
     DAY = "day"
     HOUR = "hour"
-    MINUTE = "minute"
+    HALF_HOUR = "half_hour"
+    QUARTER_HOUR = "quarter_hour"
+    FIVE_MINUTE = "five_minute"
 
     @classmethod
-    def from_cron(klass, cron: str, sample_size: int = 10) -> IntervalUnit:
+    def from_cron(klass, cron: str, sample_size: int = DEFAULT_SAMPLE_SIZE) -> IntervalUnit:
         croniter = CroniterCache(cron)
         samples = [croniter.get_next() for _ in range(sample_size)]
         min_interval = min(b - a for a, b in zip(samples, samples[1:]))
         for unit, seconds in sorted(INTERVAL_SECONDS.items(), key=lambda x: x[1], reverse=True):
-            if seconds <= min_interval:
+            if seconds <= min_interval.total_seconds():
                 return unit
-        raise ConfigError(f"Invalid cron '{cron}': must have a cadence of 1 minute or more.")
+        raise ConfigError(f"Invalid cron '{cron}': must have a cadence of 5 minutes or more.")
 
     @property
     def is_date_granularity(self) -> bool:
@@ -67,12 +72,16 @@ def is_hour(self) -> bool:
 
     @property
     def is_minute(self) -> bool:
-        return self == IntervalUnit.MINUTE
+        return self in (IntervalUnit.FIVE_MINUTE, IntervalUnit.QUARTER_HOUR, IntervalUnit.HALF_HOUR)
 
     @property
     def _cron_expr(self) -> str:
-        if self == IntervalUnit.MINUTE:
-            return "* * * * *"
+        if self == IntervalUnit.FIVE_MINUTE:
+            return "*/5 * * * *"
+        if self == IntervalUnit.QUARTER_HOUR:
+            return "*/15 * * * *"
+        if self == IntervalUnit.HALF_HOUR:
+            return "*/30 * * * *"
         if self == IntervalUnit.HOUR:
             return "0 * * * *"
         if self == IntervalUnit.DAY:
@@ -132,7 +141,9 @@ def seconds(self) -> int:
     IntervalUnit.MONTH: 60 * 60 * 24 * 28,
     IntervalUnit.DAY: 60 * 60 * 24,
     IntervalUnit.HOUR: 60 * 60,
-    IntervalUnit.MINUTE: 60,
+    IntervalUnit.HALF_HOUR: 60 * 30,
+    IntervalUnit.QUARTER_HOUR: 60 * 15,
+    IntervalUnit.FIVE_MINUTE: 60 * 5,
 }
 
 
@@ -265,7 +276,7 @@ def croniter(self, value: TimeLike) -> CroniterCache:
         if self._croniter is None:
             self._croniter = CroniterCache(self.cron, value)
         else:
-            self._croniter.curr = value
+            self._croniter.curr = to_datetime(value)
         return self._croniter
 
     def cron_next(self, value: TimeLike) -> TimeLike:
@@ -315,7 +326,7 @@ def text_diff(self, other: Node) -> str:
         """
         raise NotImplementedError
 
-    def _inferred_interval_unit(self, sample_size: int = 10) -> IntervalUnit:
+    def _inferred_interval_unit(self, sample_size: int = DEFAULT_SAMPLE_SIZE) -> IntervalUnit:
         """Infers the interval unit from the cron expression.
 
         The interval unit is used to determine the lag applied to start_date and end_date for node rendering and intervals.

sqlmesh/core/plan/definition.py

@@ -4,6 +4,7 @@
 import sys
 import typing as t
 from collections import defaultdict
+from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
 
@@ -39,7 +40,6 @@
     yesterday_ds,
 )
 from sqlmesh.utils.errors import NoChangesPlanError, PlanError, SQLMeshError
-from sqlmesh.utils.pydantic import PydanticModel
 
 logger = logging.getLogger(__name__)
 
@@ -783,7 +783,9 @@ def is_finished(self) -> bool:
         return self == PlanStatus.FINISHED
 
 
-class SnapshotIntervals(PydanticModel, frozen=True):
+# millions of these can be created, pydantic has significant overhead
+@dataclass
+class SnapshotIntervals:
     snapshot_name: str
     intervals: Intervals
 
@@ -795,11 +797,12 @@ def format_intervals(self, unit: t.Optional[IntervalUnit] = None) -> str:
         return format_intervals(self.merged_intervals, unit)
 
 
+@dataclass
 class LoadedSnapshotIntervals(SnapshotIntervals):
+    change_category: SnapshotChangeCategory
     interval_unit: t.Optional[IntervalUnit]
     node_name: str
     view_name: t.Optional[str] = None
-    change_category: SnapshotChangeCategory
 
     @classmethod
     def from_snapshot(cls, snapshot: Snapshot) -> LoadedSnapshotIntervals:

sqlmesh/core/snapshot/definition.py

@@ -1217,7 +1217,7 @@ def missing_intervals(
         snapshot_start_date = start_dt
         snapshot_end_date = end_date
         if interval:
-            snapshot_start_date, snapshot_end_date = interval
+            snapshot_start_date, snapshot_end_date = (to_datetime(i) for i in interval)
             snapshot = snapshot.copy()
             snapshot.intervals = snapshot.intervals.copy()
             snapshot.remove_interval(interval, execution_time)

sqlmesh/migrations/v0027_minute_interval_to_five.py

@@ -0,0 +1,56 @@
+"""Change any interval unit of minute to five_minute."""
+import json
+
+import pandas as pd
+from sqlglot import exp
+
+from sqlmesh.utils.migration import index_text_type
+
+
+def migrate(state_sync):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+
+    new_snapshots = []
+
+    for name, identifier, version, snapshot, kind_name in engine_adapter.fetchall(
+        exp.select("name", "identifier", "version", "snapshot", "kind_name").from_(snapshots_table),
+        quote_identifiers=True,
+    ):
+        parsed_snapshot = json.loads(snapshot)
+
+        node = parsed_snapshot["node"]
+
+        if node.get("interval_unit") == "minute":
+            node["interval_unit"] = "five_minute"
+
+        new_snapshots.append(
+            {
+                "name": name,
+                "identifier": identifier,
+                "version": version,
+                "snapshot": json.dumps(parsed_snapshot),
+                "kind_name": kind_name,
+            }
+        )
+
+    if new_snapshots:
+        engine_adapter.delete_from(snapshots_table, "TRUE")
+
+        text_type = index_text_type(engine_adapter.dialect)
+
+        engine_adapter.insert_append(
+            snapshots_table,
+            pd.DataFrame(new_snapshots),
+            columns_to_types={
+                "name": exp.DataType.build(text_type),
+                "identifier": exp.DataType.build(text_type),
+                "version": exp.DataType.build(text_type),
+                "snapshot": exp.DataType.build("text"),
+                "kind_name": exp.DataType.build(text_type),
+            },
+            contains_json=True,
+        )

sqlmesh/utils/cron.py

@@ -1,31 +1,20 @@
 from __future__ import annotations
 
 import typing as t
-from datetime import timedelta
+from datetime import datetime, timedelta
 
 from croniter import croniter
 from sqlglot.helper import first
 
-from sqlmesh.utils import ttl_cache
 from sqlmesh.utils.date import TimeLike, now, to_datetime
 
 
-@ttl_cache()
-def cron_next(cron: str, time: TimeLike) -> float:
-    return croniter(cron, to_datetime(time)).get_next()
-
-
-@ttl_cache()
-def cron_prev(cron: str, time: TimeLike) -> float:
-    return croniter(cron, to_datetime(time)).get_prev()
-
-
 class CroniterCache:
-    ESTIMATE_SAMPLES_NUM = 10
+    ESTIMATE_SAMPLES_NUM = 5
 
     def __init__(self, cron: str, time: t.Optional[TimeLike] = None):
         self.cron = cron
-        self.curr: TimeLike = now() if time is None else time
+        self.curr: datetime = to_datetime(now() if time is None else time)
         self._interval_seconds: t.Optional[int] = None
 
     @property
@@ -39,30 +28,30 @@ def interval_seconds(self) -> int:
         """
         if self._interval_seconds is None:
             seconds = set()
-            curr = to_datetime(self.curr)
+            curr = self.curr
 
             for _ in range(self.ESTIMATE_SAMPLES_NUM):
                 prev = curr
-                curr = to_datetime(cron_next(self.cron, curr))
+                curr = to_datetime(croniter(self.cron, curr).get_next())
                 seconds.add(curr - prev)
 
             if len(seconds) == 1:
-                self._interval_seconds = first(seconds).seconds
+                self._interval_seconds = int(first(seconds).total_seconds())
             else:
                 self._interval_seconds = 0
 
         return self._interval_seconds
 
-    def get_next(self, estimate: bool = False) -> float:
+    def get_next(self, estimate: bool = False) -> datetime:
         if estimate and self.interval_seconds:
-            self.curr = to_datetime(self.curr) + timedelta(seconds=self.interval_seconds)
+            self.curr = self.curr + timedelta(seconds=self.interval_seconds)
         else:
-            self.curr = cron_next(self.cron, self.curr)
-        return t.cast(float, self.curr)
+            self.curr = to_datetime(croniter(self.cron, self.curr).get_next())
+        return self.curr
 
-    def get_prev(self, estimate: bool = False) -> float:
+    def get_prev(self, estimate: bool = False) -> datetime:
         if estimate and self.interval_seconds:
-            self.curr = to_datetime(self.curr) - timedelta(seconds=self.interval_seconds)
+            self.curr = self.curr - timedelta(seconds=self.interval_seconds)
         else:
-            self.curr = cron_prev(self.cron, self.curr)
-        return t.cast(float, self.curr)
+            self.curr = to_datetime(croniter(self.cron, self.curr).get_prev())
+        return self.curr

sqlmesh/utils/date.py

@@ -13,8 +13,6 @@
 import dateparser
 from sqlglot import exp
 
-from sqlmesh.utils import ttl_cache
-
 UTC = timezone.utc
 TimeLike = t.Union[date, datetime, str, int, float]
 MILLIS_THRESHOLD = time.time() + 100 * 365 * 24 * 3600
@@ -109,7 +107,6 @@ def to_timestamp(value: TimeLike, relative_base: t.Optional[datetime] = None) ->
     return int(to_datetime(value, relative_base=relative_base).timestamp() * 1000)
 
 
-@ttl_cache()
 def to_datetime(value: TimeLike, relative_base: t.Optional[datetime] = None) -> datetime:
     """Converts a value into a UTC datetime object.
 
@@ -150,7 +147,7 @@ def to_datetime(value: TimeLike, relative_base: t.Optional[datetime] = None) ->
         raise ValueError(f"Could not convert `{value}` to datetime.")
 
     if dt.tzinfo:
-        return dt.astimezone(UTC)
+        return dt if dt.tzinfo == UTC else dt.astimezone(UTC)
     return dt.replace(tzinfo=UTC)
 
 

tests/core/engine_adapter/test_bigquery.py

@@ -305,7 +305,7 @@ def test_create_table_date_partition(
 @pytest.mark.parametrize(
     "partition_by_cols, partition_column_type, partition_interval_unit, partition_by_statement",
     [
-        ([exp.to_column("ds")], "date", IntervalUnit.MINUTE, "`ds`"),
+        ([exp.to_column("ds")], "date", IntervalUnit.FIVE_MINUTE, "`ds`"),
         ([exp.to_column("ds")], "date", IntervalUnit.HOUR, "`ds`"),
         ([exp.to_column("ds")], "date", IntervalUnit.DAY, "`ds`"),
         ([exp.to_column("ds")], "date", IntervalUnit.MONTH, "DATE_TRUNC(`ds`, MONTH)"),
@@ -333,7 +333,7 @@ def test_create_table_date_partition(
         (
             [d.parse_one("TIMESTAMP_TRUNC(ds, DAY)", dialect="bigquery")],
             "timestamp",
-            IntervalUnit.MINUTE,
+            IntervalUnit.FIVE_MINUTE,
             "TIMESTAMP_TRUNC(`ds`, DAY)",
         ),
     ],

tests/core/test_model.py

@@ -1798,9 +1798,9 @@ def test_incremental_unmanaged_validation():
 def test_custom_interval_unit():
     assert (
         load_sql_based_model(
-            d.parse("MODEL (name db.table, interval_unit MINUTE); SELECT a FROM tbl;")
+            d.parse("MODEL (name db.table, interval_unit FIVE_MINUTE); SELECT a FROM tbl;")
         ).interval_unit
-        == IntervalUnit.MINUTE
+        == IntervalUnit.FIVE_MINUTE
     )
 
     assert (
@@ -1834,10 +1834,10 @@ def test_custom_interval_unit():
     assert (
         load_sql_based_model(
             d.parse(
-                "MODEL (name db.table, cron '0 5 * * *', interval_unit 'minute'); SELECT a FROM tbl;"
+                "MODEL (name db.table, cron '0 5 * * *', interval_unit 'quarter_hour'); SELECT a FROM tbl;"
             )
         ).interval_unit
-        == IntervalUnit.MINUTE
+        == IntervalUnit.QUARTER_HOUR
     )
 
     with pytest.raises(
@@ -2134,7 +2134,7 @@ def test_scd_type_2_defaults():
                 unique_key id,
             ),
         );
-        SELECT 
+        SELECT
             1 as id,
             '2020-01-01' as ds,
             '2020-01-01' as test_updated_at,
@@ -2172,7 +2172,7 @@ def test_scd_type_2_overrides():
                 disable_restatement False,
             ),
         );
-        SELECT 
+        SELECT
             1 as id,
             '2020-01-01' as ds,
             '2020-01-01' as test_updated_at,