Feat: Introduce INCREMENTAL_BY_PARTITION model kind (#2687)

Author: izeigerman

setup.py

@@ -65,6 +65,8 @@
             "cryptography~=42.0.4",
             "dbt-core",
             "dbt-duckdb>=1.7.1",
+            "dbt-snowflake",
+            "dbt-bigquery",
             "Faker",
             "google-auth",
             "google-cloud-bigquery",
@@ -93,10 +95,8 @@
             "typing-extensions",
         ],
         "cicdtest": [
-            "dbt-bigquery",
             "dbt-databricks",
             "dbt-redshift",
-            "dbt-snowflake",
             "dbt-sqlserver>=1.7.0",
             "dbt-trino",
         ],

sqlmesh/core/dialect.py

@@ -491,6 +491,7 @@ def parse(self: Parser) -> t.Optional[exp.Expression]:
                     if kind in (
                         ModelKindName.INCREMENTAL_BY_TIME_RANGE,
                         ModelKindName.INCREMENTAL_BY_UNIQUE_KEY,
+                        ModelKindName.INCREMENTAL_BY_PARTITION,
                         ModelKindName.SEED,
                         ModelKindName.VIEW,
                         ModelKindName.SCD_TYPE_2,

sqlmesh/core/engine_adapter/base.py

@@ -1041,9 +1041,18 @@ def insert_overwrite_by_partition(
         partitioned_by: t.List[exp.Expression],
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
     ) -> None:
-        raise NotImplementedError(
-            "Insert Overwrite by Partition (not time) is not supported by this engine"
-        )
+        if self.INSERT_OVERWRITE_STRATEGY.is_insert_overwrite:
+            target_table = exp.to_table(table_name)
+            source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
+                query_or_df, columns_to_types, target_table=target_table
+            )
+            self._insert_overwrite_by_condition(
+                table_name, source_queries, columns_to_types=columns_to_types
+            )
+        else:
+            self._replace_by_key(
+                table_name, query_or_df, columns_to_types, partitioned_by, is_unique_key=False
+            )
 
     def insert_overwrite_by_time_partition(
         self,
@@ -2006,6 +2015,49 @@ def _truncate_table(self, table_name: TableName) -> None:
         table = exp.to_table(table_name)
         self.execute(f"TRUNCATE TABLE {table.sql(dialect=self.dialect, identify=True)}")
 
+    def _replace_by_key(
+        self,
+        target_table: TableName,
+        source_table: QueryOrDF,
+        columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
+        key: t.Sequence[exp.Expression],
+        is_unique_key: bool,
+    ) -> None:
+        if columns_to_types is None:
+            columns_to_types = self.columns(target_table)
+
+        temp_table = self._get_temp_table(target_table)
+        key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *key)
+        column_names = list(columns_to_types or [])
+
+        with self.transaction():
+            self.ctas(temp_table, source_table, columns_to_types=columns_to_types, exists=False)
+
+            try:
+                delete_query = exp.select(key_exp).from_(temp_table)
+                insert_query = self._select_columns(columns_to_types).from_(temp_table)
+                if not is_unique_key:
+                    delete_query = delete_query.distinct()
+                else:
+                    insert_query = insert_query.distinct(*key)
+
+                insert_statement = exp.insert(
+                    insert_query,
+                    target_table,
+                    columns=column_names,
+                )
+                delete_filter = key_exp.isin(query=delete_query)
+
+                if not self.INSERT_OVERWRITE_STRATEGY.is_replace_where:
+                    self.execute(exp.delete(target_table).where(delete_filter))
+                else:
+                    insert_statement.set("where", delete_filter)
+                    insert_statement.set("this", exp.to_table(target_table))
+
+                self.execute(insert_statement)
+            finally:
+                self.drop_table(temp_table)
+
     def _build_create_comment_table_exp(
         self, table: exp.Table, table_comment: str, table_kind: str
     ) -> exp.Comment | str:

sqlmesh/core/engine_adapter/mixins.py

