From 63840aad25b84b7f24847dc2e86f1b0c27661b52 Mon Sep 17 00:00:00 2001 From: Jgprog117 <93980271+Jgprog117@users.noreply.github.com> Date: Sat, 31 Jan 2026 15:19:22 +0100 Subject: [PATCH 01/12] Add Decimal support to min_max generator Fixes #1013 The dq_generate_min_max method now supports Python's Decimal type in addition to int and float for min/max validation checks. Changes: - Import Decimal from decimal module - Update _is_num() to include Decimal in isinstance check - Add comprehensive unit tests for Decimal, int, and float types This enables proper data quality checks for decimal-precise financial and scientific data where floating-point precision issues would cause false positives. Signed-off-by: Jgprog117 --- src/databricks/labs/dqx/profiler/generator.py | 3 +- tests/unit/test_generator_numeric.py | 61 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_generator_numeric.py diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index e26e1a139..643ade394 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -3,6 +3,7 @@ import logging import datetime import json +from decimal import Decimal from collections.abc import Callable from pyspark.sql import SparkSession from databricks.sdk import WorkspaceClient @@ -253,7 +254,7 @@ def dq_generate_min_max(column: str, criticality: str = "error", **params: dict) return None def _is_num(value): - return isinstance(value, (int, float)) + return isinstance(value, (int, float, Decimal)) def _is_temporal(value): return isinstance(value, (datetime.date, datetime.datetime)) diff --git a/tests/unit/test_generator_numeric.py b/tests/unit/test_generator_numeric.py new file mode 100644 index 000000000..01b0453aa --- /dev/null +++ b/tests/unit/test_generator_numeric.py @@ -0,0 +1,61 @@ +from decimal import Decimal + +from databricks.labs.dqx.profiler.generator import DQGenerator + + +def test_decimal_both_bounds_is_in_range(): + """Test min_max generator with Decimal type for both bounds.""" + result = DQGenerator.dq_generate_min_max( + "price_col", **{"min": Decimal("0.01"), "max": Decimal("999.99")} + ) + assert result["check"]["function"] == "is_in_range" + args = result["check"]["arguments"] + assert args["column"] == "price_col" + assert args["min_limit"] == Decimal("0.01") + assert args["max_limit"] == Decimal("999.99") + + +def test_decimal_only_min_is_not_less_than(): + """Test min_max generator with Decimal type for minimum bound only.""" + result = DQGenerator.dq_generate_min_max("amount_col", **{"min": Decimal("10.50"), "max": None}) + assert result["check"]["function"] == "is_not_less_than" + args = result["check"]["arguments"] + assert args["column"] == "amount_col" + assert args["limit"] == Decimal("10.50") + + +def test_decimal_only_max_is_not_greater_than(): + """Test min_max generator with Decimal type for maximum bound only.""" + result = DQGenerator.dq_generate_min_max("total_col", **{"min": None, "max": Decimal("1000.00")}) + assert result["check"]["function"] == "is_not_greater_than" + args = result["check"]["arguments"] + assert args["column"] == "total_col" + assert args["limit"] == Decimal("1000.00") + + +def test_int_both_bounds_is_in_range(): + """Test min_max generator with int type for both bounds.""" + result = DQGenerator.dq_generate_min_max("age_col", **{"min": 0, "max": 120}) + assert result["check"]["function"] == "is_in_range" + args = result["check"]["arguments"] + assert args["column"] == "age_col" + assert args["min_limit"] == 0 + assert args["max_limit"] == 120 + + +def test_float_both_bounds_is_in_range(): + """Test min_max generator with float type for both bounds.""" + result = DQGenerator.dq_generate_min_max("temperature_col", **{"min": -273.15, "max": 1000.0}) + assert result["check"]["function"] == "is_in_range" + args = result["check"]["arguments"] + assert args["column"] == "temperature_col" + assert args["min_limit"] == -273.15 + assert args["max_limit"] == 1000.0 + + +def test_mixed_numeric_types_not_generated(): + """Test that mixing int/float with Decimal in the same rule returns None.""" + # int min, Decimal max - should not generate a rule due to type mismatch + result = DQGenerator.dq_generate_min_max("mixed_col", **{"min": 10, "max": Decimal("100.00")}) + # The _same_family check should prevent this combination + assert result is None or result["check"]["function"] in ["is_not_greater_than"] From 34639972c9b2600d861ea36c3b8861d69c3d57b4 Mon Sep 17 00:00:00 2001 From: Jgprog117 <93980271+Jgprog117@users.noreply.github.com> Date: Sat, 31 Jan 2026 15:33:25 +0100 Subject: [PATCH 02/12] Trigger CI checks From 0d63ce87d04b71d57d6346352d06f8c557e384b9 Mon Sep 17 00:00:00 2001 From: Jgprog117 <93980271+Jgprog117@users.noreply.github.com> Date: Sun, 1 Feb 2026 17:29:43 +0100 Subject: [PATCH 03/12] Fix Decimal unit tests and add integration test for min_max generator (#1013) --- tests/integration/test_generator.py | 23 +++++++++++++++++++++++ tests/unit/test_generator_numeric.py | 21 ++++++++++++++------- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/tests/integration/test_generator.py b/tests/integration/test_generator.py index 3b20995a6..95bf16f95 100644 --- a/tests/integration/test_generator.py +++ b/tests/integration/test_generator.py @@ -1,5 +1,6 @@ import logging import datetime +from decimal import Decimal from databricks.labs.dqx.profiler.generator import DQGenerator from databricks.labs.dqx.profiler.profiler import DQProfile @@ -35,6 +36,12 @@ description='Real min/max values were used', parameters={'max': 333323.00, 'min': 1.23}, ), + DQProfile( + name='min_max', + column='price', + description='Real min/max values were used', + parameters={'min': Decimal("0.01"), 'max': Decimal("999.99")}, + ), ] @@ -95,6 +102,14 @@ def test_generate_dq_rules(ws, spark): "name": "d1_isnt_in_range", "criticality": "error", }, + { + "check": { + "function": "is_in_range", + "arguments": {"column": "price", "min_limit": Decimal("0.01"), "max_limit": Decimal("999.99")}, + }, + "name": "price_isnt_in_range", + "criticality": "error", + }, ] assert expectations == expected, f"Actual expectations: {expectations}" @@ -156,6 +171,14 @@ def test_generate_dq_rules_warn(ws, spark): "name": "d1_isnt_in_range", "criticality": "warn", }, + { + "check": { + "function": "is_in_range", + "arguments": {"column": "price", "min_limit": Decimal("0.01"), "max_limit": Decimal("999.99")}, + }, + "name": "price_isnt_in_range", + "criticality": "warn", + }, ] assert expectations == expected, f"Actual expectations: {expectations}" diff --git a/tests/unit/test_generator_numeric.py b/tests/unit/test_generator_numeric.py index 01b0453aa..52d886222 100644 --- a/tests/unit/test_generator_numeric.py +++ b/tests/unit/test_generator_numeric.py @@ -17,7 +17,9 @@ def test_decimal_both_bounds_is_in_range(): def test_decimal_only_min_is_not_less_than(): """Test min_max generator with Decimal type for minimum bound only.""" - result = DQGenerator.dq_generate_min_max("amount_col", **{"min": Decimal("10.50"), "max": None}) + result = DQGenerator.dq_generate_min_max( + "amount_col", **{"min": Decimal("10.50"), "max": None} + ) assert result["check"]["function"] == "is_not_less_than" args = result["check"]["arguments"] assert args["column"] == "amount_col" @@ -53,9 +55,14 @@ def test_float_both_bounds_is_in_range(): assert args["max_limit"] == 1000.0 -def test_mixed_numeric_types_not_generated(): - """Test that mixing int/float with Decimal in the same rule returns None.""" - # int min, Decimal max - should not generate a rule due to type mismatch - result = DQGenerator.dq_generate_min_max("mixed_col", **{"min": 10, "max": Decimal("100.00")}) - # The _same_family check should prevent this combination - assert result is None or result["check"]["function"] in ["is_not_greater_than"] +def test_mixed_int_and_decimal_is_in_range(): + """Test that mixing int and Decimal produces is_in_range since both are numeric.""" + result = DQGenerator.dq_generate_min_max( + "mixed_col", **{"min": 10, "max": Decimal("100.00")} + ) + assert result is not None + assert result["check"]["function"] == "is_in_range" + args = result["check"]["arguments"] + assert args["column"] == "mixed_col" + assert args["min_limit"] == 10 + assert args["max_limit"] == Decimal("100.00") From 66a05a883b2ace8fbb0f067083dea29df910f2ce Mon Sep 17 00:00:00 2001 From: Jgprog117 <93980271+Jgprog117@users.noreply.github.com> Date: Wed, 4 Feb 2026 11:07:49 +0100 Subject: [PATCH 04/12] Add support for Decimal to checks and change to single-line function calls --- src/databricks/labs/dqx/check_funcs.py | 32 +++++++++++++------------- tests/integration/test_generator.py | 28 +++++++++++----------- tests/unit/test_generator_numeric.py | 12 +++------- 3 files changed, 33 insertions(+), 39 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index ed1eb9654..8f66e044e 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -3,6 +3,7 @@ import warnings import ipaddress import uuid +from decimal import Decimal from collections.abc import Callable, Sequence from enum import Enum from itertools import zip_longest @@ -609,7 +610,7 @@ def is_not_in_near_future(column: str | Column, offset: int = 0, curr_timestamp: @register_rule("row") def is_equal_to( column: str | Column, - value: int | float | str | datetime.date | datetime.datetime | Column | None = None, + value: int | float | Decimal | str | datetime.date | datetime.datetime | Column | None = None, abs_tolerance: float | None = None, rel_tolerance: float | None = None, ) -> Column: @@ -671,7 +672,7 @@ def is_equal_to( @register_rule("row") def is_not_equal_to( column: str | Column, - value: int | float | str | datetime.date | datetime.datetime | Column | None = None, + value: int | float | Decimal | str | datetime.date | datetime.datetime | Column | None = None, abs_tolerance: float | None = None, rel_tolerance: float | None = None, ) -> Column: @@ -733,7 +734,7 @@ def is_not_equal_to( @register_rule("row") def is_not_less_than( - column: str | Column, limit: int | float | datetime.date | datetime.datetime | str | Column | None = None + column: str | Column, limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None ) -> Column: """Checks whether the values in the input column are not less than the provided limit. @@ -763,7 +764,7 @@ def is_not_less_than( @register_rule("row") def is_not_greater_than( - column: str | Column, limit: int | float | datetime.date | datetime.datetime | str | Column | None = None + column: str | Column, limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None ) -> Column: """Checks whether the values in the input column are not greater than the provided limit. @@ -794,8 +795,8 @@ def is_not_greater_than( @register_rule("row") def is_in_range( column: str | Column, - min_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None, - max_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None, + min_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None, + max_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None, ) -> Column: """Checks whether the values in the input column are in the provided limits (inclusive of both boundaries). @@ -832,8 +833,8 @@ def is_in_range( @register_rule("row") def is_not_in_range( column: str | Column, - min_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None, - max_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None, + min_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None, + max_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None, ) -> Column: """Checks whether the values in the input column are outside the provided limits (inclusive of both boundaries). @@ -1585,7 +1586,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> @register_rule("dataset") def is_aggr_not_greater_than( column: str | Column, - limit: int | float | str | Column, + limit: int | float | Decimal | str | Column, aggr_type: str = "count", group_by: list[str | Column] | None = None, row_filter: str | None = None, @@ -1630,7 +1631,7 @@ def is_aggr_not_greater_than( @register_rule("dataset") def is_aggr_not_less_than( column: str | Column, - limit: int | float | str | Column, + limit: int | float | Decimal | str | Column, aggr_type: str = "count", group_by: list[str | Column] | None = None, row_filter: str | None = None, @@ -1675,7 +1676,7 @@ def is_aggr_not_less_than( @register_rule("dataset") def is_aggr_equal( column: str | Column, - limit: int | float | str | Column, + limit: int | float | Decimal | str | Column, aggr_type: str = "count", group_by: list[str | Column] | None = None, row_filter: str | None = None, @@ -1726,7 +1727,7 @@ def is_aggr_equal( @register_rule("dataset") def is_aggr_not_equal( column: str | Column, - limit: int | float | str | Column, + limit: int | float | Decimal | str | Column, aggr_type: str = "count", group_by: list[str | Column] | None = None, row_filter: str | None = None, @@ -2780,7 +2781,6 @@ def _add_column_diffs( """ columns_changed = [] if compare_columns: - for col_name in compare_columns: is_numeric = isinstance(df.schema[col_name].dataType, types.NumericType) @@ -2991,7 +2991,7 @@ def _build_aggregate_check_metadata( def _is_aggr_compare( column: str | Column, - limit: int | float | str | Column, + limit: int | float | Decimal | str | Column, aggr_type: str, aggr_params: dict[str, Any] | None, group_by: list[str | Column] | None, @@ -3218,7 +3218,7 @@ def _cleanup_alias_name(column: str) -> str: def get_limit_expr( - limit: int | float | datetime.date | datetime.datetime | str | Column | None = None, + limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None, ) -> Column: """ Generate a Spark Column expression for a limit value. @@ -3247,7 +3247,7 @@ def get_limit_expr( parsed_dt = datetime.datetime.fromisoformat(limit) # Check if the string contains time component - has_time = ':' in limit + has_time = ":" in limit if has_time: return F.to_timestamp(F.lit(parsed_dt)) diff --git a/tests/integration/test_generator.py b/tests/integration/test_generator.py index 95bf16f95..947a004d1 100644 --- a/tests/integration/test_generator.py +++ b/tests/integration/test_generator.py @@ -31,16 +31,16 @@ ), DQProfile(name="is_random", column="vendor_id", parameters={"in": ["1", "4", "2"]}), DQProfile( - name='min_max', - column='d1', - description='Real min/max values were used', - parameters={'max': 333323.00, 'min': 1.23}, + name="min_max", + column="d1", + description="Real min/max values were used", + parameters={"max": 333323.00, "min": 1.23}, ), DQProfile( - name='min_max', - column='price', - description='Real min/max values were used', - parameters={'min': Decimal("0.01"), 'max': Decimal("999.99")}, + name="min_max", + column="price", + description="Real min/max values were used", + parameters={"min": Decimal("0.01"), "max": Decimal("999.99")}, ), ] @@ -332,9 +332,9 @@ def test_generate_is_unique_dq_rule(ws, spark): generator = DQGenerator(ws, spark) test_is_unique_rules = [ DQProfile( - name='is_unique', - column='col1,col2', - description='LLM-detected primary key columns: col1, col2', + name="is_unique", + column="col1,col2", + description="LLM-detected primary key columns: col1, col2", parameters={"nulls_distinct": False, "confidence": "high"}, ), ] @@ -355,9 +355,9 @@ def test_generate_is_unique_dq_rule_default_criticality(ws, spark): generator = DQGenerator(ws, spark) test_is_unique_rules = [ DQProfile( - name='is_unique', - column='col1', - description='LLM-detected primary key columns: col1, col2', + name="is_unique", + column="col1", + description="LLM-detected primary key columns: col1, col2", parameters={"nulls_distinct": True, "confidence": "low"}, ), ] diff --git a/tests/unit/test_generator_numeric.py b/tests/unit/test_generator_numeric.py index 52d886222..0d2b402e7 100644 --- a/tests/unit/test_generator_numeric.py +++ b/tests/unit/test_generator_numeric.py @@ -5,9 +5,7 @@ def test_decimal_both_bounds_is_in_range(): """Test min_max generator with Decimal type for both bounds.""" - result = DQGenerator.dq_generate_min_max( - "price_col", **{"min": Decimal("0.01"), "max": Decimal("999.99")} - ) + result = DQGenerator.dq_generate_min_max("price_col", **{"min": Decimal("0.01"), "max": Decimal("999.99")}) assert result["check"]["function"] == "is_in_range" args = result["check"]["arguments"] assert args["column"] == "price_col" @@ -17,9 +15,7 @@ def test_decimal_both_bounds_is_in_range(): def test_decimal_only_min_is_not_less_than(): """Test min_max generator with Decimal type for minimum bound only.""" - result = DQGenerator.dq_generate_min_max( - "amount_col", **{"min": Decimal("10.50"), "max": None} - ) + result = DQGenerator.dq_generate_min_max("amount_col", **{"min": Decimal("10.50"), "max": None}) assert result["check"]["function"] == "is_not_less_than" args = result["check"]["arguments"] assert args["column"] == "amount_col" @@ -57,9 +53,7 @@ def test_float_both_bounds_is_in_range(): def test_mixed_int_and_decimal_is_in_range(): """Test that mixing int and Decimal produces is_in_range since both are numeric.""" - result = DQGenerator.dq_generate_min_max( - "mixed_col", **{"min": 10, "max": Decimal("100.00")} - ) + result = DQGenerator.dq_generate_min_max("mixed_col", **{"min": 10, "max": Decimal("100.00")}) assert result is not None assert result["check"]["function"] == "is_in_range" args = result["check"]["arguments"] From 2f996ac3f40bf1a907c59bdefabfe9a346ab5a52 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 16:10:26 +0100 Subject: [PATCH 05/12] added tests to cover Decimal handling, updated dlt tests, added support for decimal when tolerance is used --- src/databricks/labs/dqx/check_funcs.py | 4 +- tests/integration/test_dataset_checks.py | 44 +++++-- tests/integration/test_dlt_rules_generator.py | 10 +- tests/integration/test_row_checks.py | 107 ++++++++++++------ 4 files changed, 113 insertions(+), 52 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index c94b17bf0..aecb3f286 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -648,7 +648,7 @@ def is_equal_to( col_str_norm, col_expr_str, col_expr = get_normalized_column_and_expr(column) value_expr = get_limit_expr(value) - if (abs_tolerance > 0.0 or rel_tolerance > 0.0) and isinstance(value, (int, float)): + if (abs_tolerance > 0.0 or rel_tolerance > 0.0) and isinstance(value, (int, float, Decimal)): # Use tolerance-based comparison for numeric columns tolerance_match = _match_values_with_tolerance(col_expr, value_expr, abs_tolerance, rel_tolerance) condition = ~tolerance_match @@ -711,7 +711,7 @@ def is_not_equal_to( col_str_norm, col_expr_str, col_expr = get_normalized_column_and_expr(column) value_expr = get_limit_expr(value) - if (abs_tolerance > 0.0 or rel_tolerance > 0.0) and isinstance(value, (int, float)): + if (abs_tolerance > 0.0 or rel_tolerance > 0.0) and isinstance(value, (int, float, Decimal)): # Use tolerance-based comparison for numeric columns tolerance_match = _match_values_with_tolerance(col_expr, value_expr, abs_tolerance, rel_tolerance) condition = tolerance_match diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index bb184803b..efcd2e4c8 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -472,9 +472,10 @@ def test_is_aggr_not_greater_than(spark: SparkSession): is_aggr_not_greater_than("a", limit=F.lit(0), aggr_type="count", row_filter="b is not null", group_by=["a"]), is_aggr_not_greater_than(F.col("b"), limit=F.lit(0), aggr_type="count", group_by=[F.col("b")]), is_aggr_not_greater_than("b", limit=0.0, aggr_type="avg"), - is_aggr_not_greater_than("b", limit=0.0, aggr_type="sum"), + is_aggr_not_greater_than("b", limit=Decimal("0.0"), aggr_type="sum"), is_aggr_not_greater_than("b", limit=0.0, aggr_type="min"), - is_aggr_not_greater_than("b", limit=0.0, aggr_type="max"), + is_aggr_not_greater_than("b", limit=Decimal("0.0"), aggr_type="max"), + is_aggr_not_greater_than("b", limit="2.0", aggr_type="avg"), ] actual = _apply_checks(test_df, checks) @@ -487,7 +488,8 @@ def test_is_aggr_not_greater_than(spark: SparkSession): "b_avg_greater_than_limit STRING, " "b_sum_greater_than_limit STRING, " "b_min_greater_than_limit STRING, " - "b_max_greater_than_limit STRING" + "b_max_greater_than_limit STRING, " + "b_avg_greater_than_limit STRING" ) expected = spark.createDataFrame( @@ -504,6 +506,7 @@ def test_is_aggr_not_greater_than(spark: SparkSession): "Sum value 4 in column 'b' is greater than limit: 0.0", "Min value 1 in column 'b' is greater than limit: 0.0", "Max value 3 in column 'b' is greater than limit: 0.0", + None, ], [ "a", @@ -516,6 +519,7 @@ def test_is_aggr_not_greater_than(spark: SparkSession): "Sum value 4 in column 'b' is greater than limit: 0.0", "Min value 1 in column 'b' is greater than limit: 0.0", "Max value 3 in column 'b' is greater than limit: 0.0", + None, ], [ "b", @@ -528,6 +532,7 @@ def test_is_aggr_not_greater_than(spark: SparkSession): "Sum value 4 in column 'b' is greater than limit: 0.0", "Min value 1 in column 'b' is greater than limit: 0.0", "Max value 3 in column 'b' is greater than limit: 0.0", + None, ], ], expected_schema, @@ -554,9 +559,10 @@ def test_is_aggr_not_less_than(spark: SparkSession): ), is_aggr_not_less_than(F.col("b"), limit=F.lit(2), aggr_type="count", group_by=["b"]), is_aggr_not_less_than("b", limit=3.0, aggr_type="avg"), - is_aggr_not_less_than("b", limit=5.0, aggr_type="sum"), + is_aggr_not_less_than("b", limit=Decimal("5.0"), aggr_type="sum"), is_aggr_not_less_than("b", limit=2.0, aggr_type="min"), - is_aggr_not_less_than("b", limit=4.0, aggr_type="max"), + is_aggr_not_less_than("b", limit=Decimal("4.0"), aggr_type="max"), + is_aggr_not_less_than("b", limit="1.0", aggr_type="avg"), ] actual = _apply_checks(test_df, checks) @@ -568,7 +574,8 @@ def test_is_aggr_not_less_than(spark: SparkSession): "b_avg_less_than_limit STRING, " "b_sum_less_than_limit STRING, " "b_min_less_than_limit STRING, " - "b_max_less_than_limit STRING" + "b_max_less_than_limit STRING, " + "b_avg_less_than_limit STRING" ) expected = spark.createDataFrame( @@ -584,6 +591,7 @@ def test_is_aggr_not_less_than(spark: SparkSession): "Sum value 4 in column 'b' is less than limit: 5.0", "Min value 1 in column 'b' is less than limit: 2.0", "Max value 3 in column 'b' is less than limit: 4.0", + None, ], [ "a", @@ -596,6 +604,7 @@ def test_is_aggr_not_less_than(spark: SparkSession): "Sum value 4 in column 'b' is less than limit: 5.0", "Min value 1 in column 'b' is less than limit: 2.0", "Max value 3 in column 'b' is less than limit: 4.0", + None, ], [ "b", @@ -608,6 +617,7 @@ def test_is_aggr_not_less_than(spark: SparkSession): "Sum value 4 in column 'b' is less than limit: 5.0", "Min value 1 in column 'b' is less than limit: 2.0", "Max value 3 in column 'b' is less than limit: 4.0", + None, ], ], expected_schema, @@ -655,9 +665,10 @@ def test_is_aggr_equal(spark: SparkSession): is_aggr_equal("a", limit=F.lit(1), aggr_type="count", row_filter="b is not null", group_by=["a"]), is_aggr_equal(F.col("b"), limit=F.lit(2), aggr_type="count", group_by=[F.col("b")]), is_aggr_equal("b", limit=2.0, aggr_type="avg"), - is_aggr_equal("b", limit=10.0, aggr_type="sum"), + is_aggr_equal("b", limit=Decimal("10.0"), aggr_type="sum"), is_aggr_equal("b", limit=1.0, aggr_type="min"), - is_aggr_equal("b", limit=5.0, aggr_type="max"), + is_aggr_equal("b", limit=Decimal("5.0"), aggr_type="max"), + is_aggr_equal("b", limit="2.0", aggr_type="avg"), ] actual = _apply_checks(test_df, checks) @@ -670,7 +681,8 @@ def test_is_aggr_equal(spark: SparkSession): "b_avg_not_equal_to_limit STRING, " "b_sum_not_equal_to_limit STRING, " "b_min_not_equal_to_limit STRING, " - "b_max_not_equal_to_limit STRING" + "b_max_not_equal_to_limit STRING, " + "b_avg_not_equal_to_limit STRING" ) expected = spark.createDataFrame( @@ -686,6 +698,7 @@ def test_is_aggr_equal(spark: SparkSession): "Sum value 4 in column 'b' is not equal to limit: 10.0", None, "Max value 3 in column 'b' is not equal to limit: 5.0", + None, ], [ "a", @@ -698,6 +711,7 @@ def test_is_aggr_equal(spark: SparkSession): "Sum value 4 in column 'b' is not equal to limit: 10.0", None, "Max value 3 in column 'b' is not equal to limit: 5.0", + None, ], [ "b", @@ -710,6 +724,7 @@ def test_is_aggr_equal(spark: SparkSession): "Sum value 4 in column 'b' is not equal to limit: 10.0", None, "Max value 3 in column 'b' is not equal to limit: 5.0", + None, ], ], expected_schema, @@ -904,9 +919,10 @@ def test_is_aggr_not_equal(spark: SparkSession): is_aggr_not_equal("a", limit=F.lit(1), aggr_type="count", row_filter="b is not null", group_by=["a"]), is_aggr_not_equal(F.col("b"), limit=F.lit(2), aggr_type="count", group_by=[F.col("b")]), is_aggr_not_equal("b", limit=2.0, aggr_type="avg"), - is_aggr_not_equal("b", limit=10.0, aggr_type="sum"), + is_aggr_not_equal("b", limit=Decimal("10.0"), aggr_type="sum"), is_aggr_not_equal("b", limit=1.0, aggr_type="min"), - is_aggr_not_equal("b", limit=5.0, aggr_type="max"), + is_aggr_not_equal("b", limit=Decimal("5.0"), aggr_type="max"), + is_aggr_not_equal("b", limit="2.0", aggr_type="avg"), ] actual = _apply_checks(test_df, checks) @@ -919,7 +935,8 @@ def test_is_aggr_not_equal(spark: SparkSession): "b_avg_equal_to_limit STRING, " "b_sum_equal_to_limit STRING, " "b_min_equal_to_limit STRING, " - "b_max_equal_to_limit STRING" + "b_max_equal_to_limit STRING, " + "b_avg_equal_to_limit STRING" ) expected = spark.createDataFrame( @@ -935,6 +952,7 @@ def test_is_aggr_not_equal(spark: SparkSession): None, "Min value 1 in column 'b' is equal to limit: 1.0", None, + "Average value 2.0 in column 'b' is equal to limit: 2.0", ], [ "a", @@ -947,6 +965,7 @@ def test_is_aggr_not_equal(spark: SparkSession): None, "Min value 1 in column 'b' is equal to limit: 1.0", None, + "Average value 2.0 in column 'b' is equal to limit: 2.0", ], [ "b", @@ -959,6 +978,7 @@ def test_is_aggr_not_equal(spark: SparkSession): None, "Min value 1 in column 'b' is equal to limit: 1.0", None, + "Average value 2.0 in column 'b' is equal to limit: 2.0", ], ], expected_schema, diff --git a/tests/integration/test_dlt_rules_generator.py b/tests/integration/test_dlt_rules_generator.py index f60171fca..6a857e356 100644 --- a/tests/integration/test_dlt_rules_generator.py +++ b/tests/integration/test_dlt_rules_generator.py @@ -19,6 +19,7 @@ def test_generate_dlt_sql_expect(ws, set_utc_timezone): "CONSTRAINT product_launch_date_min_max EXPECT (product_launch_date >= '2020-01-02')", "CONSTRAINT product_expiry_ts_min_max EXPECT (product_expiry_ts <= '2020-01-02T03:04:05.000000')", "CONSTRAINT d1_min_max EXPECT (d1 >= 1.23 and d1 <= 333323.0)", + "CONSTRAINT price_min_max EXPECT (price >= 0.01 and price <= 999.99)", ] assert expectations == expected @@ -34,6 +35,7 @@ def test_generate_dlt_sql_drop(ws): "CONSTRAINT product_launch_date_min_max EXPECT (product_launch_date >= '2020-01-02') ON VIOLATION DROP ROW", "CONSTRAINT product_expiry_ts_min_max EXPECT (product_expiry_ts <= '2020-01-02T03:04:05.000000') ON VIOLATION DROP ROW", "CONSTRAINT d1_min_max EXPECT (d1 >= 1.23 and d1 <= 333323.0) ON VIOLATION DROP ROW", + "CONSTRAINT price_min_max EXPECT (price >= 0.01 and price <= 999.99) ON VIOLATION DROP ROW", ] assert expectations == expected @@ -49,6 +51,7 @@ def test_generate_dlt_sql_fail(ws): "CONSTRAINT product_launch_date_min_max EXPECT (product_launch_date >= '2020-01-02') ON VIOLATION FAIL UPDATE", "CONSTRAINT product_expiry_ts_min_max EXPECT (product_expiry_ts <= '2020-01-02T03:04:05.000000') ON VIOLATION FAIL UPDATE", "CONSTRAINT d1_min_max EXPECT (d1 >= 1.23 and d1 <= 333323.0) ON VIOLATION FAIL UPDATE", + "CONSTRAINT price_min_max EXPECT (price >= 0.01 and price <= 999.99) ON VIOLATION FAIL UPDATE", ] assert expectations == expected @@ -57,7 +60,7 @@ def test_generate_dlt_python_expect(ws): generator = DQDltGenerator(ws) expectations = generator.generate_dlt_rules(test_rules, language="Python") expected = """@dlt.expect_all( -{"vendor_id_is_not_null": "vendor_id is not null", "vendor_id_is_in": "vendor_id in ('1', '4', '2')", "vendor_id_is_not_null_or_empty": "vendor_id is not null and trim(vendor_id) <> ''", "rate_code_id_min_max": "rate_code_id >= 1 and rate_code_id <= 265", "product_launch_date_min_max": "product_launch_date >= '2020-01-02'", "product_expiry_ts_min_max": "product_expiry_ts <= '2020-01-02T03:04:05.000000'", "d1_min_max": "d1 >= 1.23 and d1 <= 333323.0"} +{"vendor_id_is_not_null": "vendor_id is not null", "vendor_id_is_in": "vendor_id in ('1', '4', '2')", "vendor_id_is_not_null_or_empty": "vendor_id is not null and trim(vendor_id) <> ''", "rate_code_id_min_max": "rate_code_id >= 1 and rate_code_id <= 265", "product_launch_date_min_max": "product_launch_date >= '2020-01-02'", "product_expiry_ts_min_max": "product_expiry_ts <= '2020-01-02T03:04:05.000000'", "d1_min_max": "d1 >= 1.23 and d1 <= 333323.0", "price_min_max": "price >= 0.01 and price <= 999.99"} )""" assert expectations == expected @@ -66,7 +69,7 @@ def test_generate_dlt_python_drop(ws): generator = DQDltGenerator(ws) expectations = generator.generate_dlt_rules(test_rules, language="Python", action="drop") expected = """@dlt.expect_all_or_drop( -{"vendor_id_is_not_null": "vendor_id is not null", "vendor_id_is_in": "vendor_id in ('1', '4', '2')", "vendor_id_is_not_null_or_empty": "vendor_id is not null and trim(vendor_id) <> ''", "rate_code_id_min_max": "rate_code_id >= 1 and rate_code_id <= 265", "product_launch_date_min_max": "product_launch_date >= '2020-01-02'", "product_expiry_ts_min_max": "product_expiry_ts <= '2020-01-02T03:04:05.000000'", "d1_min_max": "d1 >= 1.23 and d1 <= 333323.0"} +{"vendor_id_is_not_null": "vendor_id is not null", "vendor_id_is_in": "vendor_id in ('1', '4', '2')", "vendor_id_is_not_null_or_empty": "vendor_id is not null and trim(vendor_id) <> ''", "rate_code_id_min_max": "rate_code_id >= 1 and rate_code_id <= 265", "product_launch_date_min_max": "product_launch_date >= '2020-01-02'", "product_expiry_ts_min_max": "product_expiry_ts <= '2020-01-02T03:04:05.000000'", "d1_min_max": "d1 >= 1.23 and d1 <= 333323.0", "price_min_max": "price >= 0.01 and price <= 999.99"} )""" assert expectations == expected @@ -75,7 +78,7 @@ def test_generate_dlt_python_fail(ws): generator = DQDltGenerator(ws) expectations = generator.generate_dlt_rules(test_rules, language="Python", action="fail") expected = """@dlt.expect_all_or_fail( -{"vendor_id_is_not_null": "vendor_id is not null", "vendor_id_is_in": "vendor_id in ('1', '4', '2')", "vendor_id_is_not_null_or_empty": "vendor_id is not null and trim(vendor_id) <> ''", "rate_code_id_min_max": "rate_code_id >= 1 and rate_code_id <= 265", "product_launch_date_min_max": "product_launch_date >= '2020-01-02'", "product_expiry_ts_min_max": "product_expiry_ts <= '2020-01-02T03:04:05.000000'", "d1_min_max": "d1 >= 1.23 and d1 <= 333323.0"} +{"vendor_id_is_not_null": "vendor_id is not null", "vendor_id_is_in": "vendor_id in ('1', '4', '2')", "vendor_id_is_not_null_or_empty": "vendor_id is not null and trim(vendor_id) <> ''", "rate_code_id_min_max": "rate_code_id >= 1 and rate_code_id <= 265", "product_launch_date_min_max": "product_launch_date >= '2020-01-02'", "product_expiry_ts_min_max": "product_expiry_ts <= '2020-01-02T03:04:05.000000'", "d1_min_max": "d1 >= 1.23 and d1 <= 333323.0", "price_min_max": "price >= 0.01 and price <= 999.99"} )""" assert expectations == expected @@ -129,5 +132,6 @@ def test_generate_dlt_python_dict(ws): "product_launch_date_min_max": "product_launch_date >= '2020-01-02'", "product_expiry_ts_min_max": "product_expiry_ts <= '2020-01-02T03:04:05.000000'", "d1_min_max": "d1 >= 1.23 and d1 <= 333323.0", + "price_min_max": "price >= 0.01 and price <= 999.99", } assert expectations == expected diff --git a/tests/integration/test_row_checks.py b/tests/integration/test_row_checks.py index a4c79f3fa..5ce65f47f 100644 --- a/tests/integration/test_row_checks.py +++ b/tests/integration/test_row_checks.py @@ -1237,18 +1237,22 @@ def test_col_is_not_in_range(spark, set_utc_timezone): is_not_in_range("c", "2025-01-01 00:00:00", "2025-01-03 00:00:00"), is_not_in_range("d", "c", F.expr("cast(b as timestamp) + INTERVAL 2 DAY")), is_not_in_range("e", 1, 3), + is_not_in_range("e", Decimal("1.00"), Decimal("3.00")), is_not_in_range(F.try_element_at("f", F.lit(1)), 1, 3), is_not_in_range("g", 0.2, 0.5), + is_not_in_range("a", 1.5, 3.5), + is_not_in_range("e", 1.5, 3.5), ) checked_schema = ( "a_in_range: string, a_in_range: string, b_in_range: string, b_in_range: string, " "c_in_range: string, c_in_range: string, c_in_range: string, d_in_range: string, e_in_range: string, " - "try_element_at_f_1_in_range: string, g_in_range: string" + "e_in_range: string, try_element_at_f_1_in_range: string, g_in_range: string, " + "a_in_range: string, e_in_range: string" ) expected = spark.createDataFrame( [ - [None, None, None, None, None, None, None, None, None, None, None], + [None, None, None, None, None, None, None, None, None, None, None, None, None, None], [ "Value '1' in Column 'a' in range: [1, 3]", "Value '1' in Column 'a' in range: [1, 3]", @@ -1259,8 +1263,11 @@ def test_col_is_not_in_range(spark, set_utc_timezone): "Value '2025-01-03 00:00:00' in Column 'c' in range: [2025-01-01 00:00:00, 2025-01-03 00:00:00]", None, "Value '1.00' in Column 'e' in range: [1, 3]", + "Value '1.00' in Column 'e' in range: [1.00, 3.00]", "Value '1' in Column 'try_element_at(f, 1)' in range: [1, 3]", "Value '0.3' in Column 'g' in range: [0.2, 0.5]", + "Value '1' in Column 'a' in range: [1.5, 3.5]", + "Value '1.00' in Column 'e' in range: [1.5, 3.5]", ], [ "Value '3' in Column 'a' in range: [1, 3]", @@ -1272,10 +1279,13 @@ def test_col_is_not_in_range(spark, set_utc_timezone): "Value '2025-02-03 00:00:00' in Column 'd' in range: [2025-02-01 00:00:00, 2025-02-03 00:00:00]", "Value '2025-02-03 00:00:00' in Column 'd' in range: [2025-02-01 00:00:00, 2025-02-03 00:00:00]", "Value '3.00' in Column 'e' in range: [1, 3]", + "Value '3.00' in Column 'e' in range: [1.00, 3.00]", "Value '3' in Column 'try_element_at(f, 1)' in range: [1, 3]", None, + "Value '3' in Column 'a' in range: [1.5, 3.5]", + "Value '3.00' in Column 'e' in range: [1.5, 3.5]", ], - [None, None, None, None, None, None, None, None, None, None, None], + [None, None, None, None, None, None, None, None, None, None, None, None, None, None], ], checked_schema, ) @@ -3119,13 +3129,13 @@ def test_is_data_fresh_cur(spark, set_utc_timezone): def test_col_is_not_equal_to(spark, set_utc_timezone): - schema = "a: int, b: int, c: date, d: timestamp, e: decimal(10,2), f: array" + schema = "a: int, b: int, c: date, d: timestamp, e: decimal(10,2), f: array, g: float" test_df = spark.createDataFrame( [ - [1, 1, datetime(2025, 1, 1).date(), datetime(2025, 1, 1), Decimal("1.00"), [1]], - [2, 1, datetime(2025, 2, 1).date(), datetime(2025, 2, 1), Decimal("1.01"), [2]], - [1, 2, None, None, Decimal("0.99"), [1]], - [None, None, None, None, None, [None]], + [1, 1, datetime(2025, 1, 1).date(), datetime(2025, 1, 1), Decimal("1.00"), [1], 1.5], + [2, 1, datetime(2025, 2, 1).date(), datetime(2025, 2, 1), Decimal("1.01"), [2], 2.5], + [1, 2, None, None, Decimal("0.99"), [1], 1.5], + [None, None, None, None, None, [None], None], ], schema, ) @@ -3140,12 +3150,13 @@ def test_col_is_not_equal_to(spark, set_utc_timezone): is_not_equal_to("d", "2025-01-01 00:00:00"), is_not_equal_to("e", Decimal("1.00")), is_not_equal_to(F.try_element_at("f", F.lit(1)), 1), + is_not_equal_to("g", 1.5).alias("g_equal_to_float_value"), ) expected_schema = ( "a_equal_to_literal: string, a_equal_to_str_literal: string, a_equal_to_column: string, " "c_equal_to_value: string, c_equal_to_value: string, d_equal_to_value: string, d_equal_to_value: string, " - "e_equal_to_value: string, try_element_at_f_1_equal_to_value: string" + "e_equal_to_value: string, try_element_at_f_1_equal_to_value: string, g_equal_to_float_value: string" ) expected = spark.createDataFrame( @@ -3160,8 +3171,9 @@ def test_col_is_not_equal_to(spark, set_utc_timezone): "Value '2025-01-01 00:00:00' in Column 'd' is equal to value: 2025-01-01 00:00:00", "Value '1.00' in Column 'e' is equal to value: 1.00", "Value '1' in Column 'try_element_at(f, 1)' is equal to value: 1", + "Value '1.5' in Column 'g' is equal to value: 1.5", ], - [None, None, None, None, None, None, None, None, None], + [None, None, None, None, None, None, None, None, None, None], [ "Value '1' in Column 'a' is equal to value: 1", "Value '1' in Column 'a' is equal to value: 1", @@ -3172,8 +3184,9 @@ def test_col_is_not_equal_to(spark, set_utc_timezone): None, None, "Value '1' in Column 'try_element_at(f, 1)' is equal to value: 1", + "Value '1.5' in Column 'g' is equal to value: 1.5", ], - [None, None, None, None, None, None, None, None, None], + [None, None, None, None, None, None, None, None, None, None], ], expected_schema, ) @@ -3190,19 +3203,18 @@ def test_col_is_not_equal_to_with_tolerance(spark, set_utc_timezone): - rel_tolerance: abs(a - b) <= rel_tolerance * max(abs(a), abs(b)) - If BOTH tolerances are provided, values are equal if EITHER condition is met (OR logic) - **IMPORTANT LIMITATION**: Tolerance is only applied when comparing against int/float values. - When value is Decimal, tolerance parameters are silently ignored and exact equality is used. + Tolerance is applied when comparing against int, float, or Decimal values. Returns error message when values ARE equal (within tolerance if provided), otherwise returns None. Note: This is the inverse of is_equal_to - it fails when values match. """ - schema = "a: int, b: int, c: float" + schema = "a: int, b: int, c: float, d: decimal(10,2)" test_df = spark.createDataFrame( [ - [1, 1, 100.0], - [2, 1, 101.0], - [3, 2, 99.5], - [None, None, None], + [1, 1, 100.0, Decimal("1.00")], + [2, 1, 101.0, Decimal("1.01")], + [3, 2, 99.5, Decimal("0.99")], + [None, None, None, None], ], schema, ) @@ -3211,10 +3223,13 @@ def test_col_is_not_equal_to_with_tolerance(spark, set_utc_timezone): is_not_equal_to("a", 1, rel_tolerance=0.5).alias("a_not_equal_to_value_with_rel_tol"), is_not_equal_to("c", 100.0, abs_tolerance=1).alias("c_not_equal_to_value_with_abs_tol"), is_not_equal_to("c", 100.0, rel_tolerance=0.01).alias("c_not_equal_to_value_with_rel_tol"), + is_not_equal_to("d", Decimal("1.00"), abs_tolerance=0.01).alias("d_not_equal_to_value_with_abs_tol"), + is_not_equal_to("d", Decimal("1.00"), rel_tolerance=0.01).alias("d_not_equal_to_value_with_rel_tol"), ) expected_schema = ( "a_not_equal_to_value_with_abs_tol: string, a_not_equal_to_value_with_rel_tol: string, " - "c_not_equal_to_value_with_abs_tol: string, c_not_equal_to_value_with_rel_tol: string" + "c_not_equal_to_value_with_abs_tol: string, c_not_equal_to_value_with_rel_tol: string, " + "d_not_equal_to_value_with_abs_tol: string, d_not_equal_to_value_with_rel_tol: string" ) expected = spark.createDataFrame( [ @@ -3224,6 +3239,8 @@ def test_col_is_not_equal_to_with_tolerance(spark, set_utc_timezone): "Value '1' in Column 'a' is equal to value: 1", # abs(1-1)=0 <= 0.5*max(1,1)=0.5 ✓ equal, so fails is_not_equal "Value '100.0' in Column 'c' is equal to value: 100.0", # abs(100.0-100.0)=0 <= 1.0 ✓ equal, so fails is_not_equal "Value '100.0' in Column 'c' is equal to value: 100.0", # abs(100.0-100.0)=0 <= 0.01*max(100.0,100.0)=1.0 ✓ equal, so fails is_not_equal + "Value '1.00' in Column 'd' is equal to value: 1.00", # abs(1.00-1.00)=0 <= 0.01 ✓ equal, so fails is_not_equal + "Value '1.00' in Column 'd' is equal to value: 1.00", # abs(1.00-1.00)=0 <= 0.01*max(1.00,1.00)=0.01 ✓ equal, so fails is_not_equal ], # Row 2: a=2, c=101.0, d=1.01 [ @@ -3231,6 +3248,8 @@ def test_col_is_not_equal_to_with_tolerance(spark, set_utc_timezone): "Value '2' in Column 'a' is equal to value: 1", # abs(2-1)=1 <= 0.5*max(2,1)=1.0 ✓ equal, so fails is_not_equal "Value '101.0' in Column 'c' is equal to value: 100.0", # abs(101.0-100.0)=1.0 <= 1.0 ✓ equal, so fails is_not_equal "Value '101.0' in Column 'c' is equal to value: 100.0", # abs(101.0-100.0)=1.0 <= 0.01*max(101.0,100.0)=1.01 ✓ equal, so fails is_not_equal + "Value '1.01' in Column 'd' is equal to value: 1.00", # abs(1.01-1.00)=0.01 <= 0.01 ✓ equal, so fails is_not_equal + "Value '1.01' in Column 'd' is equal to value: 1.00", # abs(1.01-1.00)=0.01 <= 0.01*max(1.01,1.00)=0.0101 ✓ equal, so fails is_not_equal ], # Row 3: a=3, c=99.5, d=0.99 [ @@ -3238,6 +3257,8 @@ def test_col_is_not_equal_to_with_tolerance(spark, set_utc_timezone): None, # abs(3-1)=2 > 0.5*max(3,1)=1.5 ✗ not equal, so passes is_not_equal "Value '99.5' in Column 'c' is equal to value: 100.0", # abs(99.5-100.0)=0.5 <= 1.0 ✓ equal, so fails is_not_equal "Value '99.5' in Column 'c' is equal to value: 100.0", # abs(99.5-100.0)=0.5 <= 0.01*max(99.5,100.0)=1.0 ✓ equal, so fails is_not_equal + "Value '0.99' in Column 'd' is equal to value: 1.00", # abs(0.99-1.00)=0.01 <= 0.01 ✓ equal, so fails is_not_equal + "Value '0.99' in Column 'd' is equal to value: 1.00", # abs(0.99-1.00)=0.01 <= 0.01*max(0.99,1.00)=0.01 ✓ equal, so fails is_not_equal ], # Row 4: All null values [ @@ -3245,6 +3266,8 @@ def test_col_is_not_equal_to_with_tolerance(spark, set_utc_timezone): None, None, None, + None, + None, ], ], expected_schema, @@ -3253,13 +3276,13 @@ def test_col_is_not_equal_to_with_tolerance(spark, set_utc_timezone): def test_col_is_equal_to(spark, set_utc_timezone): - schema = "a: int, b: int, c: date, d: timestamp, e: decimal(10,2), f: array" + schema = "a: int, b: int, c: date, d: timestamp, e: decimal(10,2), f: array, g: float" test_df = spark.createDataFrame( [ - [1, 1, datetime(2025, 1, 1).date(), datetime(2025, 1, 1), Decimal("1.00"), [1]], - [2, 1, datetime(2025, 2, 1).date(), datetime(2025, 2, 1), Decimal("1.01"), [2]], - [1, 2, None, None, Decimal("0.99"), [1]], - [None, None, None, None, None, [None]], + [1, 1, datetime(2025, 1, 1).date(), datetime(2025, 1, 1), Decimal("1.00"), [1], 1.5], + [2, 1, datetime(2025, 2, 1).date(), datetime(2025, 2, 1), Decimal("1.01"), [2], 2.5], + [1, 2, None, None, Decimal("0.99"), [1], 1.5], + [None, None, None, None, None, [None], None], ], schema, ) @@ -3274,18 +3297,20 @@ def test_col_is_equal_to(spark, set_utc_timezone): is_equal_to("d", "2025-01-01 00:00:00"), is_equal_to("e", Decimal("1.00")), is_equal_to(F.try_element_at("f", F.lit(1)), 1), + is_equal_to("g", 1.5).alias("g_not_equal_to_float_value"), ) expected_schema = ( "a_not_equal_to_value: string, a_not_equal_to_str_value: string, a_not_equal_to_value_col: string, " "c_not_equal_to_value: string, c_not_equal_to_value: string, " "d_not_equal_to_value: string, d_not_equal_to_value: string, " - "e_not_equal_to_value: string, try_element_at_f_1_not_equal_to_value: string" + "e_not_equal_to_value: string, try_element_at_f_1_not_equal_to_value: string, " + "g_not_equal_to_float_value: string" ) expected = spark.createDataFrame( [ - [None, None, None, None, None, None, None, None, None], + [None, None, None, None, None, None, None, None, None, None], [ "Value '2' in Column 'a' is not equal to value: 1", "Value '2' in Column 'a' is not equal to value: 1", @@ -3296,6 +3321,7 @@ def test_col_is_equal_to(spark, set_utc_timezone): "Value '2025-02-01 00:00:00' in Column 'd' is not equal to value: 2025-01-01 00:00:00", "Value '1.01' in Column 'e' is not equal to value: 1.00", "Value '2' in Column 'try_element_at(f, 1)' is not equal to value: 1", + "Value '2.5' in Column 'g' is not equal to value: 1.5", ], [ None, @@ -3307,8 +3333,9 @@ def test_col_is_equal_to(spark, set_utc_timezone): None, "Value '0.99' in Column 'e' is not equal to value: 1.00", None, + None, ], - [None, None, None, None, None, None, None, None, None], + [None, None, None, None, None, None, None, None, None, None], ], expected_schema, ) @@ -3325,19 +3352,18 @@ def test_col_is_equal_to_with_tolerance(spark, set_utc_timezone): - rel_tolerance: abs(a - b) <= rel_tolerance * max(abs(a), abs(b)) - If BOTH tolerances are provided, values are equal if EITHER condition is met (OR logic) - **IMPORTANT LIMITATION**: Tolerance is only applied when comparing against int/float values. - When value is Decimal, tolerance parameters are silently ignored and exact equality is used. + Tolerance is applied when comparing against int, float, or Decimal values. Returns None when values are equal (within tolerance if provided), otherwise returns error message. Note: Error messages do NOT include tolerance information, only the actual value and expected value. """ - schema = "a: int, b: int, c: float" + schema = "a: int, b: int, c: float, d: decimal(10,2)" test_df = spark.createDataFrame( [ - [1, 1, 100.0], - [2, 1, 101.0], - [3, 2, 99.5], - [None, None, None], + [1, 1, 100.0, Decimal("1.00")], + [2, 1, 101.0, Decimal("1.01")], + [3, 2, 99.5, Decimal("0.99")], + [None, None, None, None], ], schema, ) @@ -3346,10 +3372,13 @@ def test_col_is_equal_to_with_tolerance(spark, set_utc_timezone): is_equal_to("a", 1, rel_tolerance=0.5).alias("a_equal_to_value_with_rel_tol"), is_equal_to("c", 100.0, abs_tolerance=1).alias("c_equal_to_value_with_abs_tol"), is_equal_to("c", 100.0, rel_tolerance=0.01).alias("c_equal_to_value_with_rel_tol"), + is_equal_to("d", Decimal("1.00"), abs_tolerance=0.01).alias("d_equal_to_value_with_abs_tol"), + is_equal_to("d", Decimal("1.00"), rel_tolerance=0.01).alias("d_equal_to_value_with_rel_tol"), ) expected_schema = ( "a_equal_to_value_with_abs_tol: string, a_equal_to_value_with_rel_tol: string, " - "c_equal_to_value_with_abs_tol: string, c_equal_to_value_with_rel_tol: string" + "c_equal_to_value_with_abs_tol: string, c_equal_to_value_with_rel_tol: string, " + "d_equal_to_value_with_abs_tol: string, d_equal_to_value_with_rel_tol: string" ) expected = spark.createDataFrame( [ @@ -3359,6 +3388,8 @@ def test_col_is_equal_to_with_tolerance(spark, set_utc_timezone): None, # abs(1-1)=0 <= 0.5*max(1,1)=0.5 ✓ None, # abs(100.0-100.0)=0 <= 1.0 ✓ None, # abs(100.0-100.0)=0 <= 0.01*max(100.0,100.0)=1.0 ✓ + None, # abs(1.00-1.00)=0 <= 0.01 ✓ + None, # abs(1.00-1.00)=0 <= 0.01*max(1.00,1.00)=0.01 ✓ ], # Row 2: a=2, c=101.0, d=1.01 [ @@ -3366,6 +3397,8 @@ def test_col_is_equal_to_with_tolerance(spark, set_utc_timezone): None, # abs(2-1)=1 <= 0.5*max(2,1)=1.0 ✓ None, # abs(101.0-100.0)=1.0 <= 1.0 ✓ None, # abs(101.0-100.0)=1.0 <= 0.01*max(101.0,100.0)=1.01 ✓ + None, # abs(1.01-1.00)=0.01 <= 0.01 ✓ equal, so passes is_equal_to + None, # abs(1.01-1.00)=0.01 <= 0.01*max(1.01,1.00)=0.0101 ✓ equal, so passes is_equal_to ], # Row 3: a=3, c=99.5, d=0.99 [ @@ -3373,6 +3406,8 @@ def test_col_is_equal_to_with_tolerance(spark, set_utc_timezone): "Value '3' in Column 'a' is not equal to value: 1", # abs(3-1)=2 > 0.5*max(3,1)=1.5 ✗ None, # abs(99.5-100.0)=0.5 <= 1.0 ✓ None, # abs(99.5-100.0)=0.5 <= 0.01*max(99.5,100.0)=1.0 ✓ + None, # abs(0.99-1.00)=0.01 <= 0.01 ✓ equal, so passes is_equal_to + None, # abs(0.99-1.00)=0.01 <= 0.01*max(0.99,1.00)=0.01 ✓ equal, so passes is_equal_to ], # Row 4: All null values [ @@ -3380,6 +3415,8 @@ def test_col_is_equal_to_with_tolerance(spark, set_utc_timezone): None, None, None, + None, + None, ], ], expected_schema, From 2fe6e8f7b9fe47e0a084e30d273946ae6fe211df Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 16:28:32 +0100 Subject: [PATCH 06/12] updated tests and docstrings --- src/databricks/labs/dqx/check_funcs.py | 6 ++---- tests/integration/test_row_checks.py | 4 ++-- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index aecb3f286..201cf8254 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -618,8 +618,7 @@ def is_equal_to( Args: column (str | Column): Column to check. Can be a string column name or a column expression. - value (int | float | str | datetime.date | datetime.datetime | Column | None, optional): - The value to compare with. Can be a literal or a Spark Column. Defaults to None. + value value: The value to compare with. Can be a number, date, timestamp literal or a Spark Column. Defaults to None. 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: @@ -680,8 +679,7 @@ def is_not_equal_to( Args: column (str | Column): Column to check. Can be a string column name or a column expression. - value (int | float | str | datetime.date | datetime.datetime | Column | None, optional): - The value to compare with. Can be a literal or a Spark Column. Defaults to None. + value: The value to compare with. Can be a number, date, timestamp literal or a Spark Column. Defaults to None. 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: diff --git a/tests/integration/test_row_checks.py b/tests/integration/test_row_checks.py index 5ce65f47f..e4e751d7e 100644 --- a/tests/integration/test_row_checks.py +++ b/tests/integration/test_row_checks.py @@ -1266,8 +1266,8 @@ def test_col_is_not_in_range(spark, set_utc_timezone): "Value '1.00' in Column 'e' in range: [1.00, 3.00]", "Value '1' in Column 'try_element_at(f, 1)' in range: [1, 3]", "Value '0.3' in Column 'g' in range: [0.2, 0.5]", - "Value '1' in Column 'a' in range: [1.5, 3.5]", - "Value '1.00' in Column 'e' in range: [1.5, 3.5]", + None, + None, ], [ "Value '3' in Column 'a' in range: [1, 3]", From 2ac1bdde1c52646ce9ae011b40405b1f50f14eca Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 19:09:32 +0100 Subject: [PATCH 07/12] added serialization/deserialization of decimal values updated tests and docstrings --- src/databricks/labs/dqx/check_funcs.py | 2 +- src/databricks/labs/dqx/checks_serializer.py | 73 +++++++++++++++++-- src/databricks/labs/dqx/checks_storage.py | 14 +++- src/databricks/labs/dqx/utils.py | 26 ++++++- tests/integration/test_build_rules.py | 16 +++- .../test_save_and_load_checks_from_table.py | 42 ++++++++++- .../test_save_checks_to_workspace_file.py | 16 ++++ tests/unit/test_utils.py | 4 +- 8 files changed, 176 insertions(+), 17 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 201cf8254..af582a168 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -618,7 +618,7 @@ def is_equal_to( Args: column (str | Column): Column to check. Can be a string column name or a column expression. - value value: The value to compare with. Can be a number, date, timestamp literal or a Spark Column. Defaults to None. + value: The value to compare with. Can be a number, date, timestamp literal or a Spark Column. Defaults to None. 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: diff --git a/src/databricks/labs/dqx/checks_serializer.py b/src/databricks/labs/dqx/checks_serializer.py index 2ffe493f7..603376bdb 100644 --- a/src/databricks/labs/dqx/checks_serializer.py +++ b/src/databricks/labs/dqx/checks_serializer.py @@ -1,6 +1,7 @@ import logging import json import warnings +from decimal import Decimal from typing import Any from collections.abc import Callable from pathlib import Path @@ -17,7 +18,7 @@ DQForEachColRule, CHECK_FUNC_REGISTRY, ) -from databricks.labs.dqx.utils import safe_json_load +from databricks.labs.dqx.utils import safe_json_load, normalize_bound_args from databricks.labs.dqx.errors import InvalidCheckError CHECKS_TABLE_SCHEMA = ( @@ -131,7 +132,7 @@ def deserialize_checks_to_dataframe( if dq_rule_check.columns is not None: arguments["columns"] = dq_rule_check.columns - json_arguments = {k: json.dumps(v) for k, v in arguments.items()} + json_arguments = {k: json.dumps(normalize_bound_args(v)) for k, v in arguments.items()} dq_rule_rows.append( [ dq_rule_check.name, @@ -260,6 +261,58 @@ def serialize_checks(checks: list[DQRule]) -> list[dict]: return dq_rules +def normalize_checks(checks: list[dict]) -> list[dict]: + """ + Recursively normalize checks dictionary to make it JSON/YAML serializable. + Converts Decimal values to the special format for round-trip preservation. + + Args: + checks: List of check dictionaries that may contain non-serializable values. + + Returns: + List of normalized check dictionaries. + """ + + def normalize_value(val: Any) -> Any: + """Recursively normalize a value.""" + if isinstance(val, dict): + return {k: normalize_value(v) for k, v in val.items()} + # Handle None explicitly since normalize_bound_args doesn't handle it + if val is None: + return None + # For everything else, use normalize_bound_args which handles lists, tuples, Decimal, etc. + return normalize_bound_args(val) + + return [normalize_value(check) for check in checks] + + +def denormalize_checks(checks: list[dict]) -> list[dict]: + """ + Recursively convert Decimal markers back to Decimal objects after deserialization. + Converts {"__decimal__": "0.01"} back to Decimal("0.01"). + + Args: + checks: List of check dictionaries that may contain Decimal markers. + + Returns: + List of check dictionaries with Decimal markers converted to Decimal objects. + """ + + def denormalize_value(val: Any) -> Any: + """Recursively convert Decimal markers back to Decimal objects.""" + if isinstance(val, dict): + # Check if this is a Decimal marker + if "__decimal__" in val and len(val) == 1: + return Decimal(val["__decimal__"]) + # Otherwise, recursively process the dict + return {k: denormalize_value(v) for k, v in val.items()} + if isinstance(val, (list, tuple)): + return type(val)(denormalize_value(v) for v in val) + return val + + return [denormalize_value(check) for check in checks] + + def serialize_checks_to_bytes(checks: list[dict], file_path: Path) -> bytes: """ Serializes a list of checks to bytes in json or yaml (default) format. @@ -271,17 +324,27 @@ def serialize_checks_to_bytes(checks: list[dict], file_path: Path) -> bytes: Serialized checks as bytes. """ serializer = FILE_SERIALIZERS.get(file_path.suffix.lower(), yaml.safe_dump) # default to yaml - return serializer(checks).encode("utf-8") + # Normalize checks to ensure all values are JSON/YAML serializable + normalized_checks = normalize_checks(checks) + return serializer(normalized_checks).encode("utf-8") def get_file_deserializer(filepath: str) -> Callable: """ Get the deserializer function based on file. + The returned function automatically denormalizes Decimal markers back to Decimal objects. Args: filepath: Path to the file. Returns: - Deserializer function. + Deserializer function that also handles Decimal denormalization. """ ext = Path(filepath).suffix.lower() - return FILE_DESERIALIZERS.get(ext.lower(), yaml.safe_load) # default to yaml + base_deserializer = FILE_DESERIALIZERS.get(ext.lower(), yaml.safe_load) # default to yaml + + def deserializer_with_denormalization(file_like) -> list[dict]: + """Wrapper that deserializes and then denormalizes special markers (e.g. __decimal__).""" + checks = base_deserializer(file_like) or [] + return denormalize_checks(checks) + + return deserializer_with_denormalization diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index cce1de22c..5f98b9e5a 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -43,6 +43,8 @@ from databricks.sdk import WorkspaceClient from databricks.labs.dqx.checks_serializer import ( + normalize_checks, + denormalize_checks, serialize_checks_from_dataframe, deserialize_checks_to_dataframe, serialize_checks_to_bytes, @@ -242,6 +244,8 @@ def get_table_definition(schema_name: str, table_name: str) -> Table: def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) -> list[dict]: """ Normalize the checks to be compatible with the Lakebase table. + This includes normalizing Decimal values for JSON serialization and structuring + the checks for the Lakebase table schema. Args: checks: List of dq rules (checks) to normalize. @@ -250,8 +254,12 @@ def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) - Returns: List of normalized dq rules (checks). """ + # First normalize Decimal values for JSON serialization + normalized_for_serialization = normalize_checks(checks) + + # Then normalize the structure for Lakebase table normalized_checks = [] - for check in checks: + for check in normalized_for_serialization: user_metadata = check.get("user_metadata") normalized_check = { "name": check.get("name"), @@ -347,7 +355,9 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine f"for run_config_name='{config.run_config_name}'. " f"Make sure the profiler has run successfully and saved checks to this location." ) - return [dict(check) for check in checks] + checks_dict = [dict(check) for check in checks] + # Denormalize Decimal markers back to Decimal objects + return denormalize_checks(checks_dict) def _check_for_undefined_table_error(self, e: ProgrammingError, config: LakebaseChecksStorageConfig) -> NoReturn: """ diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index 39d8bb2d0..d7641d9f5 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -3,6 +3,7 @@ import datetime import logging import re +from decimal import Decimal from importlib.util import find_spec from typing import Any from fnmatch import fnmatch @@ -125,14 +126,17 @@ def normalize_bound_args(val: Any) -> Any: """ Normalize a value or collection of values for consistent processing. - Handles primitives, dates, and column-like objects. Lists, tuples, and sets are + Handles primitives, dates, Decimal, and column-like objects. Lists, tuples, and sets are recursively normalized with type preserved. + For Decimal values, uses a special JSON-serializable format to preserve type information + for round-trip deserialization. + Args: val: Value or collection of values to normalize. Returns: - Normalized value or collection. + Normalized value or collection. Decimal values are converted to a special dict format. Raises: TypeError: If a column type is unsupported. @@ -147,6 +151,10 @@ def normalize_bound_args(val: Any) -> Any: if isinstance(val, (datetime.date, datetime.datetime)): return str(val) + if isinstance(val, Decimal): + # Use a special format to preserve Decimal type information for round-trip + return {"__decimal__": str(val)} + if ConnectColumn is not None: column_types: tuple[type[Any], ...] = (Column, ConnectColumn) else: @@ -206,13 +214,25 @@ def safe_json_load(value: str): Safely load a JSON string, returning the original value if it fails to parse. This allows to specify string value without a need to escape the quotes. + Also handles special Decimal format: {"__decimal__": "0.01"} is converted back to Decimal. + Args: value: The value to parse as JSON. + + Returns: + Parsed JSON value, or original value if parsing fails. Decimal markers are converted to Decimal objects. """ try: - return json.loads(value) # load as json if possible + parsed = json.loads(value) # load as json if possible + # Check if this is a Decimal marker and convert back to Decimal + if isinstance(parsed, dict) and "__decimal__" in parsed and len(parsed) == 1: + return Decimal(parsed["__decimal__"]) + return parsed except json.JSONDecodeError: return value + except (ValueError, TypeError): + # If Decimal conversion fails, return the parsed value as-is + return value def safe_strip_file_from_path(path: str) -> str: diff --git a/tests/integration/test_build_rules.py b/tests/integration/test_build_rules.py index d5fb4c2e4..2d256e374 100644 --- a/tests/integration/test_build_rules.py +++ b/tests/integration/test_build_rules.py @@ -1,3 +1,5 @@ +from decimal import Decimal + from databricks.labs.dqx.checks_serializer import deserialize_checks_to_dataframe, serialize_checks_from_dataframe SCHEMA = "a: int, b: int, c: int" @@ -23,7 +25,7 @@ def test_build_quality_rules_from_dataframe_round_trip(spark): { "name": "column_col2_not_less_than", "criticality": "warn", - "check": {"function": "is_not_greater_than", "arguments": {"column": "test_col2", "limit": 1}}, + "check": {"function": "is_not_greater_than", "arguments": {"column": "test_col2", "limit": 1.01}}, }, { "name": "column_in_list", @@ -115,6 +117,18 @@ def test_build_quality_rules_from_dataframe_round_trip(spark): "arguments": {"column": "c"}, }, }, + { + "name": "price_isnt_in_range", + "criticality": "error", + "check": { + "function": "is_in_range", + "arguments": { + "column": "price", + "min_limit": Decimal("0.01"), + "max_limit": Decimal("999.99"), + }, + }, + }, ] df = deserialize_checks_to_dataframe(spark, test_checks) diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index be01680cb..247d89ce3 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -1,5 +1,6 @@ import json from dataclasses import dataclass +from decimal import Decimal import pytest from chispa.dataframe_comparer import assert_df_equality # type: ignore @@ -27,7 +28,7 @@ { "name": "column_not_less_than", "criticality": "warn", - "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1}}, + "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1.01}}, "filter": "Col_3 >1", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, @@ -36,6 +37,14 @@ "name": "column_in_list", "check": {"function": "is_in_list", "arguments": {"column": "col_2", "allowed": [1, 2]}}, }, + { + "criticality": "warn", + "name": "col_3_is_in_range", + "check": { + "function": "is_in_range", + "arguments": {"column": "col_3", "min_limit": Decimal("0.01"), "max_limit": Decimal("999.99")}, + }, + }, ] EXPECTED_CHECKS = [ @@ -54,7 +63,7 @@ { "name": "column_not_less_than", "criticality": "warn", - "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1}}, + "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1.01}}, "filter": "Col_3 >1", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, @@ -63,6 +72,18 @@ "criticality": "warn", "check": {"function": "is_in_list", "arguments": {"column": "col_2", "allowed": [1, 2]}}, }, + { + "name": "col_3_is_in_range", + "criticality": "warn", + "check": { + "function": "is_in_range", + "arguments": { + "column": "col_3", + "min_limit": Decimal("0.01"), + "max_limit": Decimal("999.99"), + }, + }, + }, ] @@ -141,7 +162,7 @@ def test_save_checks_to_table_with_unresolved_for_each_column(ws, make_schema, m { "name": "column_not_less_than", "criticality": "warn", - "check": {"function": "is_not_less_than", "arguments": {"limit": "1", "column": "\"col_2\""}}, + "check": {"function": "is_not_less_than", "arguments": {"limit": "1.01", "column": "\"col_2\""}}, "filter": "Col_3 >1", "run_config_name": "default", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, @@ -154,6 +175,21 @@ def test_save_checks_to_table_with_unresolved_for_each_column(ws, make_schema, m "run_config_name": "default", "user_metadata": None, }, + { + "name": "col_3_is_in_range", + "criticality": "warn", + "check": { + "function": "is_in_range", + "arguments": { + "column": '"col_3"', + "min_limit": '{"__decimal__": "0.01"}', + "max_limit": '{"__decimal__": "999.99"}', + }, + }, + "filter": None, + "run_config_name": "default", + "user_metadata": None, + }, ] expected_checks_df = spark.createDataFrame(expected_raw_checks, CHECKS_TABLE_SCHEMA) diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index 796157e1f..d748e93f7 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -1,4 +1,5 @@ import json +from decimal import Decimal from unittest.mock import patch import pytest import yaml @@ -19,6 +20,13 @@ "criticality": "error", "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, }, + { + "name": "column_not_less_than", + "criticality": "error", + "check": {"function": "is_not_less_than", "arguments": {"column": "col2", "limit": 1.01}}, + "filter": "Col3 >1", + "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, + }, { "check": { "function": "is_not_null", @@ -33,6 +41,14 @@ "name": "cost_is_null", "criticality": "error", }, + { + "criticality": "warn", + "name": "col3_is_in_range", + "check": { + "function": "is_in_range", + "arguments": {"column": "col3", "min_limit": Decimal("0.01"), "max_limit": Decimal("999.99")}, + }, + }, ] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e7644bbac..5502d03de 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -284,8 +284,8 @@ def test_safe_json_load_empty_string(): def test_safe_json_load_non_string_arg(): - with pytest.raises(TypeError): - safe_json_load(123) + result = safe_json_load(123) + assert result == 123 @pytest.mark.parametrize( From a52ea86f1dc2e48c5b3c7aa37caadde7716328e2 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 21:30:27 +0100 Subject: [PATCH 08/12] added serialization/deserialization of decimal values refactor updated tests and docstrings --- src/databricks/labs/dqx/checks_serializer.py | 674 +++++++++++-------- src/databricks/labs/dqx/checks_storage.py | 57 +- src/databricks/labs/dqx/config.py | 4 +- src/databricks/labs/dqx/engine.py | 6 +- src/databricks/labs/dqx/utils.py | 31 +- tests/integration/test_build_rules.py | 14 +- tests/integration/test_summary_metrics.py | 22 +- tests/unit/test_build_rules.py | 40 +- tests/unit/test_utils.py | 40 +- 9 files changed, 511 insertions(+), 377 deletions(-) diff --git a/src/databricks/labs/dqx/checks_serializer.py b/src/databricks/labs/dqx/checks_serializer.py index 603376bdb..9fe3dfaf0 100644 --- a/src/databricks/labs/dqx/checks_serializer.py +++ b/src/databricks/labs/dqx/checks_serializer.py @@ -1,10 +1,10 @@ import logging import json import warnings +from abc import ABC, abstractmethod from decimal import Decimal -from typing import Any +from typing import Any, TextIO from collections.abc import Callable -from pathlib import Path import yaml from pyspark.sql import DataFrame, SparkSession @@ -17,8 +17,9 @@ DQDatasetRule, DQForEachColRule, CHECK_FUNC_REGISTRY, + normalize_bound_args, ) -from databricks.labs.dqx.utils import safe_json_load, normalize_bound_args +from databricks.labs.dqx.utils import safe_json_load from databricks.labs.dqx.errors import InvalidCheckError CHECKS_TABLE_SCHEMA = ( @@ -26,325 +27,426 @@ " arguments MAP>, filter STRING, run_config_name STRING, user_metadata MAP" ) -FILE_SERIALIZERS: dict[str, Callable[[list[dict[Any, Any]]], str]] = { - ".json": json.dumps, - ".yml": yaml.safe_dump, - ".yaml": yaml.safe_dump, -} - - -FILE_DESERIALIZERS: dict[str, Callable] = { - ".json": json.load, - ".yaml": yaml.safe_load, - ".yml": yaml.safe_load, -} - logger = logging.getLogger(__name__) -def serialize_checks_from_dataframe(df: DataFrame, run_config_name: str = "default") -> list[dict]: +class ChecksNormalizer: """ - Converts a list of quality checks defined in a DataFrame to a list of quality checks - defined as Python dictionaries. - - Args: - df: DataFrame with data quality check rules. Each row should define a check. Rows should - have the following columns: - - *name* - Name that will be given to a resulting column. Autogenerated if not provided. - - *criticality* (optional) - Possible values are *error* (data going only into "bad" dataframe) and *warn* (data is going into both dataframes). - - *check* - DQX check function used in the check; A *StructType* column defining the data quality check. - - *filter* - Expression for filtering data quality checks. - - *run_config_name* (optional) - Run configuration name for storing checks across runs. - - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. - run_config_name: Run configuration name for filtering quality rules, e.g. input table or job name (use "default" if not provided). - - Returns: - List of data quality check specifications as a Python dictionary + Handles normalization and denormalization of check dictionaries. + E.g. responsible for converting Decimal values to/from serializable format. """ - check_rows = df.where(f"run_config_name = '{run_config_name}'").collect() - collect_limit = 500 - if len(check_rows) > collect_limit: - warnings.warn( - f"Collecting large number of rows from DataFrame: {len(check_rows)}", - category=UserWarning, - stacklevel=2, - ) - - checks = [] - for row in check_rows: - check_dict = { - "name": row.name, - "criticality": row.criticality, - "check": { - "function": row.check["function"], - "arguments": ( - {k: safe_json_load(v) for k, v in row.check["arguments"].items()} - if row.check["arguments"] is not None - else {} - ), - }, - } - if "for_each_column" in row.check and row.check["for_each_column"]: - check_dict["check"]["for_each_column"] = row.check["for_each_column"] - if row.filter is not None: - check_dict["filter"] = row.filter - if row.user_metadata is not None: - check_dict["user_metadata"] = row.user_metadata - checks.append(check_dict) - return checks - - -def deserialize_checks_to_dataframe( - spark: SparkSession, - checks: list[dict], - run_config_name: str = "default", -) -> DataFrame: - """ - Converts a list of quality checks defined as Python dictionaries to a DataFrame. - - Args: - spark: Spark session. - checks: list of check specifications as Python dictionaries. Each check consists of the following fields: - - *check* - Column expression to evaluate. This expression should return string value if it's evaluated to - true (it will be used as an error/warning message) or *null* if it's evaluated to *false* - - *name* - Name that will be given to a resulting column. Autogenerated if not provided - - *criticality* (optional) - Possible values are *error* (data going only into "bad" dataframe) and *warn* - (data is going into both dataframes) - - *filter* (optional) - Expression for filtering data quality checks - - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. - run_config_name: Run configuration name for storing quality checks across runs, e.g. input table or job name (use "default" if not provided) - - Returns: - DataFrame with data quality check rules - - Raises: - InvalidCheckError: If any check is invalid or unsupported. - """ - dq_rule_checks: list[DQRule] = deserialize_checks(checks) - dq_rule_rows = [] - for dq_rule_check in dq_rule_checks: - arguments = dict(dq_rule_check.check_func_kwargs) + @staticmethod + def normalize(checks: list[dict]) -> list[dict]: + """ + Recursively normalize checks dictionary to make it JSON/YAML serializable. - if dq_rule_check.column is not None: - arguments["column"] = dq_rule_check.column + Args: + checks: List of check dictionaries that may contain non-serializable values. - if dq_rule_check.columns is not None: - arguments["columns"] = dq_rule_check.columns + Returns: + List of normalized check dictionaries. + """ - json_arguments = {k: json.dumps(normalize_bound_args(v)) for k, v in arguments.items()} - dq_rule_rows.append( - [ - dq_rule_check.name, - dq_rule_check.criticality, - {"function": dq_rule_check.check_func.__name__, "arguments": json_arguments}, - dq_rule_check.filter, - run_config_name, - dq_rule_check.user_metadata, - ] - ) - return spark.createDataFrame(dq_rule_rows, CHECKS_TABLE_SCHEMA) + def normalize_value(val: Any) -> Any: + """Recursively normalize a value.""" + if isinstance(val, dict): + return {k: normalize_value(v) for k, v in val.items()} + # normalize_bound_args handles None, primitives, lists, tuples, Decimal, etc. + return normalize_bound_args(val) + return [normalize_value(check) for check in checks] -def deserialize_checks(checks: list[dict], custom_checks: dict[str, Callable] | None = None) -> list[DQRule]: - """ - Converts a list of quality checks defined as Python dictionaries to a list of `DQRule` objects. - - Args: - checks: list of dictionaries describing checks. Each check is a dictionary - consisting of following fields: - - *check* - Column expression to evaluate. This expression should return string value if it's evaluated to true - or *null* if it's evaluated to *false* - - *name* - name that will be given to a resulting column. Autogenerated if not provided - - *criticality* (optional) - possible values are *error* (data going only into "bad" dataframe), - and *warn* (data is going into both dataframes) - - *filter* (optional) - Expression for filtering data quality checks - - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. - custom_checks: dictionary with custom check functions (e.g. *globals()* of the calling module). - If not specified, then only built-in functions are used for the checks. - - Returns: - list of data quality check rules - - Raises: - InvalidCheckError: If any dictionary is invalid or unsupported. - """ - status = ChecksValidator.validate_checks(checks, custom_checks) - if status.has_errors: - raise InvalidCheckError(str(status)) - - dq_rule_checks: list[DQRule] = [] - for check_def in checks: - logger.debug(f"Processing check definition: {check_def}") - - check = check_def.get("check", {}) - name = check_def.get("name", None) - func_name = check.get("function") - func = resolve_check_function(func_name, custom_checks, fail_on_missing=True) - assert func # should already be validated - - func_args = check.get("arguments", {}) - for_each_column = check.get("for_each_column") - column = func_args.get("column") # should be defined for single-column checks only - columns = func_args.get("columns") # should be defined for multi-column checks only - assert not (column and columns) # should already be validated - criticality = check_def.get("criticality", "error") - filter_str = check_def.get("filter") - user_metadata = check_def.get("user_metadata") - - # Exclude `column` and `columns` from check_func_kwargs - # as these are always included in the check function call - check_func_kwargs = {k: v for k, v in func_args.items() if k not in {"column", "columns"}} - - # treat non-registered function as row-level checks - if for_each_column: - dq_rule_checks += DQForEachColRule( - columns=for_each_column, - name=name, - check_func=func, - criticality=criticality, - filter=filter_str, - check_func_kwargs=check_func_kwargs, - user_metadata=user_metadata, - ).get_rules() - else: - rule_type = CHECK_FUNC_REGISTRY.get(func_name) - if rule_type == "dataset": - dq_rule_checks.append( - DQDatasetRule( - column=column, - columns=columns, - check_func=func, - check_func_kwargs=check_func_kwargs, - name=name, - criticality=criticality, - filter=filter_str, - user_metadata=user_metadata, - ) - ) - else: # default to row-level rule - dq_rule_checks.append( - DQRowRule( - column=column, - columns=columns, - check_func=func, - check_func_kwargs=check_func_kwargs, - name=name, - criticality=criticality, - filter=filter_str, - user_metadata=user_metadata, - ) - ) + @staticmethod + def denormalize_value(val: Any) -> Any: + """Recursively convert special markers (e.g. Decimal) back to original objects.""" + if isinstance(val, dict): + # Check if this is a Decimal marker + if "__decimal__" in val and len(val) == 1: + return Decimal(val["__decimal__"]) + # Otherwise, recursively process the dict + return {k: ChecksNormalizer.denormalize_value(v) for k, v in val.items()} + if isinstance(val, (list, tuple)): + return type(val)(ChecksNormalizer.denormalize_value(v) for v in val) + return val - return dq_rule_checks + @staticmethod + def denormalize(checks: list[dict]) -> list[dict]: + """ + Recursively convert special markers back to objects after deserialization. + Converts {"__decimal__": "0.01"} back to Decimal("0.01"). + Args: + checks: List of check dictionaries that may contain special markers. -def serialize_checks(checks: list[DQRule]) -> list[dict]: + Returns: + List of check dictionaries with special markers converted to objects. + """ + return [ChecksNormalizer.denormalize_value(check) for check in checks] + + +class FileFormatSerializer(ABC): + """ + Abstract base class for file format serializers. """ - Converts a list of quality checks defined as *DQRule* objects to a list of quality checks - defined as Python dictionaries. - Args: - checks: List of DQRule instances to convert. + @abstractmethod + def serialize(self, data: list[dict]) -> str: + """Serialize data to string format.""" - Returns: - List of dictionaries representing the DQRule instances. + @abstractmethod + def deserialize(self, file_like: TextIO) -> list[dict]: + """Deserialize data from file-like object.""" - Raises: - InvalidCheckError: If any item in the list is not a DQRule instance. - """ - dq_rules = [] - for check in checks: - if not isinstance(check, DQRule): - raise InvalidCheckError(f"Expected DQRule instance, got {type(check).__name__}") - dq_rules.append(check.to_dict()) - return dq_rules +class JsonSerializer(FileFormatSerializer): + """JSON format serializer implementation.""" -def normalize_checks(checks: list[dict]) -> list[dict]: - """ - Recursively normalize checks dictionary to make it JSON/YAML serializable. - Converts Decimal values to the special format for round-trip preservation. + def serialize(self, data: list[dict]) -> str: + """Serialize data to JSON string.""" + return json.dumps(data) - Args: - checks: List of check dictionaries that may contain non-serializable values. + def deserialize(self, file_like: TextIO) -> list[dict]: + """Deserialize data from JSON file.""" + return json.load(file_like) or [] - Returns: - List of normalized check dictionaries. - """ - def normalize_value(val: Any) -> Any: - """Recursively normalize a value.""" - if isinstance(val, dict): - return {k: normalize_value(v) for k, v in val.items()} - # Handle None explicitly since normalize_bound_args doesn't handle it - if val is None: - return None - # For everything else, use normalize_bound_args which handles lists, tuples, Decimal, etc. - return normalize_bound_args(val) +class YamlSerializer(FileFormatSerializer): + """YAML format serializer implementation.""" + + def serialize(self, data: list[dict]) -> str: + """Serialize data to YAML string.""" + return yaml.safe_dump(data) - return [normalize_value(check) for check in checks] + def deserialize(self, file_like: TextIO) -> list[dict]: + """Deserialize data from YAML file.""" + return yaml.safe_load(file_like) or [] -def denormalize_checks(checks: list[dict]) -> list[dict]: +class SerializerFactory: + """ + Factory for creating appropriate serializers based on file extension. """ - Recursively convert Decimal markers back to Decimal objects after deserialization. - Converts {"__decimal__": "0.01"} back to Decimal("0.01"). - Args: - checks: List of check dictionaries that may contain Decimal markers. + _serializers: dict[str, type[FileFormatSerializer]] = { + ".json": JsonSerializer, + ".yaml": YamlSerializer, + ".yml": YamlSerializer, + } + + @classmethod + def get_supported_extensions(cls) -> tuple[str, ...]: + """ + Get tuple of supported file extensions. + + Returns: + Tuple of supported file extensions (e.g., (".json", ".yaml", ".yml")). + """ + return tuple(cls._serializers.keys()) + + @classmethod + def create_serializer(cls, extension: str | None = None) -> FileFormatSerializer: + """ + Create a serializer based on file extension. + + Args: + extension: File extension (e.g., ".json", ".yaml", ".yml"). + If None or empty, defaults to YAML. + + Returns: + Appropriate serializer instance. Defaults to YAML if extension not recognized or not provided. + """ + if not extension: + return YamlSerializer() + ext = extension.lower() + serializer_class = cls._serializers.get(ext, YamlSerializer) + return serializer_class() + + @classmethod + def register_format(cls, extension: str, serializer_class: type[FileFormatSerializer]) -> None: + """ + Register a new file format serializer. + + Args: + extension: File extension + serializer_class: Serializer class implementing FileFormatSerializer interface. + """ + cls._serializers[extension.lower()] = serializer_class + + +class ChecksSerializer: + """ + Handles serialization of DQRule objects to dictionaries and file formats. + """ - Returns: - List of check dictionaries with Decimal markers converted to Decimal objects. + @staticmethod + def serialize(checks: list[DQRule]) -> list[dict]: + """ + Converts a list of quality checks defined as *DQRule* objects to a list of quality checks + defined as Python dictionaries. + + Args: + checks: List of DQRule instances to convert. + + Returns: + List of dictionaries representing the DQRule instances. + + Raises: + InvalidCheckError: If any item in the list is not a DQRule instance. + """ + dq_rules = [] + for check in checks: + if not isinstance(check, DQRule): + raise InvalidCheckError(f"Expected DQRule instance, got {type(check).__name__}") + dq_rules.append(check.to_dict()) + return dq_rules + + @staticmethod + def serialize_to_bytes(checks: list[dict], extension: str) -> bytes: + """ + Serializes a list of checks to bytes in json or yaml (default) format. + + Args: + checks: List of checks to serialize. + extension: File extension (e.g., ".json", ".yaml", ".yml"). + Returns: + Serialized checks as bytes. + """ + serializer = SerializerFactory.create_serializer(extension) + normalized_checks = ChecksNormalizer.normalize(checks) + serialized_str = serializer.serialize(normalized_checks) + return serialized_str.encode("utf-8") + + +class ChecksDeserializer: + """ + Handles deserialization of dictionaries to DQRule objects and from file formats. """ - def denormalize_value(val: Any) -> Any: - """Recursively convert Decimal markers back to Decimal objects.""" - if isinstance(val, dict): - # Check if this is a Decimal marker - if "__decimal__" in val and len(val) == 1: - return Decimal(val["__decimal__"]) - # Otherwise, recursively process the dict - return {k: denormalize_value(v) for k, v in val.items()} - if isinstance(val, (list, tuple)): - return type(val)(denormalize_value(v) for v in val) - return val + def __init__(self, custom_checks: dict[str, Callable] | None = None): + """ + Initialize the deserializer. + + Args: + custom_checks: Dictionary with custom check functions. + """ + self.custom_checks = custom_checks + + def deserialize(self, checks: list[dict]) -> list[DQRule]: + """ + Converts a list of quality checks defined as Python dictionaries to a list of `DQRule` objects. + + Args: + checks: list of dictionaries describing checks. Each check is a dictionary + consisting of following fields: + - *check* - Column expression to evaluate. This expression should return string value if it's evaluated to true + or *null* if it's evaluated to *false* + - *name* - name that will be given to a resulting column. Autogenerated if not provided + - *criticality* (optional) - possible values are *error* (data going only into "bad" dataframe), + and *warn* (data is going into both dataframes) + - *filter* (optional) - Expression for filtering data quality checks + - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. + + Returns: + list of data quality check rules + + Raises: + InvalidCheckError: If any dictionary is invalid or unsupported. + """ + status = ChecksValidator.validate_checks(checks, self.custom_checks) + if status.has_errors: + raise InvalidCheckError(str(status)) + + dq_rule_checks: list[DQRule] = [] + for check_def in checks: + logger.debug(f"Processing check definition: {check_def}") + + check = check_def.get("check", {}) + name = check_def.get("name", None) + func_name = check.get("function") + func = resolve_check_function(func_name, self.custom_checks, fail_on_missing=True) + assert func # should already be validated + + func_args = check.get("arguments", {}) + for_each_column = check.get("for_each_column") + column = func_args.get("column") # should be defined for single-column checks only + columns = func_args.get("columns") # should be defined for multi-column checks only + assert not (column and columns) # should already be validated + criticality = check_def.get("criticality", "error") + filter_str = check_def.get("filter") + user_metadata = check_def.get("user_metadata") + + # Exclude `column` and `columns` from check_func_kwargs + # as these are always included in the check function call + check_func_kwargs = {k: v for k, v in func_args.items() if k not in {"column", "columns"}} + + # treat non-registered function as row-level checks + if for_each_column: + dq_rule_checks += DQForEachColRule( + columns=for_each_column, + name=name, + check_func=func, + criticality=criticality, + filter=filter_str, + check_func_kwargs=check_func_kwargs, + user_metadata=user_metadata, + ).get_rules() + else: + rule_type = CHECK_FUNC_REGISTRY.get(func_name) + if rule_type == "dataset": + dq_rule_checks.append( + DQDatasetRule( + column=column, + columns=columns, + check_func=func, + check_func_kwargs=check_func_kwargs, + name=name, + criticality=criticality, + filter=filter_str, + user_metadata=user_metadata, + ) + ) + else: # default to row-level rule + dq_rule_checks.append( + DQRowRule( + column=column, + columns=columns, + check_func=func, + check_func_kwargs=check_func_kwargs, + name=name, + criticality=criticality, + filter=filter_str, + user_metadata=user_metadata, + ) + ) - return [denormalize_value(check) for check in checks] + return dq_rule_checks + @staticmethod + def deserialize_from_file(extension: str, file_like: TextIO) -> list[dict]: + """ + Deserialize checks from a file-like object based on file extension. + Automatically denormalizes special markers back to objects. -def serialize_checks_to_bytes(checks: list[dict], file_path: Path) -> bytes: - """ - Serializes a list of checks to bytes in json or yaml (default) format. + Args: + extension: File extension (e.g., ".json", ".yaml", ".yml"). + file_like: File-like object to read from. - Args: - checks: List of checks to serialize. - file_path: Path to the file where the checks will be serialized. - Returns: - Serialized checks as bytes. - """ - serializer = FILE_SERIALIZERS.get(file_path.suffix.lower(), yaml.safe_dump) # default to yaml - # Normalize checks to ensure all values are JSON/YAML serializable - normalized_checks = normalize_checks(checks) - return serializer(normalized_checks).encode("utf-8") + Returns: + List of check dictionaries with special markers converted to objects. + """ + serializer = SerializerFactory.create_serializer(extension) + checks = serializer.deserialize(file_like) + return ChecksNormalizer.denormalize(checks) -def get_file_deserializer(filepath: str) -> Callable: +class DataFrameConverter: """ - Get the deserializer function based on file. - The returned function automatically denormalizes Decimal markers back to Decimal objects. - - Args: - filepath: Path to the file. - Returns: - Deserializer function that also handles Decimal denormalization. + Handles conversion between DataFrames and check dictionaries. """ - ext = Path(filepath).suffix.lower() - base_deserializer = FILE_DESERIALIZERS.get(ext.lower(), yaml.safe_load) # default to yaml - - def deserializer_with_denormalization(file_like) -> list[dict]: - """Wrapper that deserializes and then denormalizes special markers (e.g. __decimal__).""" - checks = base_deserializer(file_like) or [] - return denormalize_checks(checks) - - return deserializer_with_denormalization + + @staticmethod + def from_dataframe(df: DataFrame, run_config_name: str = "default") -> list[dict]: + """ + Converts a list of quality checks defined in a DataFrame to a list of quality checks + defined as Python dictionaries. + + Args: + df: DataFrame with data quality check rules. Each row should define a check. Rows should + have the following columns: + - *name* - Name that will be given to a resulting column. Autogenerated if not provided. + - *criticality* (optional) - Possible values are *error* (data going only into "bad" dataframe) and *warn* (data is going into both dataframes). + - *check* - DQX check function used in the check; A *StructType* column defining the data quality check. + - *filter* - Expression for filtering data quality checks. + - *run_config_name* (optional) - Run configuration name for storing checks across runs. + - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. + run_config_name: Run configuration name for filtering quality rules, e.g. input table or job name (use "default" if not provided). + + Returns: + List of data quality check specifications as a Python dictionary + """ + check_rows = df.where(f"run_config_name = '{run_config_name}'").collect() + collect_limit = 500 + if len(check_rows) > collect_limit: + warnings.warn( + f"Collecting large number of rows from DataFrame: {len(check_rows)}", + category=UserWarning, + stacklevel=2, + ) + + checks = [] + for row in check_rows: + check_dict = { + "name": row.name, + "criticality": row.criticality, + "check": { + "function": row.check["function"], + "arguments": ( + {k: safe_json_load(v) for k, v in row.check["arguments"].items()} + if row.check["arguments"] is not None + else {} + ), + }, + } + if "for_each_column" in row.check and row.check["for_each_column"]: + check_dict["check"]["for_each_column"] = row.check["for_each_column"] + if row.filter is not None: + check_dict["filter"] = row.filter + if row.user_metadata is not None: + check_dict["user_metadata"] = row.user_metadata + # Denormalize special markers back to objects + checks.append(ChecksNormalizer.denormalize_value(check_dict)) + return checks + + @staticmethod + def to_dataframe( + spark: SparkSession, + checks: list[dict], + run_config_name: str = "default", + ) -> DataFrame: + """ + Converts a list of quality checks defined as Python dictionaries to a DataFrame. + + Args: + spark: Spark session. + checks: list of check specifications as Python dictionaries. Each check consists of the following fields: + - *check* - Column expression to evaluate. This expression should return string value if it's evaluated to + true (it will be used as an error/warning message) or *null* if it's evaluated to *false* + - *name* - Name that will be given to a resulting column. Autogenerated if not provided + - *criticality* (optional) - Possible values are *error* (data going only into "bad" dataframe) and *warn* + (data is going into both dataframes) + - *filter* (optional) - Expression for filtering data quality checks + - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. + run_config_name: Run configuration name for storing quality checks across runs, e.g. input table or job name (use "default" if not provided) + + Returns: + DataFrame with data quality check rules + + Raises: + InvalidCheckError: If any check is invalid or unsupported. + """ + deserializer = ChecksDeserializer() + dq_rule_checks: list[DQRule] = deserializer.deserialize(checks) + + dq_rule_rows = [] + for dq_rule_check in dq_rule_checks: + arguments = dict(dq_rule_check.check_func_kwargs) + + if dq_rule_check.column is not None: + arguments["column"] = dq_rule_check.column + + if dq_rule_check.columns is not None: + arguments["columns"] = dq_rule_check.columns + + json_arguments = {k: json.dumps(normalize_bound_args(v)) for k, v in arguments.items()} + dq_rule_rows.append( + [ + dq_rule_check.name, + dq_rule_check.criticality, + {"function": dq_rule_check.check_func.__name__, "arguments": json_arguments}, + dq_rule_check.filter, + run_config_name, + dq_rule_check.user_metadata, + ] + ) + return spark.createDataFrame(dq_rule_rows, CHECKS_TABLE_SCHEMA) diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py index 5f98b9e5a..068005134 100644 --- a/src/databricks/labs/dqx/checks_storage.py +++ b/src/databricks/labs/dqx/checks_storage.py @@ -43,14 +43,13 @@ from databricks.sdk import WorkspaceClient from databricks.labs.dqx.checks_serializer import ( - normalize_checks, - denormalize_checks, - serialize_checks_from_dataframe, - deserialize_checks_to_dataframe, - serialize_checks_to_bytes, - get_file_deserializer, - FILE_SERIALIZERS, + ChecksSerializer, + ChecksDeserializer, + SerializerFactory, + DataFrameConverter, + ChecksNormalizer, ) +from databricks.labs.dqx.utils import get_file_extension from databricks.labs.dqx.config_serializer import ConfigSerializer from databricks.labs.dqx.installer.mixins import InstallationMixin from databricks.labs.dqx.io import TABLE_PATTERN @@ -111,7 +110,7 @@ def load(self, config: TableChecksStorageConfig) -> list[dict]: if not self.spark.catalog.tableExists(config.location): raise NotFound(f"Checks table {config.location} does not exist in the workspace") rules_df = self.spark.read.table(config.location) - return serialize_checks_from_dataframe(rules_df, run_config_name=config.run_config_name) or [] + return DataFrameConverter.from_dataframe(rules_df, run_config_name=config.run_config_name) or [] @telemetry_logger("save_checks", "table") def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None: @@ -126,7 +125,7 @@ def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None: InvalidCheckError: If any check is invalid or unsupported. """ logger.info(f"Saving quality rules (checks) to table '{config.location}'") - rules_df = deserialize_checks_to_dataframe(self.spark, checks, run_config_name=config.run_config_name) + rules_df = DataFrameConverter.to_dataframe(self.spark, checks, run_config_name=config.run_config_name) rules_df.write.option("replaceWhere", f"run_config_name = '{config.run_config_name}'").saveAsTable( config.location, mode=config.mode ) @@ -244,7 +243,7 @@ def get_table_definition(schema_name: str, table_name: str) -> Table: def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) -> list[dict]: """ Normalize the checks to be compatible with the Lakebase table. - This includes normalizing Decimal values for JSON serialization and structuring + This includes normalizing special values for JSON serialization and structuring the checks for the Lakebase table schema. Args: @@ -254,9 +253,9 @@ def _normalize_checks(checks: list[dict], config: LakebaseChecksStorageConfig) - Returns: List of normalized dq rules (checks). """ - # First normalize Decimal values for JSON serialization - normalized_for_serialization = normalize_checks(checks) - + # First normalize special values for JSON serialization + normalized_for_serialization = ChecksNormalizer.normalize(checks) + # Then normalize the structure for Lakebase table normalized_checks = [] for check in normalized_for_serialization: @@ -356,8 +355,8 @@ def _load_checks_from_lakebase(self, config: LakebaseChecksStorageConfig, engine f"Make sure the profiler has run successfully and saved checks to this location." ) checks_dict = [dict(check) for check in checks] - # Denormalize Decimal markers back to Decimal objects - return denormalize_checks(checks_dict) + # Denormalize special markers back to objects + return ChecksNormalizer.denormalize(checks_dict) def _check_for_undefined_table_error(self, e: ProgrammingError, config: LakebaseChecksStorageConfig) -> NoReturn: """ @@ -490,8 +489,6 @@ def load(self, config: WorkspaceFileChecksStorageConfig) -> list[dict]: file_path = config.location logger.info(f"Loading quality rules (checks) from '{file_path}' in the workspace.") - deserializer = get_file_deserializer(file_path) - try: file_bytes = self.ws.workspace.download(file_path).read() file_content = file_bytes.decode("utf-8") @@ -499,7 +496,8 @@ def load(self, config: WorkspaceFileChecksStorageConfig) -> list[dict]: raise NotFound(f"Checks file {file_path} missing: {e}") from e try: - return deserializer(StringIO(file_content)) or [] + extension = get_file_extension(file_path) + return ChecksDeserializer.deserialize_from_file(extension, StringIO(file_content)) or [] except (yaml.YAMLError, json.JSONDecodeError) as e: raise InvalidCheckError(f"Invalid checks in file: {file_path}: {e}") from e @@ -517,7 +515,8 @@ def save(self, checks: list[dict], config: WorkspaceFileChecksStorageConfig) -> workspace_dir = str(file_path.parent) self.ws.workspace.mkdirs(workspace_dir) - content = serialize_checks_to_bytes(checks, file_path) + extension = get_file_extension(file_path) + content = ChecksSerializer.serialize_to_bytes(checks, extension) self.ws.workspace.upload(config.location, content, format=ImportFormat.AUTO, overwrite=True) @@ -543,11 +542,10 @@ def load(self, config: FileChecksStorageConfig) -> list[dict]: file_path = config.location logger.info(f"Loading quality rules (checks) from '{file_path}'.") - deserializer = get_file_deserializer(file_path) - try: + extension = get_file_extension(file_path) with open(file_path, "r", encoding="utf-8") as f: - return deserializer(f) or [] + return ChecksDeserializer.deserialize_from_file(extension, f) or [] except FileNotFoundError as e: raise FileNotFoundError(f"Checks file {file_path} missing: {e}") from e except (yaml.YAMLError, json.JSONDecodeError) as e: @@ -569,7 +567,8 @@ def save(self, checks: list[dict], config: FileChecksStorageConfig) -> None: os.makedirs(file_path.parent, exist_ok=True) try: - content = serialize_checks_to_bytes(checks, file_path) + extension = get_file_extension(file_path) + content = ChecksSerializer.serialize_to_bytes(checks, extension) with open(file_path, "wb") as file: file.write(content) except FileNotFoundError: @@ -707,8 +706,6 @@ def load(self, config: VolumeFileChecksStorageConfig) -> list[dict]: file_path = config.location logger.info(f"Loading quality rules (checks) from '{file_path}' in a volume.") - deserializer = get_file_deserializer(file_path) - try: file_download = self.ws.files.download(file_path) if not file_download.contents: @@ -722,7 +719,8 @@ def load(self, config: VolumeFileChecksStorageConfig) -> list[dict]: raise NotFound(f"Checks file {file_path} missing: {e}") from e try: - return deserializer(StringIO(file_content)) or [] + extension = get_file_extension(file_path) + return ChecksDeserializer.deserialize_from_file(extension, StringIO(file_content)) or [] except (yaml.YAMLError, json.JSONDecodeError) as e: raise InvalidCheckError(f"Invalid checks in file: {file_path}: {e}") from e @@ -740,7 +738,8 @@ def save(self, checks: list[dict], config: VolumeFileChecksStorageConfig) -> Non volume_dir = str(file_path.parent) self.ws.files.create_directory(volume_dir) - content = serialize_checks_to_bytes(checks, file_path) + extension = get_file_extension(file_path) + content = ChecksSerializer.serialize_to_bytes(checks, extension) binary_data = BytesIO(content) self.ws.files.upload(config.location, binary_data, overwrite=True) @@ -905,4 +904,6 @@ def is_table_location(location: str) -> bool: Returns: bool: True if the location is a valid table name and not a file path, False otherwise. """ - return bool(TABLE_PATTERN.match(location)) and not location.lower().endswith(tuple(FILE_SERIALIZERS.keys())) + return bool(TABLE_PATTERN.match(location)) and not location.lower().endswith( + SerializerFactory.get_supported_extensions() + ) diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index ff58f72e8..a5d8f3de6 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -2,7 +2,7 @@ from functools import cached_property from dataclasses import dataclass, field, asdict -from databricks.labs.dqx.checks_serializer import FILE_SERIALIZERS +from databricks.labs.dqx.checks_serializer import SerializerFactory from databricks.labs.dqx.errors import InvalidConfigError, InvalidParameterError __all__ = [ @@ -351,7 +351,7 @@ def __post_init__(self): raise InvalidParameterError("Invalid path: Path is missing a schema name") if len(parts) < 5 or not parts[4]: raise InvalidParameterError("Invalid path: Path is missing a volume name") - if len(parts) < 6 or not parts[-1].lower().endswith(tuple(FILE_SERIALIZERS.keys())): + if len(parts) < 6 or not parts[-1].lower().endswith(SerializerFactory.get_supported_extensions()): raise InvalidParameterError("Invalid path: Path must include a file name after the volume") diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 61de3bd69..71175d5a1 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -16,7 +16,7 @@ from databricks.labs.dqx.base import DQEngineBase, DQEngineCoreBase from databricks.labs.dqx.checks_resolver import resolve_custom_check_functions_from_path -from databricks.labs.dqx.checks_serializer import deserialize_checks +from databricks.labs.dqx.checks_serializer import ChecksDeserializer from databricks.labs.dqx.config_serializer import ConfigSerializer from databricks.labs.dqx.checks_storage import ( FileChecksStorageHandler, @@ -211,7 +211,7 @@ def apply_checks_by_metadata( A DataFrame with errors and warnings result columns and an optional Observation which tracks data quality summary metrics. Summary metrics are returned by any `DQEngine` with an `observer` specified. """ - dq_rule_checks = deserialize_checks(checks, custom_check_functions) + dq_rule_checks = ChecksDeserializer(custom_check_functions).deserialize(checks) return self.apply_checks(df, dq_rule_checks, ref_dfs) @@ -243,7 +243,7 @@ def apply_checks_by_metadata_and_split( Raises: InvalidCheckError: If any of the checks are invalid. """ - dq_rule_checks = deserialize_checks(checks, custom_check_functions) + dq_rule_checks = ChecksDeserializer(custom_check_functions).deserialize(checks) good_df, bad_df, *observations = self.apply_checks_and_split(df, dq_rule_checks, ref_dfs) diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index d7641d9f5..2c4a4044c 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -7,6 +7,7 @@ from importlib.util import find_spec from typing import Any from fnmatch import fnmatch +from pathlib import Path from pyspark.sql import Column @@ -141,6 +142,9 @@ def normalize_bound_args(val: Any) -> Any: Raises: TypeError: If a column type is unsupported. """ + if val is None: + return None + if isinstance(val, (list, tuple, set)): normalized = [normalize_bound_args(v) for v in val] return normalized @@ -214,25 +218,13 @@ def safe_json_load(value: str): Safely load a JSON string, returning the original value if it fails to parse. This allows to specify string value without a need to escape the quotes. - Also handles special Decimal format: {"__decimal__": "0.01"} is converted back to Decimal. - Args: value: The value to parse as JSON. - - Returns: - Parsed JSON value, or original value if parsing fails. Decimal markers are converted to Decimal objects. """ try: - parsed = json.loads(value) # load as json if possible - # Check if this is a Decimal marker and convert back to Decimal - if isinstance(parsed, dict) and "__decimal__" in parsed and len(parsed) == 1: - return Decimal(parsed["__decimal__"]) - return parsed + return json.loads(value) except json.JSONDecodeError: return value - except (ValueError, TypeError): - # If Decimal conversion fails, return the parsed value as-is - return value def safe_strip_file_from_path(path: str) -> str: @@ -482,3 +474,16 @@ def missing_required_packages(packages: list[str]) -> bool: True if any package is missing, False otherwise. """ return not all(find_spec(spec) for spec in packages) + + +def get_file_extension(file_path: str | os.PathLike) -> str: + """ + Extract file extension from a file path. + + Args: + file_path: File path as string or path-like object. + + Returns: + File extension (e.g., ".json", ".yaml", ".yml") or empty string if no extension. + """ + return Path(file_path).suffix diff --git a/tests/integration/test_build_rules.py b/tests/integration/test_build_rules.py index 2d256e374..fe5da3443 100644 --- a/tests/integration/test_build_rules.py +++ b/tests/integration/test_build_rules.py @@ -1,6 +1,6 @@ from decimal import Decimal -from databricks.labs.dqx.checks_serializer import deserialize_checks_to_dataframe, serialize_checks_from_dataframe +from databricks.labs.dqx.checks_serializer import DataFrameConverter SCHEMA = "a: int, b: int, c: int" @@ -131,8 +131,8 @@ def test_build_quality_rules_from_dataframe_round_trip(spark): }, ] - df = deserialize_checks_to_dataframe(spark, test_checks) - checks = serialize_checks_from_dataframe(df) + df = DataFrameConverter.to_dataframe(spark, test_checks) + checks = DataFrameConverter.from_dataframe(df) assert checks == test_checks, "The loaded checks do not match the expected checks." @@ -157,12 +157,12 @@ def test_build_quality_rules_from_dataframe_with_run_config(spark): "check": {"function": "is_not_less_than", "arguments": {"column": "test_col", "limit": "5"}}, }, ] - default_checks_df = deserialize_checks_to_dataframe(spark, default_checks) - workflow_checks_df = deserialize_checks_to_dataframe(spark, workflow_checks, run_config_name="workflow_001") + default_checks_df = DataFrameConverter.to_dataframe(spark, default_checks) + workflow_checks_df = DataFrameConverter.to_dataframe(spark, workflow_checks, run_config_name="workflow_001") df = default_checks_df.union(workflow_checks_df) - checks = serialize_checks_from_dataframe(df, run_config_name="workflow_001") + checks = DataFrameConverter.from_dataframe(df, run_config_name="workflow_001") assert checks == workflow_checks, "The loaded checks do not match the expected workflow checks." - checks = serialize_checks_from_dataframe(df) + checks = DataFrameConverter.from_dataframe(df) assert checks == default_checks, "The loaded checks do not match the expected default checks." diff --git a/tests/integration/test_summary_metrics.py b/tests/integration/test_summary_metrics.py index ab2e630dc..b2d7d9110 100644 --- a/tests/integration/test_summary_metrics.py +++ b/tests/integration/test_summary_metrics.py @@ -6,7 +6,7 @@ from chispa.dataframe_comparer import assert_df_equality # type: ignore from databricks.labs.dqx.config import InputConfig, OutputConfig, ExtraParams -from databricks.labs.dqx.checks_serializer import deserialize_checks +from databricks.labs.dqx.checks_serializer import ChecksDeserializer from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.io import save_dataframe_as_table from databricks.labs.dqx.metrics_observer import DQMetricsObserver, OBSERVATION_TABLE_SCHEMA @@ -77,7 +77,7 @@ def test_observer_metrics_before_action(ws, spark, apply_checks_method): ) if apply_checks_method == DQEngine.apply_checks: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) _, observation = dq_engine.apply_checks(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata: _, observation = dq_engine.apply_checks_by_metadata(test_df, TEST_CHECKS) @@ -105,7 +105,7 @@ def test_observer_metrics(ws, spark, apply_checks_method): ) if apply_checks_method == DQEngine.apply_checks: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) checked_df, observation = dq_engine.apply_checks(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata: checked_df, observation = dq_engine.apply_checks_by_metadata(test_df, TEST_CHECKS) @@ -178,7 +178,7 @@ def test_observer_custom_metrics(ws, spark, apply_checks_method): ) if apply_checks_method == DQEngine.apply_checks: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) checked_df, observation = dq_engine.apply_checks(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata: checked_df, observation = dq_engine.apply_checks_by_metadata(test_df, TEST_CHECKS) @@ -227,7 +227,7 @@ def test_engine_without_observer_no_metrics_saved(ws, spark, make_schema, make_r metrics_config = OutputConfig(location=metrics_table_name, mode="overwrite") if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, output_config=output_config, metrics_config=metrics_config ) @@ -999,7 +999,7 @@ def test_observer_metrics_output(skip_if_classic_compute, apply_checks_method, s checks_location = "fake.yml" if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, @@ -1155,7 +1155,7 @@ def test_observer_metrics_output_with_quarantine( checks_location = "fake.yml" if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, @@ -1304,7 +1304,7 @@ def test_save_results_in_table_batch_with_metrics( ) if apply_checks_method == DQEngine.apply_checks_and_split: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) output_df, quarantine_df, observation = dq_engine.apply_checks_and_split(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata_and_split: output_df, quarantine_df, observation = dq_engine.apply_checks_by_metadata_and_split(test_df, TEST_CHECKS) @@ -1425,7 +1425,7 @@ def test_save_results_in_table_streaming_with_metrics( test_df = spark.readStream.table(input_table_name) if apply_checks_method == DQEngine.apply_checks_and_split: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) output_df, quarantine_df, _ = dq_engine.apply_checks_and_split(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata_and_split: output_df, quarantine_df, _ = dq_engine.apply_checks_by_metadata_and_split(test_df, TEST_CHECKS) @@ -1576,7 +1576,7 @@ def test_streaming_observer_metrics_output(apply_checks_method, spark, ws, make_ test_df.write.mode("overwrite").saveAsTable(input_config.location) if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, @@ -1745,7 +1745,7 @@ def test_streaming_observer_metrics_output_and_quarantine( checks_location = "fake.yml" if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = deserialize_checks(TEST_CHECKS) + checks = ChecksDeserializer().deserialize(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, diff --git a/tests/unit/test_build_rules.py b/tests/unit/test_build_rules.py index 836a35173..951bfdc65 100644 --- a/tests/unit/test_build_rules.py +++ b/tests/unit/test_build_rules.py @@ -3,8 +3,6 @@ import logging import datetime import json -from pathlib import Path -from unittest.mock import Mock import yaml import pytest import pyspark.sql.functions as F @@ -43,11 +41,7 @@ register_rule, DQDatasetRule, ) -from databricks.labs.dqx.checks_serializer import ( - deserialize_checks, - serialize_checks, - serialize_checks_to_bytes, -) +from databricks.labs.dqx.checks_serializer import ChecksDeserializer, ChecksSerializer from databricks.labs.dqx.errors import InvalidCheckError, InvalidParameterError SCHEMA = "a: int, b: int, c: int" @@ -761,7 +755,7 @@ def test_build_rules_by_metadata(): }, ] - actual_rules = deserialize_checks(checks) + actual_rules = ChecksDeserializer().deserialize(checks) expected_rules = [ DQRowRule( @@ -1034,14 +1028,14 @@ def test_build_checks_by_metadata_when_check_spec_is_missing() -> None: checks: list[dict] = [{}] # missing check spec with pytest.raises(InvalidCheckError, match="'check' field is missing"): - deserialize_checks(checks) + ChecksDeserializer().deserialize(checks) def test_build_checks_by_metadata_when_function_spec_is_missing() -> None: checks: list[dict] = [{"check": {}}] # missing func spec with pytest.raises(InvalidCheckError, match="'function' field is missing in the 'check' block"): - deserialize_checks(checks) + ChecksDeserializer().deserialize(checks) def test_build_checks_by_metadata_when_arguments_are_missing(): @@ -1058,14 +1052,14 @@ def test_build_checks_by_metadata_when_arguments_are_missing(): InvalidCheckError, match="No arguments provided for function 'is_not_null_and_not_empty' in the 'arguments' block", ): - deserialize_checks(checks) + ChecksDeserializer().deserialize(checks) def test_build_checks_by_metadata_when_function_does_not_exist(): checks = [{"check": {"function": "function_does_not_exists", "arguments": {"column": "a"}}}] with pytest.raises(InvalidCheckError, match="function 'function_does_not_exists' is not defined"): - deserialize_checks(checks) + ChecksDeserializer().deserialize(checks) def test_build_checks_by_metadata_logging_debug_calls(caplog): @@ -1078,7 +1072,7 @@ def test_build_checks_by_metadata_logging_debug_calls(caplog): logger = logging.getLogger("databricks.labs.dqx.checks_resolver") logger.setLevel(logging.DEBUG) with caplog.at_level("DEBUG"): - deserialize_checks(checks) + ChecksDeserializer().deserialize(checks) assert "Resolving function: is_not_null_and_not_empty" in caplog.text @@ -1434,7 +1428,7 @@ def test_convert_dq_rules_to_metadata(): check_func_kwargs={"window_minutes": 60, "min_records_per_window": 1}, ), ] - actual_metadata = serialize_checks(checks) + actual_metadata = ChecksSerializer.serialize(checks) expected_metadata = [ { @@ -1729,7 +1723,7 @@ def test_convert_dq_rules_to_metadata(): def test_convert_dq_rules_to_metadata_when_empty() -> None: checks: list = [] - actual_metadata = serialize_checks(checks) + actual_metadata = ChecksSerializer.serialize(checks) expected_metadata: list[dict] = [] assert actual_metadata == expected_metadata @@ -1737,7 +1731,7 @@ def test_convert_dq_rules_to_metadata_when_empty() -> None: def test_convert_dq_rules_to_metadata_when_not_dq_rule() -> None: checks: list = [1] with pytest.raises(InvalidCheckError, match="Expected DQRule instance, got int"): - serialize_checks(checks) + ChecksSerializer.serialize(checks) def test_dq_rules_to_dict_when_column_expression_is_complex() -> None: @@ -1776,7 +1770,7 @@ def test_deserialize_checks_filter_takes_precedence_over_row_filter_in_arguments }, ] - actual_rules = deserialize_checks(checks) + actual_rules = ChecksDeserializer().deserialize(checks) # Verify that filter attribute is set and row_filter in arguments is overridden assert len(actual_rules) == 1 @@ -1804,7 +1798,7 @@ def test_deserialize_checks_row_filter_in_arguments_when_no_filter(): }, ] - actual_rules = deserialize_checks(checks) + actual_rules = ChecksDeserializer().deserialize(checks) # Verify that rule is created without filter attribute assert len(actual_rules) == 1 @@ -1917,10 +1911,10 @@ def test_metadata_round_trip_conversion_preserves_rules() -> None: ), ] - checks_dict = serialize_checks(checks) - converted_checks = deserialize_checks(checks_dict) + checks_dict = ChecksSerializer.serialize(checks) + converted_checks = ChecksDeserializer().deserialize(checks_dict) - assert serialize_checks(converted_checks) == serialize_checks(checks) + assert ChecksSerializer.serialize(converted_checks) == ChecksSerializer.serialize(checks) @pytest.mark.parametrize( @@ -1933,7 +1927,5 @@ def test_metadata_round_trip_conversion_preserves_rules() -> None: ], ) def test_serialize_checks_to_bytes(checks, file_path_suffix, expected_output): - mock_path = Mock(spec=Path) - mock_path.suffix = file_path_suffix - result = serialize_checks_to_bytes(checks, mock_path) + result = ChecksSerializer.serialize_to_bytes(checks, file_path_suffix) assert result == expected_output diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5502d03de..2537a8181 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1,5 +1,6 @@ from datetime import date, datetime from typing import Any +from pathlib import Path from unittest.mock import Mock import pyspark.sql.functions as F import pytest @@ -13,10 +14,11 @@ safe_json_load, get_columns_as_strings, is_simple_column_expression, - normalize_bound_args, safe_strip_file_from_path, missing_required_packages, + get_file_extension, ) +from databricks.labs.dqx.rule import normalize_bound_args from databricks.labs.dqx.errors import InvalidParameterError, InvalidConfigError from databricks.labs.dqx.config import InputConfig @@ -284,8 +286,8 @@ def test_safe_json_load_empty_string(): def test_safe_json_load_non_string_arg(): - result = safe_json_load(123) - assert result == 123 + with pytest.raises(TypeError): + safe_json_load(123) @pytest.mark.parametrize( @@ -336,6 +338,10 @@ def test_normalize_bound_args_unsupported_type(): normalize_bound_args({"a": 1}) +def test_normalize_bound_args_handle_none(): + assert normalize_bound_args(None) is None + + def test_get_reference_dataframes_with_missing_ref_tables(mock_spark) -> None: assert get_reference_dataframes(mock_spark, reference_tables={}) is None assert get_reference_dataframes(mock_spark, reference_tables=None) is None @@ -380,3 +386,31 @@ def test_safe_strip_file_from_path(path: str, expected: str): ) def test_missing_required_packages(packages, expected): assert missing_required_packages(packages) == expected + + +@pytest.mark.parametrize( + "file_path, expected_extension", + [ + ("/path/to/file.json", ".json"), + ("/path/to/file.yaml", ".yaml"), + ("/path/to/file.yml", ".yml"), + ("file.json", ".json"), + ("file.yaml", ".yaml"), + ("file.yml", ".yml"), + ("/path/to/file", ""), + ("file", ""), + ("/path/to/file.JSON", ".JSON"), # Case preserved, will be lowercased by serializer + ("/path/to/file.YAML", ".YAML"), + ("/path/to/file.with.multiple.dots.json", ".json"), + ("", ""), + ], +) +def test_get_file_extension(file_path: str, expected_extension: str): + """Test get_file_extension function with various file paths.""" + assert get_file_extension(file_path) == expected_extension + + +def test_get_file_extension_with_path_object(): + """Test get_file_extension function with Path object.""" + file_path = Path("/path/to/file.json") + assert get_file_extension(file_path) == ".json" From 59286413c2c402f4e42942aaac164469a9f23bbd Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 21:48:28 +0100 Subject: [PATCH 09/12] refactor --- src/databricks/labs/dqx/checks_serializer.py | 50 +++++++++++++++++++- src/databricks/labs/dqx/engine.py | 6 +-- tests/integration/test_summary_metrics.py | 22 ++++----- tests/unit/test_build_rules.py | 34 +++++++------ 4 files changed, 81 insertions(+), 31 deletions(-) diff --git a/src/databricks/labs/dqx/checks_serializer.py b/src/databricks/labs/dqx/checks_serializer.py index 9fe3dfaf0..798e6bcf8 100644 --- a/src/databricks/labs/dqx/checks_serializer.py +++ b/src/databricks/labs/dqx/checks_serializer.py @@ -425,8 +425,7 @@ def to_dataframe( Raises: InvalidCheckError: If any check is invalid or unsupported. """ - deserializer = ChecksDeserializer() - dq_rule_checks: list[DQRule] = deserializer.deserialize(checks) + dq_rule_checks: list[DQRule] = deserialize_checks(checks) dq_rule_rows = [] for dq_rule_check in dq_rule_checks: @@ -450,3 +449,50 @@ def to_dataframe( ] ) return spark.createDataFrame(dq_rule_rows, CHECKS_TABLE_SCHEMA) + + +def serialize_checks(checks: list[DQRule]) -> list[dict]: + """ + Converts a list of quality checks defined as *DQRule* objects to a list of quality checks + defined as Python dictionaries. + + This is a convenience user-friendly function that wraps ChecksSerializer.serialize. + + Args: + checks: List of DQRule instances to convert. + + Returns: + List of dictionaries representing the DQRule instances. + + Raises: + InvalidCheckError: If any item in the list is not a DQRule instance. + """ + return ChecksSerializer.serialize(checks) + + +def deserialize_checks(checks: list[dict], custom_checks: dict[str, Callable] | None = None) -> list[DQRule]: + """ + Converts a list of quality checks defined as Python dictionaries to a list of DQRule objects. + + This is a convenience user-friendly function that wraps ChecksDeserializer.deserialize. + + Args: + checks: list of dictionaries describing checks. Each check is a dictionary + consisting of following fields: + - *check* - Column expression to evaluate. This expression should return string value if it's evaluated to true + or *null* if it's evaluated to *false* + - *name* - name that will be given to a resulting column. Autogenerated if not provided + - *criticality* (optional) - possible values are *error* (data going only into "bad" dataframe), + and *warn* (data is going into both dataframes) + - *filter* (optional) - Expression for filtering data quality checks + - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. + custom_checks: Dictionary with custom check functions. + + Returns: + list of data quality check rules + + Raises: + InvalidCheckError: If any dictionary is invalid or unsupported. + """ + deserializer = ChecksDeserializer(custom_checks) + return deserializer.deserialize(checks) diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 71175d5a1..61de3bd69 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -16,7 +16,7 @@ from databricks.labs.dqx.base import DQEngineBase, DQEngineCoreBase from databricks.labs.dqx.checks_resolver import resolve_custom_check_functions_from_path -from databricks.labs.dqx.checks_serializer import ChecksDeserializer +from databricks.labs.dqx.checks_serializer import deserialize_checks from databricks.labs.dqx.config_serializer import ConfigSerializer from databricks.labs.dqx.checks_storage import ( FileChecksStorageHandler, @@ -211,7 +211,7 @@ def apply_checks_by_metadata( A DataFrame with errors and warnings result columns and an optional Observation which tracks data quality summary metrics. Summary metrics are returned by any `DQEngine` with an `observer` specified. """ - dq_rule_checks = ChecksDeserializer(custom_check_functions).deserialize(checks) + dq_rule_checks = deserialize_checks(checks, custom_check_functions) return self.apply_checks(df, dq_rule_checks, ref_dfs) @@ -243,7 +243,7 @@ def apply_checks_by_metadata_and_split( Raises: InvalidCheckError: If any of the checks are invalid. """ - dq_rule_checks = ChecksDeserializer(custom_check_functions).deserialize(checks) + dq_rule_checks = deserialize_checks(checks, custom_check_functions) good_df, bad_df, *observations = self.apply_checks_and_split(df, dq_rule_checks, ref_dfs) diff --git a/tests/integration/test_summary_metrics.py b/tests/integration/test_summary_metrics.py index b2d7d9110..ab2e630dc 100644 --- a/tests/integration/test_summary_metrics.py +++ b/tests/integration/test_summary_metrics.py @@ -6,7 +6,7 @@ from chispa.dataframe_comparer import assert_df_equality # type: ignore from databricks.labs.dqx.config import InputConfig, OutputConfig, ExtraParams -from databricks.labs.dqx.checks_serializer import ChecksDeserializer +from databricks.labs.dqx.checks_serializer import deserialize_checks from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.io import save_dataframe_as_table from databricks.labs.dqx.metrics_observer import DQMetricsObserver, OBSERVATION_TABLE_SCHEMA @@ -77,7 +77,7 @@ def test_observer_metrics_before_action(ws, spark, apply_checks_method): ) if apply_checks_method == DQEngine.apply_checks: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) _, observation = dq_engine.apply_checks(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata: _, observation = dq_engine.apply_checks_by_metadata(test_df, TEST_CHECKS) @@ -105,7 +105,7 @@ def test_observer_metrics(ws, spark, apply_checks_method): ) if apply_checks_method == DQEngine.apply_checks: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) checked_df, observation = dq_engine.apply_checks(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata: checked_df, observation = dq_engine.apply_checks_by_metadata(test_df, TEST_CHECKS) @@ -178,7 +178,7 @@ def test_observer_custom_metrics(ws, spark, apply_checks_method): ) if apply_checks_method == DQEngine.apply_checks: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) checked_df, observation = dq_engine.apply_checks(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata: checked_df, observation = dq_engine.apply_checks_by_metadata(test_df, TEST_CHECKS) @@ -227,7 +227,7 @@ def test_engine_without_observer_no_metrics_saved(ws, spark, make_schema, make_r metrics_config = OutputConfig(location=metrics_table_name, mode="overwrite") if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, output_config=output_config, metrics_config=metrics_config ) @@ -999,7 +999,7 @@ def test_observer_metrics_output(skip_if_classic_compute, apply_checks_method, s checks_location = "fake.yml" if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, @@ -1155,7 +1155,7 @@ def test_observer_metrics_output_with_quarantine( checks_location = "fake.yml" if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, @@ -1304,7 +1304,7 @@ def test_save_results_in_table_batch_with_metrics( ) if apply_checks_method == DQEngine.apply_checks_and_split: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) output_df, quarantine_df, observation = dq_engine.apply_checks_and_split(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata_and_split: output_df, quarantine_df, observation = dq_engine.apply_checks_by_metadata_and_split(test_df, TEST_CHECKS) @@ -1425,7 +1425,7 @@ def test_save_results_in_table_streaming_with_metrics( test_df = spark.readStream.table(input_table_name) if apply_checks_method == DQEngine.apply_checks_and_split: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) output_df, quarantine_df, _ = dq_engine.apply_checks_and_split(test_df, checks) elif apply_checks_method == DQEngine.apply_checks_by_metadata_and_split: output_df, quarantine_df, _ = dq_engine.apply_checks_by_metadata_and_split(test_df, TEST_CHECKS) @@ -1576,7 +1576,7 @@ def test_streaming_observer_metrics_output(apply_checks_method, spark, ws, make_ test_df.write.mode("overwrite").saveAsTable(input_config.location) if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, @@ -1745,7 +1745,7 @@ def test_streaming_observer_metrics_output_and_quarantine( checks_location = "fake.yml" if apply_checks_method == DQEngine.apply_checks_and_save_in_table: - checks = ChecksDeserializer().deserialize(TEST_CHECKS) + checks = deserialize_checks(TEST_CHECKS) dq_engine.apply_checks_and_save_in_table( checks=checks, input_config=input_config, diff --git a/tests/unit/test_build_rules.py b/tests/unit/test_build_rules.py index 951bfdc65..a78f44482 100644 --- a/tests/unit/test_build_rules.py +++ b/tests/unit/test_build_rules.py @@ -41,7 +41,11 @@ register_rule, DQDatasetRule, ) -from databricks.labs.dqx.checks_serializer import ChecksDeserializer, ChecksSerializer +from databricks.labs.dqx.checks_serializer import ( + ChecksSerializer, + serialize_checks, + deserialize_checks, +) from databricks.labs.dqx.errors import InvalidCheckError, InvalidParameterError SCHEMA = "a: int, b: int, c: int" @@ -755,7 +759,7 @@ def test_build_rules_by_metadata(): }, ] - actual_rules = ChecksDeserializer().deserialize(checks) + actual_rules = deserialize_checks(checks) expected_rules = [ DQRowRule( @@ -1028,14 +1032,14 @@ def test_build_checks_by_metadata_when_check_spec_is_missing() -> None: checks: list[dict] = [{}] # missing check spec with pytest.raises(InvalidCheckError, match="'check' field is missing"): - ChecksDeserializer().deserialize(checks) + deserialize_checks(checks) def test_build_checks_by_metadata_when_function_spec_is_missing() -> None: checks: list[dict] = [{"check": {}}] # missing func spec with pytest.raises(InvalidCheckError, match="'function' field is missing in the 'check' block"): - ChecksDeserializer().deserialize(checks) + deserialize_checks(checks) def test_build_checks_by_metadata_when_arguments_are_missing(): @@ -1052,14 +1056,14 @@ def test_build_checks_by_metadata_when_arguments_are_missing(): InvalidCheckError, match="No arguments provided for function 'is_not_null_and_not_empty' in the 'arguments' block", ): - ChecksDeserializer().deserialize(checks) + deserialize_checks(checks) def test_build_checks_by_metadata_when_function_does_not_exist(): checks = [{"check": {"function": "function_does_not_exists", "arguments": {"column": "a"}}}] with pytest.raises(InvalidCheckError, match="function 'function_does_not_exists' is not defined"): - ChecksDeserializer().deserialize(checks) + deserialize_checks(checks) def test_build_checks_by_metadata_logging_debug_calls(caplog): @@ -1072,7 +1076,7 @@ def test_build_checks_by_metadata_logging_debug_calls(caplog): logger = logging.getLogger("databricks.labs.dqx.checks_resolver") logger.setLevel(logging.DEBUG) with caplog.at_level("DEBUG"): - ChecksDeserializer().deserialize(checks) + deserialize_checks(checks) assert "Resolving function: is_not_null_and_not_empty" in caplog.text @@ -1428,7 +1432,7 @@ def test_convert_dq_rules_to_metadata(): check_func_kwargs={"window_minutes": 60, "min_records_per_window": 1}, ), ] - actual_metadata = ChecksSerializer.serialize(checks) + actual_metadata = serialize_checks(checks) expected_metadata = [ { @@ -1723,7 +1727,7 @@ def test_convert_dq_rules_to_metadata(): def test_convert_dq_rules_to_metadata_when_empty() -> None: checks: list = [] - actual_metadata = ChecksSerializer.serialize(checks) + actual_metadata = serialize_checks(checks) expected_metadata: list[dict] = [] assert actual_metadata == expected_metadata @@ -1731,7 +1735,7 @@ def test_convert_dq_rules_to_metadata_when_empty() -> None: def test_convert_dq_rules_to_metadata_when_not_dq_rule() -> None: checks: list = [1] with pytest.raises(InvalidCheckError, match="Expected DQRule instance, got int"): - ChecksSerializer.serialize(checks) + serialize_checks(checks) def test_dq_rules_to_dict_when_column_expression_is_complex() -> None: @@ -1770,7 +1774,7 @@ def test_deserialize_checks_filter_takes_precedence_over_row_filter_in_arguments }, ] - actual_rules = ChecksDeserializer().deserialize(checks) + actual_rules = deserialize_checks(checks) # Verify that filter attribute is set and row_filter in arguments is overridden assert len(actual_rules) == 1 @@ -1798,7 +1802,7 @@ def test_deserialize_checks_row_filter_in_arguments_when_no_filter(): }, ] - actual_rules = ChecksDeserializer().deserialize(checks) + actual_rules = deserialize_checks(checks) # Verify that rule is created without filter attribute assert len(actual_rules) == 1 @@ -1911,10 +1915,10 @@ def test_metadata_round_trip_conversion_preserves_rules() -> None: ), ] - checks_dict = ChecksSerializer.serialize(checks) - converted_checks = ChecksDeserializer().deserialize(checks_dict) + checks_dict = serialize_checks(checks) + converted_checks = deserialize_checks(checks_dict) - assert ChecksSerializer.serialize(converted_checks) == ChecksSerializer.serialize(checks) + assert serialize_checks(converted_checks) == serialize_checks(checks) @pytest.mark.parametrize( From 720bb73f8b7b1196922129355bb5f142108a4994 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 21:55:49 +0100 Subject: [PATCH 10/12] refactor --- src/databricks/labs/dqx/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index 2c4a4044c..2afdebd7d 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -137,7 +137,7 @@ def normalize_bound_args(val: Any) -> Any: val: Value or collection of values to normalize. Returns: - Normalized value or collection. Decimal values are converted to a special dict format. + Normalized value or collection. Raises: TypeError: If a column type is unsupported. From b0f2e98345bfb473a3e02afa23d7962c1eaa38bd Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 22:06:02 +0100 Subject: [PATCH 11/12] updated docstring --- src/databricks/labs/dqx/checks_serializer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/checks_serializer.py b/src/databricks/labs/dqx/checks_serializer.py index 798e6bcf8..50cded72d 100644 --- a/src/databricks/labs/dqx/checks_serializer.py +++ b/src/databricks/labs/dqx/checks_serializer.py @@ -74,7 +74,7 @@ def denormalize_value(val: Any) -> Any: def denormalize(checks: list[dict]) -> list[dict]: """ Recursively convert special markers back to objects after deserialization. - Converts {"__decimal__": "0.01"} back to Decimal("0.01"). + Converts special markers (e.g., __decimal__ format) back to Decimal objects. Args: checks: List of check dictionaries that may contain special markers. From 0bd296654de0a41dd0d54cd47f1af459786557b3 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 8 Feb 2026 22:08:27 +0100 Subject: [PATCH 12/12] updated docs --- docs/dqx/docs/guide/quality_checks_storage.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dqx/docs/guide/quality_checks_storage.mdx b/docs/dqx/docs/guide/quality_checks_storage.mdx index 225afeb1d..4bc101605 100644 --- a/docs/dqx/docs/guide/quality_checks_storage.mdx +++ b/docs/dqx/docs/guide/quality_checks_storage.mdx @@ -167,7 +167,7 @@ You can load quality checks from various storage locations. You can then apply t For the save methods to work, checks must be defined as a list of dictionaries. -If you create checks as a list of DQRule objects, you can convert them using the `serialize_checks` method, as described [here](/docs/reference/quality_checks/#converting-dqx-classes-to-metadata). +If you create checks as a list of DQRule objects, you can convert them using the `serialize_checks` function, as described [here](/docs/reference/quality_checks/#converting-dqx-classes-to-metadata).