diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 523d7d706..ed3c944f1 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/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 @@ -2162,9 +2170,10 @@ def aggr_matches_dataset( group_by: Optional list of column names or Column expressions in the checked DataFrame to compare the aggregate per group instead of dataset-wide. Only simple column expressions are supported, e.g. *F.col("region")*. A group present in the checked DataFrame but absent from the reference - is reported as a mismatch; groups present only in the reference are not surfaced. Note that - when *aggr_type* is a window-incompatible aggregate (e.g. *count_distinct*), the checked-side - grouping join is not null-safe, so a legitimately null group key may under-report. + is reported as a mismatch; groups present only in the reference are not surfaced. Group keys are + matched null-safely on both the checked and reference sides (including window-incompatible + aggregates such as *count_distinct*), so a legitimately null group key is compared like any + other rather than being dropped. ref_group_by: Optional list of group-by columns on the reference (upstream) side, matched to *group_by* by position. Defaults to *group_by* when omitted. Must have the same length as *group_by*. Requires *group_by* to be set. @@ -3159,6 +3168,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 +3720,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) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 98a30a12f..26898ae3f 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -3301,6 +3301,18 @@ def test_apply_checks_with_sql_query(ws, spark): "run_id": RUN_ID, "user_metadata": {}, }, + # null-safe merge on [b, c]: this row's key (b=NULL, c=3) now matches the query + # group (b=NULL, c=3) because NULL<=>NULL is True on b (and 3==3 on c) + { + "name": "multiple_key_check_violation", + "message": "multiple key check failed", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, ], ], [ @@ -3343,6 +3355,18 @@ def test_apply_checks_with_sql_query(ws, spark): "run_id": RUN_ID, "user_metadata": {}, }, + # null-safe merge on [b, c]: this row's key (b=NULL, c=4) now matches the query + # group (b=NULL, c=4) because NULL<=>NULL is True on b (and 4==4 on c) + { + "name": "multiple_key_check_violation", + "message": "multiple key check failed", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, ], ], [None, None, None, None, None], diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index 3fbf53f41..799c3f59b 100644 --- a/tests/integration/test_dataset_checks.py +++ b/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. @@ -2041,24 +2106,23 @@ def test_aggr_matches_dataset_count_distinct_group_by_mismatch(spark: SparkSessi def test_aggr_matches_dataset_count_distinct_group_by_null_key(spark: SparkSession): """count_distinct + group_by with a NULL group key. - count_distinct is a window-incompatible aggregate, so the grouped join is not null-safe (documented on - aggr_matches_dataset). A legitimately NULL group key therefore cannot be matched to the reference and is - surfaced with a NULL limit (mismatch) rather than being silently compared. This test pins that documented - behavior so a future change to the join is a conscious decision. + count_distinct is a window-incompatible aggregate that uses a two-stage groupBy join, but that join is + null-safe on both the checked and reference sides. A legitimately NULL group key is therefore matched to + the reference and compared like any other group rather than being dropped. This test pins that behavior so + a future change to the join is a conscious decision. """ - # NULL group: distinct b = {1, 2} (2); group "y": distinct b = {3} (1) + # NULL group: checked distinct b = {1, 2} (2), ref distinct b = {10, 20} (2) -> equal -> passes. + # Group "y": checked distinct b = {3} (1), ref distinct b = {30} (1) -> equal -> passes. test_df = spark.createDataFrame([[None, 1], [None, 2], ["y", 3]], SCHEMA) ref_df = spark.createDataFrame([[None, 10], [None, 20], ["y", 30]], SCHEMA) condition, apply_fn = aggr_matches_dataset("b", ref_df_name="ref_df", aggr_type="count_distinct", group_by=["a"]) checked = apply_fn(test_df, spark, {"ref_df": ref_df}).select("a", "b", condition.alias("cond")) - # Assert the flagging behavior (which rows are flagged) rather than the exact NULL-limit message text: - # the non-null group "y" matches (distinct 1 == 1) and passes; the NULL-key group cannot join - # null-safely for a window-incompatible aggregate, so it is surfaced (condition is non-null). + # Both groups match null-safely and their distinct counts are equal, so no row is flagged. results = {(row["a"], row["b"]): row["cond"] for row in checked.collect()} assert results[("y", 3)] is None - assert results[(None, 1)] is not None - assert results[(None, 2)] is not None + assert results[(None, 1)] is None + assert results[(None, 2)] is None def test_aggr_matches_dataset_group_by_null_key_mismatch(spark: SparkSession): diff --git a/tests/integration/test_has_no_aggr_outliers.py b/tests/integration/test_has_no_aggr_outliers.py index f7c4f5e7a..74a4ebb55 100644 --- a/tests/integration/test_has_no_aggr_outliers.py +++ b/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}"