@@ -41,31 +41,9 @@ def merge(
             raise SQLMeshError(
                 "This engine does not support MERGE expressions and therefore `when_matched` is not supported."
             )
-        if columns_to_types is None:
-            columns_to_types = self.columns(target_table)
-
-        temp_table = self._get_temp_table(target_table)
-        unique_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *unique_key)
-        column_names = list(columns_to_types or [])
-
-        with self.transaction():
-            self.ctas(temp_table, source_table, columns_to_types=columns_to_types, exists=False)
-            self.execute(
-                exp.delete(target_table).where(
-                    unique_exp.isin(query=exp.select(unique_exp).from_(temp_table))
-                )
-            )
-            self.execute(
-                exp.insert(
-                    self._select_columns(columns_to_types)
-                    .distinct(*unique_key)
-                    .from_(temp_table)
-                    .subquery(),
-                    target_table,
-                    columns=column_names,
-                )
-            )
-            self.drop_table(temp_table)
+        self._replace_by_key(
+            target_table, source_table, columns_to_types, unique_key, is_unique_key=True
+        )
 
 
 class PandasNativeFetchDFSupportMixin(EngineAdapter):

sqlmesh/core/model/__init__.py

@@ -21,6 +21,7 @@
     IncrementalByTimeRangeKind as IncrementalByTimeRangeKind,
     IncrementalByUniqueKeyKind as IncrementalByUniqueKeyKind,
     IncrementalUnmanagedKind as IncrementalUnmanagedKind,
+    IncrementalByPartitionKind as IncrementalByPartitionKind,
     ModelKind as ModelKind,
     ModelKindMixin as ModelKindMixin,
     ModelKindName as ModelKindName,

sqlmesh/core/model/kind.py

@@ -53,6 +53,10 @@ def is_incremental_by_time_range(self) -> bool:
     def is_incremental_by_unique_key(self) -> bool:
         return self.model_kind_name == ModelKindName.INCREMENTAL_BY_UNIQUE_KEY
 
+    @property
+    def is_incremental_by_partition(self) -> bool:
+        return self.model_kind_name == ModelKindName.INCREMENTAL_BY_PARTITION
+
     @property
     def is_incremental_unmanaged(self) -> bool:
         return self.model_kind_name == ModelKindName.INCREMENTAL_UNMANAGED
@@ -62,6 +66,7 @@ def is_incremental(self) -> bool:
         return (
             self.is_incremental_by_time_range
             or self.is_incremental_by_unique_key
+            or self.is_incremental_by_partition
             or self.is_incremental_unmanaged
             or self.is_scd_type_2
         )
@@ -129,6 +134,7 @@ class ModelKindName(str, ModelKindMixin, Enum):
 
     INCREMENTAL_BY_TIME_RANGE = "INCREMENTAL_BY_TIME_RANGE"
     INCREMENTAL_BY_UNIQUE_KEY = "INCREMENTAL_BY_UNIQUE_KEY"
+    INCREMENTAL_BY_PARTITION = "INCREMENTAL_BY_PARTITION"
     INCREMENTAL_UNMANAGED = "INCREMENTAL_UNMANAGED"
     FULL = "FULL"
     # Legacy alias to SCD Type 2 By Time
@@ -399,6 +405,20 @@ def data_hash_values(self) -> t.List[t.Optional[str]]:
         ]
 
 
+class IncrementalByPartitionKind(_Incremental):
+    name: Literal[ModelKindName.INCREMENTAL_BY_PARTITION] = ModelKindName.INCREMENTAL_BY_PARTITION
+    forward_only: Literal[True] = True
+    disable_restatement: SQLGlotBool = True
+
+    @property
+    def metadata_hash_values(self) -> t.List[t.Optional[str]]:
+        return [
+            *super().metadata_hash_values,
+            str(self.forward_only),
+            str(self.disable_restatement),
+        ]
+
+
 class IncrementalUnmanagedKind(_Incremental):
     name: Literal[ModelKindName.INCREMENTAL_UNMANAGED] = ModelKindName.INCREMENTAL_UNMANAGED
     insert_overwrite: SQLGlotBool = False
