Fix null-safe result joins

src/databricks/labs/dqx/check_funcs.py

@@ -1687,12 +1687,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}'")
@@ -1998,6 +1993,9 @@ def has_no_aggr_outliers(
     unique_str = uuid.uuid4().hex
     condition_col = f"__dq_outlier_cond_{aggr_col_str_norm}_{aggr_type}_{unique_str}"
     msg_col = f"__dq_outlier_msg_{aggr_col_str_norm}_{aggr_type}_{unique_str}"
+    internal_cols = {
+        name: f"__dq_{name}_{unique_str}" for name in ("grain", "metric", "rn", "current", "mu", "sigma", "n", "jkey")
+    }
 
     def apply(df: DataFrame) -> DataFrame:
         """
@@ -2024,69 +2022,79 @@ def apply(df: DataFrame) -> DataFrame:
         aggr_expr = _build_aggregate_expression(aggr_type, filtered_expr, aggr_params)
 
         group_cols = [F.col(c) if isinstance(c, str) else c for c in (group_by or [])]
-        grain_col = F.date_trunc(time_interval, F.col(time_column)).alias("__dq_grain")
+        grain_expr = F.date_trunc(time_interval, F.col(time_column)).alias(internal_cols["grain"])
 
         # Step 1: aggregate per time-grain (and group)
-        aggregate_df = df.groupBy(*group_cols, grain_col).agg(aggr_expr.alias("__dq_metric"))
+        aggregate_df = df.groupBy(*group_cols, grain_expr).agg(aggr_expr.alias(internal_cols["metric"]))
 
         # Step 2: rank grains per group, most-recent first
         window_spec = (
-            Window.partitionBy(*group_cols).orderBy(F.col("__dq_grain").desc())
+            Window.partitionBy(*group_cols).orderBy(F.col(internal_cols["grain"]).desc())
             if group_by
-            else Window.orderBy(F.col("__dq_grain").desc())
+            else Window.orderBy(F.col(internal_cols["grain"]).desc())
         )
-        ranked = aggregate_df.withColumn("__dq_rn", F.row_number().over(window_spec))
+        ranked = aggregate_df.withColumn(internal_cols["rn"], F.row_number().over(window_spec))
 
         # Step 3: most-recent bucket (rank 1) and history (ranks 2..lookback_num_intervals+1)
-        current = ranked.filter(F.col("__dq_rn") == 1).select(
+        current = ranked.filter(F.col(internal_cols["rn"]) == 1).select(
             *group_cols,
-            F.col("__dq_metric").alias("__dq_current"),
+            F.col(internal_cols["metric"]).alias(internal_cols["current"]),
+        )
+        hist = ranked.filter(
+            (F.col(internal_cols["rn"]) >= 2) & (F.col(internal_cols["rn"]) <= lookback_num_intervals + 1)
         )
-        hist = ranked.filter((F.col("__dq_rn") >= 2) & (F.col("__dq_rn") <= lookback_num_intervals + 1))
 
         # Step 4: rolling baseline stats
         stats = hist.groupBy(*group_cols).agg(
-            F.avg("__dq_metric").alias("__dq_mu"),
-            F.stddev_pop("__dq_metric").alias("__dq_sigma"),
-            F.count("*").alias("__dq_n"),
+            F.avg(internal_cols["metric"]).alias(internal_cols["mu"]),
+            F.stddev_pop(internal_cols["metric"]).alias(internal_cols["sigma"]),
+            F.count("*").alias(internal_cols["n"]),
         )
 
         # 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,
+                [internal_cols["mu"], internal_cols["sigma"], internal_cols["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).
             # A plain crossJoin with empty stats would produce 0 rows and discard all df rows.
             joined = (
-                current.withColumn("__dq_jkey", F.lit(1))
+                current.withColumn(internal_cols["jkey"], F.lit(1))
                 .join(
-                    stats.withColumn("__dq_jkey", F.lit(1)),
-                    on="__dq_jkey",
+                    stats.withColumn(internal_cols["jkey"], F.lit(1)),
+                    on=internal_cols["jkey"],
                     how="left",
                 )
-                .drop("__dq_jkey")
+                .drop(internal_cols["jkey"])
             )
 
         # Step 6: compute violation and message string (2 output columns: condition + msg)
-        # Guard on NULL __dq_n: occurs when hist is empty (fewer than 2 buckets)
-        insufficient_history = F.col("__dq_n").isNull() | (F.col("__dq_n") < warmup_num_intervals)
-        delta_expr = F.abs(F.col("__dq_current") - F.col("__dq_mu"))
+        # A NULL history count occurs when hist is empty (fewer than 2 buckets).
+        insufficient_history = F.col(internal_cols["n"]).isNull() | (F.col(internal_cols["n"]) < warmup_num_intervals)
+        delta_expr = F.abs(F.col(internal_cols["current"]) - F.col(internal_cols["mu"]))
         violation = (
             F.when(insufficient_history, F.lit(False))
-            .when(F.col("__dq_sigma").isNull() | (F.col("__dq_sigma") == 0), F.lit(False))
-            .when(F.col("__dq_current").isNull(), F.lit(False))
-            .otherwise(delta_expr > F.lit(sigma) * F.col("__dq_sigma"))
+            .when(
+                F.col(internal_cols["sigma"]).isNull() | (F.col(internal_cols["sigma"]) == 0),
+                F.lit(False),
+            )
+            .when(F.col(internal_cols["current"]).isNull(), F.lit(False))
+            .otherwise(delta_expr > F.lit(sigma) * F.col(internal_cols["sigma"]))
         )
         message = F.concat_ws(
             "",
             F.lit(f"{aggr_type}({aggr_col_str}): current="),
-            F.col("__dq_current").cast("string"),
+            F.col(internal_cols["current"]).cast("string"),
             F.lit(", baseline="),
-            F.col("__dq_mu").cast("string"),
+            F.col(internal_cols["mu"]).cast("string"),
             F.lit(", stddev="),
-            F.col("__dq_sigma").cast("string"),
+            F.col(internal_cols["sigma"]).cast("string"),
             F.lit(", delta="),
             delta_expr.cast("string"),
             F.lit(f" exceeds {sigma} x stddev (lookback={lookback_num_intervals} intervals)"),
@@ -2100,7 +2108,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
@@ -3159,6 +3167,42 @@ 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.
+
+    Preconditions are the caller's responsibility and are not validated here (all current callers
+    satisfy them):
+
+    - *join_columns* must be non-empty. An empty list makes the null-safe join condition reduce to a
+      constant TRUE, i.e. an unconditional cross join.
+    - *result_df* must have at most one row per join key. Otherwise the left join multiplies *df* rows.
+    - *result_columns* must be disjoint from *df.columns*. Otherwise the final select emits two columns
+      with the same name.
+
+    The helper aliases the two sides internally as "df" and "ref_df".
+
+    Args:
+        df: The input DataFrame whose rows and columns must be preserved.
+        result_df: The computed results, unique per combination of join key values.
+        join_columns: Non-empty list of column names shared by both DataFrames and used for matching.
+        result_columns: Non-overlapping result columns to append from *result_df*.
+
+    Returns:
+        The input rows and columns with the requested result columns appended.
+    """
+    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:
@@ -3675,7 +3719,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,
     aggr_matches_dataset,
 )
 from databricks.labs.dqx.utils import get_column_name_or_alias
