diff --git a/pyproject.toml b/pyproject.toml index f5fbf0caf..51e0df906 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -468,8 +468,8 @@ bad-functions = ["map", "input"] # ignored-parents = # Maximum number of arguments for function / method. -max-args = 10 -max-positional-arguments = 10 +max-args = 11 +max-positional-arguments = 11 # Maximum number of attributes for a class (see R0902). max-attributes = 16 @@ -481,7 +481,7 @@ max-bool-expr = 5 max-branches = 20 # Maximum number of locals for function / method body. -max-locals = 21 +max-locals = 23 # Maximum number of parents for a class (see R0901). max-parents = 7 diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index be1d9695d..19f0f3fc4 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1267,6 +1267,8 @@ def compare_datasets( null_safe_row_matching: bool | None = True, null_safe_column_value_matching: bool | None = True, row_filter: str | None = None, + abs_tolerance: float | None = None, + rel_tolerance: float | None = None, ) -> tuple[Column, Callable]: """ Dataset-level check that compares two datasets and returns a condition for changed rows, @@ -1315,6 +1317,17 @@ def compare_datasets( If enabled, (NULL, NULL) column values are equal and matching. row_filter: Optional SQL expression to filter rows in the input DataFrame. Auto-injected from the check filter. + abs_tolerance: Values are considered equal if the absolute difference is less than or equal to the tolerance. This is applicable to numeric columns. + Example: abs(a - b) <= tolerance + With tolerance=0.01: + 2.001 and 2.0099 → equal (diff = 0.0089) + 2.001 and 2.02 → not equal (diff = 0.019) + rel_tolerance: Relative tolerance for numeric comparisons. Differences within this relative tolerance are ignored. Useful if numbers vary in scale. + Example: abs(a - b) <= rel_tolerance * max(abs(a), abs(b)) + With tolerance=0.01 (1%): + 100 vs 101 → equal (diff = 1, tolerance = 1) + 2.001 vs 2.0099 → equal + Returns: Tuple[Column, Callable]: @@ -1323,6 +1336,11 @@ def compare_datasets( """ _validate_ref_params(columns, ref_columns, ref_df_name, ref_table) + abs_tolerance = 0.0 if abs_tolerance is None else abs_tolerance + rel_tolerance = 0.0 if rel_tolerance is None else rel_tolerance + if abs_tolerance < 0 or rel_tolerance < 0: + raise ValueError("Absolute and/or relative tolerances if provided must be non-negative") + # convert all input columns to strings pk_column_names = get_columns_as_strings(columns, allow_simple_expressions_only=True) ref_pk_column_names = get_columns_as_strings(ref_columns, allow_simple_expressions_only=True) @@ -1369,7 +1387,9 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> df, ref_df, pk_column_names, ref_pk_column_names, check_missing_records, null_safe_row_matching ) results = _add_row_diffs(results, pk_column_names, ref_pk_column_names, row_missing_col, row_extra_col) - results = _add_column_diffs(results, compare_columns, columns_changed_col, null_safe_column_value_matching) + results = _add_column_diffs( + results, compare_columns, columns_changed_col, null_safe_column_value_matching, abs_tolerance, rel_tolerance + ) results = _add_compare_condition( results, condition_col, row_missing_col, row_extra_col, columns_changed_col, filter_col ) @@ -1576,11 +1596,51 @@ def _add_row_diffs( return df +def _add_numeric_tolerance_condition( + col_name: str, abs_tolerance: float, rel_tolerance: float, null_safe_column_value_matching: bool | None = None +) -> Column: + df_col = F.col(f"df.{col_name}") + ref_col = F.col(f"ref_df.{col_name}") + + # Handle NULL cases explicitly based on null_safe_column_value_matching + if null_safe_column_value_matching: + # NULL safety: (NULL, NULL) should be considered equal + both_null = df_col.isNull() & ref_col.isNull() + either_null = df_col.isNull() | ref_col.isNull() + + # For non-NULL values, apply tolerance logic + tolerance_match = _match_values_with_tolerance(df_col, ref_col, abs_tolerance, rel_tolerance) + + # Values are considered equal if: + # 1. Both are NULL (null safety), OR + # 2. Neither is NULL AND they're within tolerance + values_match = both_null | (~either_null & tolerance_match) + else: + # Null safety disabled: if either value is NULL, consider them matching + either_null = df_col.isNull() | ref_col.isNull() + + tolerance_match = _match_values_with_tolerance(df_col, ref_col, abs_tolerance, rel_tolerance) + + # Values are considered equal if: either is NULL OR both non-NULL and within tolerance + values_match = either_null | tolerance_match + + # Return True if values are NOT within tolerance (indicating a difference) + return ~values_match + + +def _match_values_with_tolerance(df_col: Column, ref_col: Column, abs_tolerance: float, rel_tolerance: float) -> Column: + abs_diff = F.abs(df_col - ref_col) + tolerance_val_relative = rel_tolerance * F.greatest(F.abs(df_col), F.abs(ref_col)) + return (abs_diff <= F.lit(abs_tolerance)) | (abs_diff <= tolerance_val_relative) + + def _add_column_diffs( df: DataFrame, compare_columns: list[str], columns_changed_col: str, null_safe_column_value_matching: bool | None = True, + abs_tolerance: float = 0.0, + rel_tolerance: float = 0.0, ) -> DataFrame: """ Adds a column to the DataFrame that contains a map of changed columns and their differences. @@ -1596,29 +1656,43 @@ def _add_column_diffs( null_safe_column_value_matching: If True, treats nulls as equal when matching column values. If enabled (NULL, NULL) column values are equal and matching. If False, uses a standard inequality comparison (`!=`), where (NULL, NULL) values are not considered equal. - + abs_tolerance: Absolute tolerance for numeric comparisons. Differences within this absolute tolerance are ignored. + Example: abs(a - b) <= abs_tolerance + rel_tolerance: Relative tolerance for numeric comparisons. Differences within this relative tolerance are ignored. + Example: abs(a - b) <= rel_tolerance * max(abs(a), abs(b)) Returns: A DataFrame with the added *columns_changed_col* containing the map of changed columns and differences. """ + columns_changed = [] if compare_columns: - columns_changed = [ - F.when( - # with null-safe comparison values are matching if they are equal or both are NULL - ( - ~F.col(f"df.{col}").eqNullSafe(F.col(f"ref_df.{col}")) + + for col_name in compare_columns: + is_numeric = isinstance(df.schema[col_name].dataType, types.NumericType) + + if (abs_tolerance > 0.0 or rel_tolerance > 0.0) and is_numeric: + # Absolute and relative difference + condition = _add_numeric_tolerance_condition( + col_name, abs_tolerance, rel_tolerance, null_safe_column_value_matching + ) + else: + condition = ( + ~F.col(f"df.{col_name}").eqNullSafe(F.col(f"ref_df.{col_name}")) if null_safe_column_value_matching - else F.col(f"df.{col}") != F.col(f"ref_df.{col}") - ), - F.struct( - F.lit(col).alias("col_changed"), + else F.col(f"df.{col_name}") != F.col(f"ref_df.{col_name}") + ) + + columns_changed.append( + F.when( + condition, F.struct( - F.col(f"df.{col}").cast("string").alias("df"), - F.col(f"ref_df.{col}").cast("string").alias("ref"), - ).alias("diff"), - ), - ).otherwise(None) - for col in compare_columns - ] + F.lit(col_name).alias("col_changed"), + F.struct( + F.col(f"df.{col_name}").cast("string").alias("df"), + F.col(f"ref_df.{col_name}").cast("string").alias("ref"), + ).alias("diff"), + ), + ).otherwise(None) + ) df = df.withColumn(columns_changed_col, F.array_compact(F.array(*columns_changed))) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 66ca8b827..783cca327 100644 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -1116,6 +1116,197 @@ def test_apply_is_unique(ws, spark): assert_df_equality(checked, expected, ignore_nullable=True) +def test_compare_datasets_with_tolerance(ws, spark): + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + + schema = "id int, value double" + # Source DataFrame: has values near, just within, and just outside tolerances + src_df = spark.createDataFrame( + [ + [1, 100.00], # equal under zero tolerance + [2, 100.99], # equal under abs_tolerance=1 (diff = 0.99) + [3, 101.01], # not equal under abs_tolerance=1 (diff = 1.01) + [4, 202.0], # equal under rel_tolerance=0.01 (diff = 2, tolerance = 2.02) + [5, 204.5], # not equal under rel_tolerance=0.01 (diff = 4.5, tolerance = 2.0) + [6, None], # Null comparison + [7, None], # Null comparison + ], + schema, + ) + + # Reference DataFrame + ref_df = spark.createDataFrame( + [ + [1, 100.00], + [2, 100.00], + [3, 100.0], + [4, 200.0], + [5, 200.0], + [6, 100.00], + [7, None], + ], + schema, + ) + + pk_columns = ["id"] + + # Add check with both tolerances + checks = [ + DQDatasetRule( + name="id_compare_with_tolerance", + criticality="error", + check_func=check_funcs.compare_datasets, + columns=pk_columns, + check_func_kwargs={ + "ref_columns": pk_columns, + "ref_df_name": "ref_df", + "abs_tolerance": 1.0, # absolute tolerance of 1 + "rel_tolerance": 0.01, # relative tolerance of 1% + "null_safe_column_value_matching": True, + }, + user_metadata={"test": "tolerance"}, + ), + ] + + refs_df = {"ref_df": ref_df} + + checked = dq_engine.apply_checks(src_df, checks, refs_df) + + # Build expected results: rows only get flagged when outside of both tolerances + expected = spark.createDataFrame( + [ + [1, 100.00, None, None], # exact match, no error/warning + [2, 100.99, None, None], # diff = 0.99 <= abs_tolerance=1.0, so no error + [3, 101.01, None, None], # diff = 1.01 <= (1.0 + 0.01*100 = 2.0), so no error], + [4, 202.00, None, None], # diff = 2.0, rel_tolerance = 2.02, so within relative tolerance + [ + 5, + 204.50, + [ + { + "name": "id_compare_with_tolerance", + "message": '{"row_missing":false,"row_extra":false,"changed":{"value":{"df":"204.5","ref":"200.0"}}}', + "columns": pk_columns, + "filter": None, + "function": "compare_datasets", + "run_time": RUN_TIME, + "user_metadata": {"test": "tolerance"}, + } + ], + None, + ], + [ + 6, + None, + [ + { + "name": "id_compare_with_tolerance", + "message": '{"row_missing":false,"row_extra":false,"changed":{"value":{"ref":"100.0"}}}', + "columns": pk_columns, + "filter": None, + "function": "compare_datasets", + "run_time": RUN_TIME, + "user_metadata": {"test": "tolerance"}, + } + ], + None, + ], + [7, None, None, None], + ], + schema + REPORTING_COLUMNS, + ) + + assert_df_equality(checked.sort(pk_columns), expected.sort(pk_columns), ignore_nullable=True) + + +def test_compare_datasets_with_tolerance_with_disabled_null_safe_column_value_matching(ws, spark): + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + + schema = "id int, value double" + # Source DataFrame: has values near, just within, and just outside tolerances + src_df = spark.createDataFrame( + [ + [1, 100.00], # equal under zero tolerance + [2, 100.99], # equal under abs_tolerance=1 (diff = 0.99) + [3, 101.01], # not equal under abs_tolerance=1 (diff = 1.01) + [4, 202.0], # equal under rel_tolerance=0.01 (diff = 2, tolerance = 2.02) + [5, 204.5], # not equal under rel_tolerance=0.01 (diff = 4.5, tolerance = 2.0) + [6, None], # Null comparison + [7, None], # Null comparison + ], + schema, + ) + + # Reference DataFrame + ref_df = spark.createDataFrame( + [ + [1, 100.00], + [2, 100.00], + [3, 100.0], + [4, 200.0], + [5, 200.0], + [6, 100.00], + [7, None], + ], + schema, + ) + + pk_columns = ["id"] + + # Add check with both tolerances + checks = [ + DQDatasetRule( + name="id_compare_with_tolerance", + criticality="error", + check_func=check_funcs.compare_datasets, + columns=pk_columns, + check_func_kwargs={ + "ref_columns": pk_columns, + "ref_df_name": "ref_df", + "abs_tolerance": 1.0, # absolute tolerance of 1 + "rel_tolerance": 0.01, # relative tolerance of 1% + "null_safe_column_value_matching": False, + }, + user_metadata={"test": "tolerance"}, + ), + ] + + refs_df = {"ref_df": ref_df} + + checked = dq_engine.apply_checks(src_df, checks, refs_df) + + # Build expected results: rows only get flagged when outside of both tolerances + expected = spark.createDataFrame( + [ + [1, 100.00, None, None], # exact match, no error/warning + [2, 100.99, None, None], # diff = 0.99 <= abs_tolerance=1.0, so no error + [3, 101.01, None, None], # diff = 1.01 <= (1.0 + 0.01*100 = 2.0), so no error], + [4, 202.00, None, None], # diff = 2.0, rel_tolerance = 2.02, so within relative tolerance + [ + 5, + 204.50, + [ + { + "name": "id_compare_with_tolerance", + "message": '{"row_missing":false,"row_extra":false,"changed":{"value":{"df":"204.5","ref":"200.0"}}}', + "columns": pk_columns, + "filter": None, + "function": "compare_datasets", + "run_time": RUN_TIME, + "user_metadata": {"test": "tolerance"}, + } + ], + None, + ], + [6, None, None, None], # Nulls, should be considered equal if null_safe is disabled + [7, None, None, None], + ], + schema + REPORTING_COLUMNS, + ) + + assert_df_equality(checked.sort(pk_columns), expected.sort(pk_columns), ignore_nullable=True) + + def test_apply_checks(ws, spark): dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA) diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index 02433634b..0191fe555 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -1178,6 +1178,8 @@ def test_compare_dataset_disabled_null_safe_row_matching(spark: SparkSession): df = spark.createDataFrame( [ [1, 1, None], + [2, 2, 2], + [3, 3, None], [1, None, "val1"], ], schema, @@ -1186,6 +1188,8 @@ def test_compare_dataset_disabled_null_safe_row_matching(spark: SparkSession): df_ref = spark.createDataFrame( [ [1, 1, None], + [2, 2, None], + [3, 3, 3], [1, None, "val2"], ], schema, @@ -1199,6 +1203,7 @@ def test_compare_dataset_disabled_null_safe_row_matching(spark: SparkSession): ref_df_name="df_ref", check_missing_records=True, null_safe_row_matching=False, + null_safe_column_value_matching=True, ) actual: DataFrame = apply(df, spark, {"df_ref": df_ref}) @@ -1222,6 +1227,32 @@ def test_compare_dataset_disabled_null_safe_row_matching(spark: SparkSession): separators=(',', ':'), ), }, + { + "id1": 2, + "id2": 2, + "name": 2, + compare_status_column: json.dumps( + { + "row_missing": False, + "row_extra": False, + "changed": {"name": {"df": "2"}}, + }, + separators=(',', ':'), + ), + }, + { + "id1": 3, + "id2": 3, + "name": None, + compare_status_column: json.dumps( + { + "row_missing": False, + "row_extra": False, + "changed": {"name": {"ref": "3"}}, + }, + separators=(',', ':'), + ), + }, { "id1": 1, "id2": None, @@ -1250,6 +1281,7 @@ def test_compare_dataset_disabled_null_safe_column_value_matching(spark: SparkSe [ [1, "val1"], [2, "val2"], + [3, None], ], schema, ) @@ -1258,6 +1290,7 @@ def test_compare_dataset_disabled_null_safe_column_value_matching(spark: SparkSe [ [1, None], # should not show any diff in the name [2, "val2"], + [3, "val3"], ], schema, ) @@ -1290,6 +1323,11 @@ def test_compare_dataset_disabled_null_safe_column_value_matching(spark: SparkSe "name": "val2", compare_status_column: None, }, + { + "id": 3, + "name": None, + compare_status_column: None, + }, ], expected_schema, ) diff --git a/tests/unit/test_dataset_checks.py b/tests/unit/test_dataset_checks.py index 4c3688b68..de3e888a9 100644 --- a/tests/unit/test_dataset_checks.py +++ b/tests/unit/test_dataset_checks.py @@ -65,6 +65,29 @@ def test_compare_datasets_exceptions(ref_df_name, ref_table, ref_columns, column ) +@pytest.mark.parametrize( + "abs_tolerance, rel_tolerance", + [ + (-1, None), + (None, -1), + (-1, -1), + ], +) +def test_compare_datasets_invalid_tolerance_exceptions(abs_tolerance, rel_tolerance): + with pytest.raises(ValueError, match="Absolute and/or relative tolerances if provided must be non-negative"): + DQDatasetRule( + criticality="warn", + check_func=check_funcs.compare_datasets, + columns=["col1"], + check_func_kwargs={ + "ref_columns": ["col1"], + "ref_table": "ref_table", + "abs_tolerance": abs_tolerance, + "rel_tolerance": rel_tolerance, + }, + ) + + def test_sql_query_missing_merge_columns(): with pytest.raises(ValueError, match="merge_columns must contain at least one column"): DQDatasetRule(