@@ -411,7 +431,11 @@ def data_hash_values(self) -> t.List[t.Optional[str]]:
 
     @property
     def metadata_hash_values(self) -> t.List[t.Optional[str]]:
-        return [*super().metadata_hash_values, str(self.forward_only)]
+        return [
+            *super().metadata_hash_values,
+            str(self.forward_only),
+            str(self.disable_restatement),
+        ]
 
 
 class ViewKind(_ModelKind):
@@ -581,6 +605,7 @@ class ExternalKind(_ModelKind):
         FullKind,
         IncrementalByTimeRangeKind,
         IncrementalByUniqueKeyKind,
+        IncrementalByPartitionKind,
         IncrementalUnmanagedKind,
         SeedKind,
         ViewKind,
@@ -596,6 +621,7 @@ class ExternalKind(_ModelKind):
     ModelKindName.FULL: FullKind,
     ModelKindName.INCREMENTAL_BY_TIME_RANGE: IncrementalByTimeRangeKind,
     ModelKindName.INCREMENTAL_BY_UNIQUE_KEY: IncrementalByUniqueKeyKind,
+    ModelKindName.INCREMENTAL_BY_PARTITION: IncrementalByPartitionKind,
     ModelKindName.INCREMENTAL_UNMANAGED: IncrementalUnmanagedKind,
     ModelKindName.SEED: SeedKind,
     ModelKindName.VIEW: ViewKind,

sqlmesh/core/model/meta.py

