Repository: TobikoData/sqlmesh
Pull request #2840 (2024-06-28 19:54:13)

Fix: Check for null or duplicate primary keys when running table_diff
This adds a check for the uniqueness of the primary key (grain) or whether it contains null values and accordingly shows a warning to the user, fixes: #2810 

By default, this behavior is turned off due to the overhead of the additional query, but it can be enabled with the `check_grain` flag.

```
sqlmesh table_diff prod:dev sqlmesh_example.full_model --show-sample --check-grain

Grain should have unique and not-null audits for accurate results.

Row Counts:
└──  FULL MATCH: 3 rows (100.0%)

COMMON ROWS column comparison stats:
            pct_match
num_orders      100.0

COMMON ROWS sample data differences:
  All joined rows match
```

georgesittas: > Thank you @georgesittas addressed your comments and suggestions

Thanks 👍, I'll take another look soon.
Themiscodes: Addressed comments, refactored distinct query, condensed tests to one with cases for multiple column keys and null values

Patch:
diff --git a/docs/guides/tablediff.md b/docs/guides/tablediff.md
index c0c02080..0300d93c 100644
--- a/docs/guides/tablediff.md
+++ b/docs/guides/tablediff.md
@@ -109,6 +109,15 @@ The `COMMON ROWS sample data differences` section displays the row whose `item_i
 
 The `PROD ONLY sample rows` section shows the one row that is present in `PROD` but not in `DEV`.
 
+If we add the `--skip-grain-check` option, the grain is not validated. By default without this flag, a warning is displayed to the user if rows contain null or duplicate grains.
+
+```bash linenums="1"
+$ sqlmesh table_diff prod:dev2 sqlmesh_example.incremental_model
+
+Grain should have unique and not-null audits for accurate results.
+
+```
+
 ## Diffing tables or views
 
 Compare specific tables or views with the SQLMesh CLI interface by using the command `sqlmesh table_diff [source table]:[target table]`.
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index d032c2f9..5efcc6a6 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -417,6 +417,8 @@ Options:
   --limit INTEGER         The limit of the sample dataframe.
   --show-sample           Show a sample of the rows that differ. With many
                           columns, the output can be very wide.
+  --skip-grain-check      Disable the check for a primary key (grain) that 
+                          is missing or is not unique.
   -d, --decimals INTEGER  The number of decimal places to keep when comparing
                           floating point columns. Default: 3
   --help                  Show this message and exit.
diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md
index 3c71041d..0932ddc6 100644
--- a/docs/reference/notebook.md
+++ b/docs/reference/notebook.md
@@ -267,7 +267,7 @@ Create a schema file containing external model schemas.
 #### table_diff
 ```
 %table_diff [--on [ON ...]] [--model MODEL] [--where WHERE]
-                  [--limit LIMIT] [--show-sample]
+                  [--limit LIMIT] [--show-sample] [--skip-grain-check]
                   SOURCE:TARGET
 
 Show the diff between two tables.
@@ -278,14 +278,15 @@ positional arguments:
   SOURCE:TARGET    Source and target in `SOURCE:TARGET` format
 
 options:
-  --on <[ON ...]>  The column to join on. Can be specified multiple times. The
-                   model grain will be used if not specified.
-  --model MODEL    The model to diff against when source and target are
-                   environments and not tables.
-  --where WHERE    An optional where statement to filter results.
-  --limit LIMIT    The limit of the sample dataframe.
-  --show-sample    Show a sample of the rows that differ. With many columns,
-                   the output can be very wide.
+  --on <[ON ...]>     The column to join on. Can be specified multiple times. The
+                      model grain will be used if not specified.
+  --model MODEL       The model to diff against when source and target are
+                      environments and not tables.
+  --where WHERE       An optional where statement to filter results.
+  --limit LIMIT       The limit of the sample dataframe.
+  --show-sample       Show a sample of the rows that differ. With many columns,
+                      the output can be very wide.  
+  --skip-grain-check  Disable the check for a primary key (grain) that is missing or is not unique.
 ```
 
 #### model
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 09fcf989..64016034 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -743,6 +743,11 @@ def create_external_models(obj: Context) -> None:
     default=3,
     help="The number of decimal places to keep when comparing floating point columns. Default: 3",
 )
+@click.option(
+    "--skip-grain-check",
+    is_flag=True,
+    help="Disable the check for a primary key (grain) that is missing or is not unique.",
+)
 @click.pass_obj
 @error_handler
 @cli_analytics
diff --git a/sqlmesh/core/console.py b/sqlmesh/core/console.py
index bab6635b..59ed0861 100644
--- a/sqlmesh/core/console.py
+++ b/sqlmesh/core/console.py
@@ -254,7 +254,9 @@ class Console(abc.ABC):
         """Show table schema diff."""
 
     @abc.abstractmethod
-    def show_row_diff(self, row_diff: RowDiff, show_sample: bool = True) -> None:
+    def show_row_diff(
+        self, row_diff: RowDiff, show_sample: bool = True, skip_grain_check: bool = False
+    ) -> None:
         """Show table summary diff."""
 
     def _limit_model_names(self, tree: Tree, verbose: bool = False) -> Tree:
@@ -1065,7 +1067,9 @@ class TerminalConsole(Console):
 
         self.console.print(tree)
 
-    def show_row_diff(self, row_diff: RowDiff, show_sample: bool = True) -> None:
+    def show_row_diff(
+        self, row_diff: RowDiff, show_sample: bool = True, skip_grain_check: bool = False
+    ) -> None:
         source_name = row_diff.source
         if row_diff.source_alias:
             source_name = row_diff.source_alias.upper()
@@ -1073,6 +1077,17 @@ class TerminalConsole(Console):
         if row_diff.target_alias:
             target_name = row_diff.target_alias.upper()
 
+        if row_diff.stats["null_grain_count"] > 0 or (
+            not skip_grain_check
+            and (
+                row_diff.stats["distinct_count_s"] != row_diff.stats["s_count"]
+                or row_diff.stats["distinct_count_t"] != row_diff.stats["t_count"]
+            )
+        ):
+            self.console.print(
+                "[b][red]\nGrain should have unique and not-null audits for accurate results.[/red][/b]"
+            )
+
         tree = Tree("[b]Row Counts:[/b]")
         if row_diff.full_match_count:
             tree.add(
@@ -2032,7 +2047,9 @@ class DebuggerTerminalConsole(TerminalConsole):
     def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
         self._write(schema_diff)
 
-    def show_row_diff(self, row_diff: RowDiff, show_sample: bool = True) -> None:
+    def show_row_diff(
+        self, row_diff: RowDiff, show_sample: bool = True, skip_grain_check: bool = False
+    ) -> None:
         self._write(row_diff)
 
 
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 16643dc0..b66e9a5d 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1275,6 +1275,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         show: bool = True,
         show_sample: bool = True,
         decimals: int = 3,
+        skip_grain_check: bool = False,
     ) -> TableDiff:
         """Show a diff between two tables.
 
@@ -1289,6 +1290,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             show: Show the table diff output in the console.
             show_sample: Show the sample dataframe in the console. Requires show=True.
             decimals: The number of decimal places to keep when comparing floating point columns.
+            skip_grain_check: Skip check for rows that contain null or duplicate grains.
 
         Returns:
             The TableDiff object containing schema and summary differences.
@@ -1337,7 +1339,11 @@ class GenericContext(BaseContext, t.Generic[C]):
         )
         if show:
             self.console.show_schema_diff(table_diff.schema_diff())
-            self.console.show_row_diff(table_diff.row_diff(), show_sample=show_sample)
+            self.console.show_row_diff(
+                table_diff.row_diff(skip_grain_check=skip_grain_check),
+                show_sample=show_sample,
+                skip_grain_check=skip_grain_check,
+            )
         return table_diff
 
     @python_api_analytics
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index 82a26c5e..e68fa3ce 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -4,7 +4,7 @@ import math
 import typing as t
 
 import pandas as pd
-from sqlglot import exp
+from sqlglot import exp, parse_one
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
@@ -205,7 +205,7 @@ class TableDiff:
             model_name=self.model_name,
         )
 
-    def row_diff(self) -> RowDiff:
+    def row_diff(self, skip_grain_check: bool = False) -> RowDiff:
         if self._row_diff is None:
             s_selects = {c: exp.column(c, "s").as_(f"s__{c}") for c in self.source_schema}
             t_selects = {c: exp.column(c, "t").as_(f"t__{c}") for c in self.target_schema}
@@ -277,6 +277,20 @@ class TableDiff:
                         1,
                         0,
                     ).as_("row_joined"),
+                    exp.func(
+                        "IF",
+                        exp.or_(
+                            *(
+                                exp.and_(
+                                    exp.column(c, "s").is_(exp.Null()),
+                                    exp.column(c, "t").is_(exp.Null()),
+                                )
+                                for c in index_cols
+                            ),
+                        ),
+                        1,
+                        0,
+                    ).as_("null_grain"),
                     *comparisons,
                 )
                 .from_(exp.alias_(self.source, "s"))
@@ -309,13 +323,26 @@ class TableDiff:
             temp_table = exp.table_("diff", db="sqlmesh_temp", quoted=True)
 
             with self.adapter.temp_table(query, name=temp_table) as table:
-                summary_query = exp.select(
+                summary_sums = [
                     exp.func("SUM", "s_exists").as_("s_count"),
                     exp.func("SUM", "t_exists").as_("t_count"),
                     exp.func("SUM", "row_joined").as_("join_count"),
+                    exp.func("SUM", "null_grain").as_("null_grain_count"),
                     exp.func("SUM", "row_full_match").as_("full_match_count"),
                     *(exp.func("SUM", name(c)).as_(c.alias) for c in comparisons),
-                ).from_(table)
+                ]
+
+                if not skip_grain_check:
+                    s_grains = ", ".join((f"s__{c}" for c in index_cols))
+                    t_grains = ", ".join((f"t__{c}" for c in index_cols))
+                    summary_sums.extend(
+                        [
+                            parse_one(f"COUNT(DISTINCT({s_grains}))").as_("distinct_count_s"),
+                            parse_one(f"COUNT(DISTINCT({t_grains}))").as_("distinct_count_t"),
+                        ]
+                    )
+
+                summary_query = exp.select(*summary_sums).from_(table)
 
                 stats_df = self.adapter.fetchdf(summary_query, quote_identifiers=True)
                 stats_df["s_only_count"] = stats_df["s_count"] - stats_df["join_count"]
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index eaa1700c..8d5eb692 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -604,6 +604,11 @@ class SQLMeshMagics(Magics):
         default=3,
         help="The number of decimal places to keep when comparing floating point columns. Default: 3",
     )
+    @argument(
+        "--skip-grain-check",
+        action="store_true",
+        help="Disable the check for a primary key (grain) that is missing or is not unique.",
+    )
     @line_magic
     @pass_sqlmesh_context
     def table_diff(self, context: Context, line: str) -> None:
@@ -622,6 +627,7 @@ class SQLMeshMagics(Magics):
             limit=args.limit,
             show_sample=args.show_sample,
             decimals=args.decimals,
+            skip_grain_check=args.skip_grain_check,
         )
 
     @magic_arguments()


Test patch:
diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py
index e3555316..598609a9 100644
--- a/tests/core/test_table_diff.py
+++ b/tests/core/test_table_diff.py
@@ -1,9 +1,10 @@
 import pytest
 import pandas as pd
 from sqlglot import exp
+from sqlmesh.core import dialect as d
 
 from sqlmesh.core.config import AutoCategorizationMode, CategorizerConfig
-from sqlmesh.core.model import SqlModel
+from sqlmesh.core.model import SqlModel, load_sql_based_model
 
 
 @pytest.mark.slow
@@ -134,3 +135,85 @@ def test_data_diff_decimals(sushi_context_fixed_date):
     )
     assert diff.row_diff().full_match_count == 2
     assert diff.row_diff().partial_match_count == 1
+
+
+@pytest.mark.slow
+def test_grain_check(sushi_context_fixed_date):
+    expressions = d.parse(
+        """
+        MODEL (name memory.sushi.grain_items, kind full, grain(key_1, key_2));
+        SELECT
+            key_1,
+            key_2,
+            value,
+        FROM
+            (VALUES
+                (1, 1, 1),
+                (7, 4, 2),
+                (NULL, 3, 3),
+                (NULL, NULL, 3),
+                (1, 2, 2),
+                (4, NULL, 3),
+                (2, 3, 2),
+            ) AS t (key_1,key_2, value)
+    """
+    )
+    model_s = load_sql_based_model(expressions)
+    sushi_context_fixed_date.upsert_model(model_s)
+    sushi_context_fixed_date.plan(
+        "source_dev",
+        no_prompts=True,
+        auto_apply=True,
+        skip_tests=True,
+        start="2023-01-31",
+        end="2023-01-31",
+    )
+
+    model = sushi_context_fixed_date.models['"memory"."sushi"."grain_items"']
+
+    modified_model = model.dict()
+    modified_model["query"] = (
+        exp.select("*")
+        .from_(model.query.subquery())
+        .union(
+            "SELECT key_1, key_2, value FROM (VALUES (1, 6, 1),(1, 5, 3),(NULL, 2, 3),) AS t (key_1, key_2, value)"
+        )
+    )
+
+    modified_sqlmodel = SqlModel(**modified_model)
+    sushi_context_fixed_date.upsert_model(modified_sqlmodel)
+
+    sushi_context_fixed_date.auto_categorize_changes = CategorizerConfig(
+        sql=AutoCategorizationMode.FULL
+    )
+    sushi_context_fixed_date.plan(
+        "target_dev",
+        create_from="source_dev",
+        no_prompts=True,
+        auto_apply=True,
+        skip_tests=True,
+        start="2023-01-31",
+        end="2023-01-31",
+    )
+
+    diff = sushi_context_fixed_date.table_diff(
+        source="source_dev",
+        target="target_dev",
+        on=["key_1", "key_2"],
+        model_or_snapshot="sushi.grain_items",
+        skip_grain_check=False,
+    )
+
+    row_diff = diff.row_diff()
+    assert row_diff.full_match_count == 7
+    assert row_diff.full_match_pct == 93.33
+    assert row_diff.s_only_count == 2
+    assert row_diff.t_only_count == 5
+    assert row_diff.stats["join_count"] == 4
+    assert row_diff.stats["null_grain_count"] == 4
+    assert row_diff.stats["s_count"] != row_diff.stats["distinct_count_s"]
+    assert row_diff.stats["distinct_count_s"] == 7
+    assert row_diff.stats["t_count"] != row_diff.stats["distinct_count_t"]
+    assert row_diff.stats["distinct_count_t"] == 10
+    assert row_diff.s_sample.shape == (0, 3)
+    assert row_diff.t_sample.shape == (3, 3)