Fix null-safe result joins

src/databricks/labs/dqx/check_funcs.py

@@ -1623,12 +1623,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) ->
         # To retain the original records we need to join back to the input DataFrame.
         # Therefore, applying this check multiple times at once can potentially lead to long spark plans.
         # When applying large number of sql query checks, it may be beneficial to split it into separate runs.
-        joined_df = df.join(user_query_df_unique, on=merge_columns, how="left")
-
-        # we only care about original columns + condition
-        result_df = joined_df.select(*[joined_df[col] for col in df.columns], joined_df[unique_condition_column])
-
-        return result_df
+        return _join_results_on_null_safe_columns(df, user_query_df_unique, merge_columns, [unique_condition_column])
 
     if negate:
         message_expr = F.lit(msg) if msg else F.lit(f"Value is matching query: '{query}'")
@@ -1991,7 +1986,7 @@ def apply(df: DataFrame) -> DataFrame:
         # Step 5: join current bucket + stats (left-join so current bucket always survives)
         if group_by:
             join_keys = [c if isinstance(c, str) else get_column_name_or_alias(c) for c in group_by]
-            joined = current.join(stats, on=join_keys, how="left")
+            joined = _join_results_on_null_safe_columns(current, stats, join_keys, ["__dq_mu", "__dq_sigma", "__dq_n"])
         else:
             # Use a dummy-key left-join so the current bucket row survives even when stats
             # is empty (e.g. only 1 bucket of data → hist is empty → stats has 0 rows).
@@ -2037,7 +2032,7 @@ def apply(df: DataFrame) -> DataFrame:
         select_cols = [condition_col, msg_col]
         if group_by:
             join_keys = [c if isinstance(c, str) else get_column_name_or_alias(c) for c in group_by]
-            return df.join(result.select(*join_keys, *select_cols), on=join_keys, how="left")
+            return _join_results_on_null_safe_columns(df, result, join_keys, select_cols)
         return df.crossJoin(result.select(*select_cols))
 
     # Build alias
@@ -2922,6 +2917,20 @@ def _match_rows(
     return results
 
 
+def _join_results_on_null_safe_columns(
+    df: DataFrame, result_df: DataFrame, join_columns: list[str], result_columns: list[str]
+) -> DataFrame:
+    """Left-join computed result columns while matching null join keys."""
+    joined = _match_rows(
+        df.alias("df"),
+        result_df.alias("ref_df"),
+        join_columns,
+        join_columns,
+        check_missing_records=False,
+    )
+    return joined.select("df.*", *[F.col(f"ref_df.{column}").alias(column) for column in result_columns])
+
+
 def _add_row_diffs(
     df: DataFrame, pk_column_names: list[str], ref_pk_column_names: list[str], row_missing_col: str, row_extra_col: str
 ) -> DataFrame:
@@ -3374,7 +3383,7 @@ def apply(df: DataFrame) -> DataFrame:
                 # Note: Aliased Column expressions in group_by are not supported for window-incompatible
                 # aggregates (e.g., count_distinct). Use string column names or simple F.col() expressions.
                 join_cols = [col if isinstance(col, str) else get_column_name_or_alias(col) for col in group_by]
-                df = df.join(agg_df, on=join_cols, how="left")
+                df = _join_results_on_null_safe_columns(df, agg_df, join_cols, [metric_col])
             else:
                 # Use standard window function approach for window-compatible aggregates
                 window_spec = Window.partitionBy(*group_cols)

tests/integration/test_dataset_checks.py

@@ -21,6 +21,7 @@
     compare_datasets,
     is_data_fresh_per_time_window,
     has_valid_schema,
+    sql_query,
 )
 from databricks.labs.dqx.utils import get_column_name_or_alias
 from databricks.labs.dqx.errors import InvalidParameterError, MissingParameterError
@@ -31,6 +32,48 @@
 SCHEMA = "a: string, b: int"
 
 