@@ -300,6 +300,8 @@ def _kind_validator(cls, values: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
             for field in ("partitioned_by_", "clustered_by"):
                 if values.get(field) and not kind.is_materialized:
                     raise ValueError(f"{field} field cannot be set for {kind} models")
+            if kind.is_incremental_by_partition and not values.get("partitioned_by_"):
+                raise ValueError(f"partitioned_by field is required for {kind.name} models")
         return values
 
     @property

sqlmesh/core/snapshot/evaluator.py

@@ -865,6 +865,8 @@ def _evaluation_strategy(snapshot: SnapshotInfoLike, adapter: EngineAdapter) ->
         klass = IncrementalByTimeRangeStrategy
     elif snapshot.is_incremental_by_unique_key:
         klass = IncrementalByUniqueKeyStrategy
+    elif snapshot.is_incremental_by_partition:
+        klass = IncrementalByPartitionStrategy
     elif snapshot.is_incremental_unmanaged:
         klass = IncrementalUnmanagedStrategy
     elif snapshot.is_view:
@@ -1209,6 +1211,26 @@ def delete(self, table_name: str) -> None:
         logger.info("Dropped table '%s'", table_name)
 
 
+class IncrementalByPartitionStrategy(MaterializableStrategy):
+    def insert(
+        self,
+        snapshot: Snapshot,
+        name: str,
+        query_or_df: QueryOrDF,
+        snapshots: t.Dict[str, Snapshot],
+        deployability_index: DeployabilityIndex,
+        batch_index: int,
+        **kwargs: t.Any,
+    ) -> None:
+        model = snapshot.model
+        self.adapter.insert_overwrite_by_partition(
+            name,
+            query_or_df,
+            partitioned_by=model.partitioned_by,
+            columns_to_types=model.columns_to_types,
+        )
+
+
 class IncrementalByTimeRangeStrategy(MaterializableStrategy):
     def insert(
         self,

tests/core/engine_adapter/test_base.py

@@ -10,6 +10,7 @@
 from sqlglot import parse_one
 from sqlglot.helper import ensure_list
 
+from sqlmesh.core import dialect as d
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.engine_adapter import EngineAdapter, EngineAdapterWithIndexSupport
 from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy
@@ -2674,3 +2675,67 @@ def test_pre_ping(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable)
     ]
 
     adapter._connection_pool.get().close.assert_called_once()
+
+
+def test_insert_overwrite_by_partition_query(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "test_schema.test_table"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
+
+    adapter.insert_overwrite_by_partition(
+        table_name,
+        parse_one("SELECT a, ds, b FROM tbl"),
+        partitioned_by=[
+            d.parse_one("DATETIME_TRUNC(ds, MONTH)"),
+            d.parse_one("b"),
+        ],
+        columns_to_types={
+            "a": exp.DataType.build("int"),
+            "ds": exp.DataType.build("DATETIME"),
+            "b": exp.DataType.build("boolean"),
+        },
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        'CREATE TABLE "test_schema"."__temp_test_table_abcdefgh" AS SELECT "a", "ds", "b" FROM "tbl"',
+        'DELETE FROM "test_schema"."test_table" WHERE CONCAT_WS(\'__SQLMESH_DELIM__\', DATETIME_TRUNC("ds", MONTH), "b") IN (SELECT DISTINCT CONCAT_WS(\'__SQLMESH_DELIM__\', DATETIME_TRUNC("ds", MONTH), "b") FROM "test_schema"."__temp_test_table_abcdefgh")',
+        'INSERT INTO "test_schema"."test_table" ("a", "ds", "b") SELECT "a", "ds", "b" FROM "test_schema"."__temp_test_table_abcdefgh"',
+        'DROP TABLE IF EXISTS "test_schema"."__temp_test_table_abcdefgh"',
+    ]
+
+
+def test_insert_overwrite_by_partition_query_insert_overwrite_strategy(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+    adapter.INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.INSERT_OVERWRITE
+
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "test_schema.test_table"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
+
+    adapter.insert_overwrite_by_partition(
+        table_name,
+        parse_one("SELECT a, ds, b FROM tbl"),
+        partitioned_by=[
+            d.parse_one("DATETIME_TRUNC(ds, MONTH)"),
+            d.parse_one("b"),
+        ],
+        columns_to_types={
+            "a": exp.DataType.build("int"),
+            "ds": exp.DataType.build("DATETIME"),
+            "b": exp.DataType.build("boolean"),
+        },
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        'INSERT OVERWRITE TABLE "test_schema"."test_table" ("a", "ds", "b") SELECT "a", "ds", "b" FROM "tbl"'
+    ]

tests/core/engine_adapter/test_databricks.py

@@ -6,6 +6,7 @@
 from pytest_mock import MockFixture
 from sqlglot import exp, parse_one
 
+from sqlmesh.core import dialect as d
 from sqlmesh.core.engine_adapter import DatabricksEngineAdapter
 from tests.core.engine_adapter import to_sql_calls
 
@@ -103,3 +104,35 @@ def test_get_current_database(make_mocked_engine_adapter: t.Callable):
 
     assert adapter.get_current_database() == "test_database"
     assert to_sql_calls(adapter) == ["SELECT CURRENT_DATABASE()"]
+
+
+def test_insert_overwrite_by_partition_query(
+    make_mocked_engine_adapter: t.Callable, mocker: MockFixture, make_temp_table_name: t.Callable
+):
+    adapter = make_mocked_engine_adapter(DatabricksEngineAdapter)
+
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "test_schema.test_table"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
+
+    adapter.insert_overwrite_by_partition(
+        table_name,
+        parse_one("SELECT a, ds, b FROM tbl"),
+        partitioned_by=[
+            d.parse_one("DATETIME_TRUNC(ds, MONTH)"),
+            d.parse_one("b"),
+        ],
+        columns_to_types={
+            "a": exp.DataType.build("int"),
+            "ds": exp.DataType.build("DATETIME"),
+            "b": exp.DataType.build("boolean"),
+        },
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        "CREATE TABLE `test_schema`.`temp_test_table_abcdefgh` AS SELECT `a`, `ds`, `b` FROM `tbl`",
+        "INSERT INTO `test_schema`.`test_table` REPLACE WHERE CONCAT_WS('__SQLMESH_DELIM__', DATE_TRUNC('MONTH', `ds`), `b`) IN (SELECT DISTINCT CONCAT_WS('__SQLMESH_DELIM__', DATE_TRUNC('MONTH', `ds`), `b`) FROM `test_schema`.`temp_test_table_abcdefgh`) SELECT `a`, `ds`, `b` FROM `test_schema`.`temp_test_table_abcdefgh`",
+        "DROP TABLE IF EXISTS `test_schema`.`temp_test_table_abcdefgh`",
+    ]