@@ -32,6 +33,49 @@
 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"))
+    assert actual.count() == len(rows)
+    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(
         [
@@ -1260,6 +1304,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,24 +221,29 @@ def test_null_current_passes(spark: SparkSession):
 # ---------------------------------------------------------------------------
 
 
-def test_group_by_isolates_bands(spark: SparkSession):
+@pytest.mark.parametrize("spike_group", ["spike", None])
+@pytest.mark.parametrize("group_column", ["grp", "__dq_metric"])
+def test_group_by_isolates_bands(spark: SparkSession, spike_group: str | None, group_column: str):
     """
-    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,
+    including when its column name matches an internal column prefix.
     """
     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")
+    df = spark.createDataFrame(rows, "grp: string, event_date: date, metric: double").withColumnRenamed(
+        "grp", group_column
+    )
 
     condition, apply_fn = has_no_aggr_outliers(
         "metric",
@@ -247,12 +252,12 @@ def test_group_by_isolates_bands(spark: SparkSession):
         sigma=3.0,
         lookback_num_intervals=14,
         warmup_num_intervals=3,
-        group_by=["grp"],
+        group_by=[group_column],
     )
     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()]
+    stable_msgs = [r[-1] for r in result.filter(F.col(group_column) == "stable").collect()]
+    spike_msgs = [r[-1] for r in result.filter(F.col(group_column).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}"