feat: add batch_size support to scd type 2 kinds

Author: eakmanrq

docs/concepts/models/model_kinds.md

@@ -1241,12 +1241,13 @@ This is the most accurate representation of the menu based on the source data pr
 
 ### Shared Configuration Options
 
-| Name                    | Description                                                                                                     | Type                      |
-|-------------------------|-----------------------------------------------------------------------------------------------------------------|---------------------------|
-| unique_key              | Unique key used for identifying rows between source and target                                                  | List of strings or string |
-| valid_from_name         | The name of the `valid_from` column to create in the target table. Default: `valid_from`                        | string                    |
-| valid_to_name           | The name of the `valid_to` column to create in the target table. Default: `valid_to`                            | string                    |
-| invalidate_hard_deletes | If set to `true`, when a record is missing from the source table it will be marked as invalid. Default: `false` | bool                      |
+| Name                    | Description                                                                                                                                                                                                                                                                                                       | Type                      |
+|-------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------|
+| unique_key              | Unique key used for identifying rows between source and target                                                                                                                                                                                                                                                    | List of strings or string |
+| valid_from_name         | The name of the `valid_from` column to create in the target table. Default: `valid_from`                                                                                                                                                                                                                          | string                    |
+| valid_to_name           | The name of the `valid_to` column to create in the target table. Default: `valid_to`                                                                                                                                                                                                                              | string                    |
+| invalidate_hard_deletes | If set to `true`, when a record is missing from the source table it will be marked as invalid. Default: `false`                                                                                                                                                                                                   | bool                      |
+| `batch_size`            | The maximum number of intervals that can be evaluated in a single backfill task. If this is `None`, all intervals will be processed as part of a single task. See [Processing Source Table with Historical Data](#processing-source-table-with-historical-data) for more info on this use case. (Default: `None`) | int                       |
 
 !!! tip "Important"
 
@@ -1278,6 +1279,18 @@ This is the most accurate representation of the menu based on the source data pr
 | columns                      | The name of the columns to check for changes. `*` to represent that all columns should be checked.                                                                                                                                          | List of strings or string |
 | execution_time_as_valid_from | By default, when the model is first loaded `valid_from` is set to `1970-01-01 00:00:00` and future new rows will have `execution_time` of when the pipeline ran. This changes the behavior to always use `execution_time`. Default: `false` | bool                      |
 
+
+### Processing Source Table with Historical Data
+
+The most common case for SCD Type 2 is creating history for a table that it doesn't have it already. 
+In the example of the restaurant menu, the menu just tells you what is offered right now, but you want to know what was offered over time.
+In this case, the default setting of `None` for `batch_size` is the best option.
+
+Another use case though is processing a source table that already has history in it. 
+A common example of this is a "daily snapshot" table that is created by a source system that takes a snapshot of the data at the end of each day.
+If your source table has historical records, like a "daily snapshot" table, then set `batch_size` to `1` to process each interval (each day if a `@daily` cron) in sequential order.
+That way the historical records will be properly captured in the SCD Type 2 table.
+
 ### Querying SCD Type 2 Models
 
 #### Querying the current version of a record

sqlmesh/core/model/kind.py

@@ -672,6 +672,7 @@ class _SCDType2Kind(_Incremental):
     valid_to_name: SQLGlotColumn = Field(exp.column("valid_to"), validate_default=True)
     invalidate_hard_deletes: SQLGlotBool = False
     time_data_type: exp.DataType = Field(exp.DataType.build("TIMESTAMP"), validate_default=True)
+    batch_size: t.Optional[SQLGlotPositiveInt] = None
 
     forward_only: SQLGlotBool = True
     disable_restatement: SQLGlotBool = True
@@ -711,6 +712,7 @@ def data_hash_values(self) -> t.List[t.Optional[str]]:
             gen(self.valid_to_name),
             str(self.invalidate_hard_deletes),
             gen(self.time_data_type),
+            gen(self.batch_size) if self.batch_size is not None else None,
         ]
 
     @property

sqlmesh/migrations/v0080_add_batch_size_to_scd_type_2_models.py

@@ -0,0 +1,5 @@
+"""Add batch_size to SCD Type 2 models which changes their data hash."""
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    pass

tests/core/test_model.py

@@ -4403,6 +4403,7 @@ def test_scd_type_2_by_column_overrides():
                 forward_only False,
                 disable_restatement False,
                 invalidate_hard_deletes False,
+                batch_size 1
             ),
         );
         SELECT
@@ -4428,6 +4429,7 @@ def test_scd_type_2_by_column_overrides():
     assert scd_type_2_model.kind.is_scd_type_2
     assert scd_type_2_model.kind.is_materialized
     assert scd_type_2_model.kind.time_data_type == exp.DataType.build("TIMESTAMPTZ")
+    assert scd_type_2_model.kind.batch_size == 1
     assert not scd_type_2_model.kind.invalidate_hard_deletes
     assert not scd_type_2_model.kind.forward_only
     assert not scd_type_2_model.kind.disable_restatement

tests/core/test_scheduler.py

@@ -13,6 +13,7 @@
     IncrementalByTimeRangeKind,
     IncrementalByUniqueKeyKind,
     TimeColumn,
+    SCDType2ByColumnKind,
 )
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.scheduler import (
@@ -810,3 +811,69 @@ def signal_base(batch: DatetimeRanges):
         snapshot_b: [(to_timestamp("2023-01-01"), to_timestamp("2023-01-04"))],
         snapshot_c: [(to_timestamp("2023-01-01"), to_timestamp("2023-01-02"))],
     }
+
+
+@pytest.mark.parametrize(
+    "batch_size, expected_batches",
+    [
+        (
+            1,
+            [
+                (to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
+                (to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
+                (to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
+            ],
+        ),
+        (
+            None,
+            [
+                (to_timestamp("2023-01-01"), to_timestamp("2023-01-04")),
+            ],
+        ),
+    ],
+)
+def test_scd_type_2_batch_size(
+    mocker: MockerFixture,
+    make_snapshot,
+    get_batched_missing_intervals,
+    batch_size: t.Optional[int],
+    expected_batches: t.List[t.Tuple[int, int]],
+):
+    """
+    Test that SCD_TYPE_2_BY_COLUMN models are batched correctly based on batch_size.
+    With batch_size=1, we expect 3 separate batches for 3 days.
+    Without a specified batch_size, we expect a single batch for the entire period.
+    """
+    start = to_datetime("2023-01-01")
+    end = to_datetime("2023-01-04")
+
+    # Configure kind params
+    kind_params = {}
+    if batch_size is not None:
+        kind_params["batch_size"] = batch_size
+
+    # Create the model and snapshot
+    model = SqlModel(
+        name="test_scd_model",
+        kind=SCDType2ByColumnKind(columns="valid_to", unique_key=["id"], **kind_params),
+        cron="@daily",
+        start=start,
+        query=parse_one("SELECT id, valid_from, valid_to FROM source"),
+    )
+    snapshot = make_snapshot(model)
+
+    # Setup scheduler
+    snapshot_evaluator = SnapshotEvaluator(adapters=mocker.MagicMock(), ddl_concurrent_tasks=1)
+    scheduler = Scheduler(
+        snapshots=[snapshot],
+        snapshot_evaluator=snapshot_evaluator,
+        state_sync=mocker.MagicMock(),
+        max_workers=2,
+        default_catalog=None,
+    )
+
+    # Get batches for the time period
+    batches = get_batched_missing_intervals(scheduler, start, end, end)[snapshot]
+
+    # Verify batches match expectations
+    assert batches == expected_batches