+@pytest.mark.parametrize(
+    "schema, merge_columns, rows",
+    [
+        (
+            "row_id: string, amount: int",
+            ["row_id"],
+            [(None, 100), ("row-1", 100), ("row-2", -1)],
+        ),
+        (
+            "row_id: string, part_id: string, amount: int",
+            ["row_id", "part_id"],
+            [(None, "part-1", 100), ("row-1", None, 100), ("row-2", "part-2", -1)],
+        ),
+    ],
+)
+def test_sql_query_with_null_merge_columns(
+    spark: SparkSession,
+    schema: str,
+    merge_columns: list[str],
+    rows: list[tuple[str | int | None, ...]],
+):
+    """A true query condition must survive null keys in single and composite joins."""
+    test_df = spark.createDataFrame(rows, schema)
+    query_columns = ", ".join(merge_columns)
+    condition, apply_method = sql_query(
+        f"SELECT {query_columns}, amount > 0 AS condition FROM {{{{ input_view }}}}",
+        merge_columns=merge_columns,
+        msg="positive amount",
+    )
+
+    actual = apply_method(test_df, spark, {}).select(*merge_columns, "amount", condition.alias("violation"))
+    violations = {tuple(row[column] for column in merge_columns): row["violation"] for row in actual.collect()}
+
+    expected = {}
+    for row in rows:
+        amount = row[-1]
+        assert isinstance(amount, int)
+        expected[tuple(row[:-1])] = "positive amount" if amount > 0 else None
+
+    assert violations == expected
+
+
 def test_has_no_outliers_int_numeric_types(spark: SparkSession):
     test_df = spark.createDataFrame(
         [
@@ -1249,6 +1292,27 @@ def test_is_aggr_with_count_distinct_and_group_by(spark: SparkSession):
     assertDataFrameEqual(actual, expected, checkRowOrder=False)
 
 
+def test_is_aggr_with_count_distinct_and_null_group(spark: SparkSession):
+    """A violating null group must retain its aggregated metric after reattachment."""
+    test_df = spark.createDataFrame(
+        [[None, "val1"], [None, "val2"], ["group2", "val3"]],
+        "a: string, b: string",
+    )
+
+    actual = _apply_checks(
+        test_df,
+        [is_aggr_not_greater_than("b", limit=1, aggr_type="count_distinct", group_by=["a"])],
+    )
+
+    message = "Distinct count value 2 in column 'b' per group of columns 'a' is greater than limit: 1"
+    expected = spark.createDataFrame(
+        [[None, "val1", message], [None, "val2", message], ["group2", "val3", None]],
+        "a: string, b: string, b_count_distinct_group_by_a_greater_than_limit: string",
+    )
+
+    assertDataFrameEqual(actual, expected, checkRowOrder=False)
+
+
 def test_is_aggr_with_count_distinct_and_column_expression_in_group_by(spark: SparkSession):
     """Test count_distinct with Column expression (F.col) in group_by.
 

tests/integration/test_has_no_aggr_outliers.py

@@ -11,7 +11,7 @@
 3.  test_warmup_passes                      – only 3 days of history → pass (warmup)
 4.  test_constant_series_passes             – all identical values → sigma==0 → pass
 5.  test_null_current_passes                – most-recent bucket missing → pass
-6.  test_group_by_isolates_bands            – two groups, one spikes, one flat
+6.  test_group_by_isolates_bands            – two groups, including NULL, one spikes, one flat
 7.  test_row_filter                         – junk rows excluded before baseline
 8.  test_hour_truncation                    – time_interval="hour"
 9.  test_yaml_round_trip                    – load check from YAML via DQEngine
@@ -221,22 +221,23 @@ def test_null_current_passes(spark: SparkSession):
 # ---------------------------------------------------------------------------
 
 
-def test_group_by_isolates_bands(spark: SparkSession):
+@pytest.mark.parametrize("spike_group", ["spike", None])
+def test_group_by_isolates_bands(spark: SparkSession, spike_group: str | None):
     """
-    Two groups: 'stable' (14 days at 10 +/- 1, today=10) and 'spike' (14 days
-    at 10 +/- 1, today=50). Only the 'spike' group should violate.
+    Two groups: 'stable' (14 days at 10 +/- 1, today=10) and a nullable spike
+    group (14 days at 10 +/- 1, today=50). Only the spike group should violate.
     """
     hist_values = [9.0, 11.0] * 7  # mean=10, stddev_pop=1.0
 
-    rows = []
+    rows: list[tuple[str | None, date, float]] = []
     for i, value in enumerate(hist_values):
         event_date = BASE_DATE + timedelta(days=i)
         rows.append(("stable", event_date, value))
-        rows.append(("spike", event_date, value))
+        rows.append((spike_group, event_date, value))
 
     today = BASE_DATE + timedelta(days=14)
     rows.append(("stable", today, 10.0))  # within band
-    rows.append(("spike", today, 50.0))  # spike
+    rows.append((spike_group, today, 50.0))  # spike
 
     df = spark.createDataFrame(rows, "grp: string, event_date: date, metric: double")
 
@@ -252,7 +253,7 @@ def test_group_by_isolates_bands(spark: SparkSession):
     result = _apply_with_original(df, [(condition, apply_fn)])
 
     stable_msgs = [r[-1] for r in result.filter(F.col("grp") == "stable").collect()]
-    spike_msgs = [r[-1] for r in result.filter(F.col("grp") == "spike").collect()]
+    spike_msgs = [r[-1] for r in result.filter(F.col("grp").eqNullSafe(F.lit(spike_group))).collect()]
 
     assert all(m is None for m in stable_msgs), f"Stable group should not violate: {stable_msgs}"
     assert all(m is not None for m in spike_msgs), f"Spike group should violate: {spike_msgs}"