From 81b8a989704fd8a18763633da063e0df906f39c9 Mon Sep 17 00:00:00 2001 From: Matt Faltyn Date: Tue, 21 Jul 2026 19:38:29 +0200 Subject: [PATCH 1/4] Fix null-safe result joins --- src/databricks/labs/dqx/check_funcs.py | 27 +++++--- tests/integration/test_dataset_checks.py | 64 +++++++++++++++++++ .../integration/test_has_no_aggr_outliers.py | 17 ++--- 3 files changed, 91 insertions(+), 17 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 5789fc591..5249a439e 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/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) diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index cac3edf0e..a94789a66 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, ) 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. diff --git a/tests/integration/test_has_no_aggr_outliers.py b/tests/integration/test_has_no_aggr_outliers.py index f7c4f5e7a..abcf58fef 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,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}" From 2abf5c2b9cd49a5643758a0e27e2578c4b33643d Mon Sep 17 00:00:00 2001 From: Matt Faltyn Date: Fri, 24 Jul 2026 19:45:42 +0200 Subject: [PATCH 2/4] Address null-safe join review feedback --- src/databricks/labs/dqx/check_funcs.py | 81 +++++++++++++------ tests/integration/test_dataset_checks.py | 1 + .../integration/test_has_no_aggr_outliers.py | 16 ++-- 3 files changed, 66 insertions(+), 32 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 7d7e50233..b75f7b19a 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1938,6 +1938,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: """ @@ -1964,69 +1967,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 = _join_results_on_null_safe_columns(current, stats, join_keys, ["__dq_mu", "__dq_sigma", "__dq_n"]) + 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)"), @@ -3102,7 +3115,23 @@ def _match_rows( 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.""" + """ + Left-join computed result columns while matching null join keys. + + The caller must ensure that *result_df* has at most one row per join key and + that *result_columns* are disjoint from *df.columns*. Otherwise the join can + multiply input rows or create duplicate output columns. 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: 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"), diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index 5a3a7e2b7..f92a40573 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -64,6 +64,7 @@ def test_sql_query_with_null_merge_columns( ) 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 = {} diff --git a/tests/integration/test_has_no_aggr_outliers.py b/tests/integration/test_has_no_aggr_outliers.py index abcf58fef..74a4ebb55 100644 --- a/tests/integration/test_has_no_aggr_outliers.py +++ b/tests/integration/test_has_no_aggr_outliers.py @@ -222,10 +222,12 @@ def test_null_current_passes(spark: SparkSession): @pytest.mark.parametrize("spike_group", ["spike", None]) -def test_group_by_isolates_bands(spark: SparkSession, spike_group: str | 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 a nullable spike - group (14 days at 10 +/- 1, today=50). Only the spike group should violate. + 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 @@ -239,7 +241,9 @@ def test_group_by_isolates_bands(spark: SparkSession, spike_group: str | None): rows.append(("stable", today, 10.0)) # within band 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", @@ -248,12 +252,12 @@ def test_group_by_isolates_bands(spark: SparkSession, spike_group: str | None): 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").eqNullSafe(F.lit(spike_group))).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}" From ae94da0138d771b64167e1f7212b216d3fdea17b Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 28 Jul 2026 12:51:12 +0200 Subject: [PATCH 3/4] Document _join_results_on_null_safe_columns preconditions Spell out the three caller responsibilities the helper relies on but does not validate: non-empty join_columns (an empty list reduces the null-safe join to a cross join), one row per join key in result_df (else df rows fan out), and result_columns disjoint from df.columns (else duplicate output columns). All current callers satisfy these; the note is for future reuse. Co-authored-by: Isaac --- src/databricks/labs/dqx/check_funcs.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index fffb86088..e7882db40 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -3173,15 +3173,21 @@ def _join_results_on_null_safe_columns( """ Left-join computed result columns while matching null join keys. - The caller must ensure that *result_df* has at most one row per join key and - that *result_columns* are disjoint from *df.columns*. Otherwise the join can - multiply input rows or create duplicate output columns. The helper aliases - the two sides internally as "df" and "ref_df". + 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: Column names shared by both DataFrames and used for matching. + 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: From 47a5bbbe9b6b398a0d7e1c38dfd015f13029f813 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 28 Jul 2026 17:56:08 +0200 Subject: [PATCH 4/4] Update stale test expectations for null-safe result joins Two integration tests still encoded the pre-null-safe join behavior: - test_apply_checks_with_sql_query: the multiple_key_check_violation check merges on [b, c]. With the old left join, rows whose b was NULL could not match their own query group (NULL == NULL is NULL, not True), so no violation attached. The null-safe join now matches them (NULL <=> NULL), so rows (2, NULL, 3) and (1, NULL, 4) correctly gain the violation. - test_aggr_matches_dataset_count_distinct_group_by_null_key and the aggr_matches_dataset docstring claimed count_distinct's grouped join is not null-safe. Both the checked-side two-stage groupBy join and the reference-side join are in fact null-safe, so a NULL group key is compared like any other. Updated the test to assert the NULL group matches and passes, and corrected the docstring. Verified both tests pass against a live workspace. Co-authored-by: Isaac --- src/databricks/labs/dqx/check_funcs.py | 7 ++++--- tests/integration/test_apply_checks.py | 24 ++++++++++++++++++++++++ tests/integration/test_dataset_checks.py | 19 +++++++++---------- 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index e7882db40..ed3c944f1 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -2170,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. 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 f92a40573..799c3f59b 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -2106,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):