[FEATURE]: Add optional rounding for `_aggr` check_funcs

[davidwanner-8451]

Is there an existing issue for this?

• I have searched the existing issues

Problem statement

I have some checks that utilize check_funcs.is_aggr_equal and have been trying to figure out how to round the aggregation as I keep getting warnings that are strictly due to rounding diffs

Example warning:

0: {"name": "ensure_numerator_sum_to_1", "message": "Sum value 0.9999999999999999 in column 'Numerator_Percent' per group of columns 'profile_type, concatenated_keys, Segmentation' is not equal to limit: 1", "columns": ["Numerator_Percent"], "filter": null, "function": "is_aggr_equal", "run_time": "2026-01-20T15:52:16.486Z", "run_id": "c4e9c0be-18bf-4cf5-b3a9-fefb77950a2c", "user_metadata": {}}

Here is the check itself:

  DQDatasetRule(  
    name="ensure_denom_sum_to_1",
    criticality="warn",
    check_func=check_funcs.is_aggr_equal,
    column = "Denominator_Percent",
    check_func_kwargs={
      "group_by": ['profile_type', 'concatenated_keys', 'Segmentation'],
      "aggr_type": "sum",
      "limit": 1}
  ),
]

Proposed Solution

It seems like the _is_aggr_compare function could add an optional rounding:

dqx/src/databricks/labs/dqx/check_funcs.py

Lines 2755 to 2793 in 45da959

	def _is_aggr_compare(
	column: str | Column,
	limit: int | float | str | Column,
	aggr_type: str,
	aggr_params: dict[str, Any] | None,
	group_by: list[str | Column] | None,
	row_filter: str | None,
	compare_op: Callable[[Column, Column], Column],
	compare_op_label: str,
	compare_op_name: str,
	) -> tuple[Column, Callable]:
	"""
	Helper to build aggregation comparison checks with a given operator.
	Constructs a condition and closure that verify whether an aggregation on a column
	(or groups of columns) satisfies a comparison against a limit (e.g., greater than).
	Args:
	column: Column name (str) or Column expression to aggregate.
	limit: Numeric value, column name, or SQL expression for the limit. String literals must be single quoted, e.g. 'string_value'.
	aggr_type: Aggregation type. Curated functions include 'count', 'sum', 'avg', 'min', 'max',
	'count_distinct', 'stddev', 'percentile', and more. Any Databricks built-in aggregate
	function is supported (will trigger a warning for non-curated functions).
	aggr_params: Optional dictionary of parameters for aggregate functions that require them
	(e.g., percentile functions need {"percentile": 0.95}).
	group_by: Optional list of columns or Column expressions to group by.
	row_filter: Optional SQL expression to filter rows before aggregation.
	compare_op: Comparison operator (e.g., operator.gt, operator.lt).
	compare_op_label: Human-readable label for the comparison (e.g., 'greater than').
	compare_op_name: Name identifier for the comparison (e.g., 'greater_than').
	Returns:
	A tuple of:
	- A Spark Column representing the condition for the aggregation check.
	- A closure that applies the aggregation check logic.
	Raises:
	InvalidParameterError: If an aggregate returns non-numeric types or is not found.
	MissingParameterError: If required parameters for specific aggregates are not provided.

Potentially something like this being added to the aggr_expr:

 # Apply rounding if specified
 if round_decimals is not None:
    aggr_expr = F.round(aggr_expr, round_decimals)

Additional Context

No response

[mwojtyczka]

For consistency, we should add this to is_aggr_equal, is_aggr_not_equal, is_equal_to, and is_not_equal_to check functions.
We provide tolerance in the compare_datasets function with an option to specify abs_tolerance or rel_tolerance. I think a similar approach should be used here if the user provides a float as limit.

Do you want to work on this?

[davidwanner-8451]

Yes I will take a stab, although this is my first attempt at contributing so might take me a bit / require some back and forth

[mwojtyczka]

Sure thing, please familiarize yourself with https://databrickslabs.github.io/dqx/docs/dev/contributing/

[davidwanner-8451]

@mwojtyczka - trying to think through the negative example.

Assume I have a basic example like this:

from databricks.labs.dqx.check_funcs import get_normalized_column_and_expr, get_limit_expr, _match_values_with_tolerance, make_condition, is_equal_to, is_not_equal_to
from databricks.labs.dqx.rule import DQRowRule

data = [
    ("A", 200),
    ("B",200),
    ("C", 199),  # VIOLATION!
]

df = spark.createDataFrame(
    data,
    schema=[
        "id",
        "sales",
    ],
)

# List of checks
dqx_check = [
  DQRowRule(
      name="equal_200",
      criticality="warn",
      check_func=is_equal_to,
      column="sales",
      check_func_kwargs = {"value": 200}
  ),
  DQRowRule(
      name="nequal_200",
      criticality="warn",
      check_func=is_not_equal_to,
      column="sales",
      check_func_kwargs = {"value": 199}
  ),
]

valid_df, quarantine_df = dq_engine.apply_checks_and_split(df, dqx_check)

quarantine_df.display()

The is_equal_to makes sense - we could add functionality for that func to make the 199 record as follows:

  DQRowRule(
      name="equal_within1_to_200",
      criticality="error",
      check_func=is_equal_to,
      column="sales",
      check_func_kwargs = {"value": 200, "abs_tolerance": 1.0 }
  ),

(I have that working now)

However, I am confused by what the not equal version is when tolerance is included. Do you have an example you can share given the above basic example?

[mwojtyczka]

Not equal is just the opposite.

dqx_check = [
  DQRowRule(
      criticality="warn",
      check_func=is_equal_to,
      column="sales",
      check_func_kwargs = {"value": 200, "abs_tolerance": 0.1 }  # values in [199.9, 200.1] are considered equal to 200; ABS(col_expr - value_expr) > abs_tolerance
  ),
  DQRowRule(
      criticality="warn",
      check_func=is_not_equal_to,
      column="sales",
      check_func_kwargs = {"value": 200, "abs_tolerance": 0.1 }  # values outside [199.9, 200.1] are considered not equal to 200; ABS(col_expr - value_expr) <= abs_tolerance
  ),
]

[davidwanner-8451]

Hi @ghanse and @mwojtyczka - making good progress, and have a working solution for the 4 functions is_aggr_equal, is_aggr_not_equal, is_equal_to, and is_not_equal_to that I have tested a local build of the wheel on.

I am wondering what testing is required? I did not see a specific area for check_funcs.py, which is the only area I am impacting. However, I did see that there is testing of these functions happening in a few areas:

• https://github.com/databrickslabs/dqx/blob/main/tests/integration/test_apply_checks.py
• https://github.com/databrickslabs/dqx/blob/main/tests/integration/test_dataset_checks.py

Is the idea just to expand on current testing of these functions to include some tolerance calls? Or is there more that should be done?

Thanks

[ghanse]

Yes, there are a few tests you in https://github.com/databrickslabs/dqx/blob/main/tests/integration/test_apply_checks.py that can be expanded with additional checks:

• test_apply_checks_all_checks_as_yaml
• test_apply_checks_al_checks_using_classes

Then update the existing is_aggr_... tests or add new tests in https://github.com/databrickslabs/dqx/blob/main/tests/integration/test_dataset_checks.py.

[mwojtyczka]

I think it should be enough to add tests in https://github.com/databrickslabs/dqx/blob/main/tests/integration/test_dataset_checks.py

This covers detailed cases. Please add new tests for this. Don't modify existing.

[davidwanner-8451]

Sounds good, thanks all - I am planning to add a function for each of the new four that tests the tolerance.

Here is an example for is_aggr_equal - does this seem acceptable?

def test_is_aggr_equal_with_tolerance(spark: SparkSession):
    """Test is_aggr_equal with absolute and relative tolerance parameters."""
    test_df = spark.createDataFrame(
        [
            ["a", 100],
            ["b", 100],
            ["c", 1],
        ],
        SCHEMA,
    )

    checks = [
        # Test 1: Absolute tolerance - avg(b) = 67.0, limit = 67.05, abs_tol = 0.1 -> PASS (diff = 0.05 <= 0.1)
        is_aggr_equal("b", limit=67.05, aggr_type="avg", abs_tolerance=0.1),
        # Test 2: Absolute tolerance - sum(b) = 201, limit = 201.5, abs_tol = 0.1 -> FAIL (diff = 0.5 > 0.1)
        is_aggr_equal("b", limit=201.5, aggr_type="sum", abs_tolerance=0.1),
        # Test 3: Relative tolerance - avg(b) = 67.0, limit = 68, rel_tol = 0.02 (2%) -> PASS (diff = 1, rel_tol = 1.36, 1 <= 1.36)
        is_aggr_equal("b", limit=68, aggr_type="avg", rel_tolerance=0.02),
        # Test 4: Relative tolerance - max(b) = 100, limit = 105, rel_tol = 0.01 (1%) -> FAIL (diff = 5, rel_tol = 1.05, 5 > 1.05)
        is_aggr_equal("b", limit=105, aggr_type="max", rel_tolerance=0.01),
        # Test 5: Both tolerances (OR logic) - min(b) = 1, limit = 1.4, abs_tol = 0.5, rel_tol = 0.01 -> PASS via abs (diff = 0.4 <= 0.5)
        is_aggr_equal("b", limit=1.4, aggr_type="min", abs_tolerance=0.5, rel_tolerance=0.01),
        # Test 6: Both tolerances (OR logic) - count(a) = 3, limit = 4, abs_tol = 0.5, rel_tol = 0.3 -> PASS via rel (diff = 1, rel_tol = 1.2, 1 <= 1.2)
        is_aggr_equal("a", limit=4, aggr_type="count", abs_tolerance=0.5, rel_tolerance=0.3),
        # Test 7: Both tolerances fail - sum(b) = 201, limit = 210, abs_tol = 0.5, rel_tol = 0.01 -> FAIL both (diff = 9 > 0.5 AND 9 > 2.1)
        is_aggr_equal("b", limit=210, aggr_type="sum", abs_tolerance=0.5, rel_tolerance=0.01),
        # Test 8: Exact match with zero tolerance - count(a) = 3, limit = 3, abs_tol = 0.0 -> PASS
        is_aggr_equal("a", limit=3, aggr_type="count", abs_tolerance=0.0),
    ]

    actual = _apply_checks(test_df, checks)

    expected_schema = (
        f"{SCHEMA}, "
        "b_avg_not_equal_to_limit STRING, "
        "b_sum_not_equal_to_limit STRING, "
        "b_avg_not_equal_to_limit STRING, "
        "b_max_not_equal_to_limit STRING, "
        "b_min_not_equal_to_limit STRING, "
        "a_count_not_equal_to_limit STRING, "
        "b_sum_not_equal_to_limit STRING, "
        "a_count_not_equal_to_limit STRING"
    )

    expected = spark.createDataFrame(
        [
            [
                "a",
                100,
                None, # Test 1: avg passes with abs tolerance
                "Sum value 201 in column 'b' is not equal to limit: 201.5", # Test 2: sum fails abs tolerance
                None, # Test 3: avg passes with rel tolerance
                "Max value 100 in column 'b' is not equal to limit: 105", # Test 4: max fails rel tolerance
                None,  # Test 5: min passes via abs tolerance (OR logic)
                None, # Test 6: count passes via rel tolerance (OR logic)
                "Sum value 201 in column 'b' is not equal to limit: 210",  # Test 7: sum fails both tolerances
                None, # Test 8: exact match passes
            ],
            [
                "b",
                100,
                None,
                "Sum value 201 in column 'b' is not equal to limit: 201.5",
                None,
                "Max value 100 in column 'b' is not equal to limit: 105",
                None,
                None,
                "Sum value 201 in column 'b' is not equal to limit: 210",
                None,
            ],
            [
                "c",
                1,
                None,  
                "Sum value 201 in column 'b' is not equal to limit: 201.5",  
                None,  
                "Max value 100 in column 'b' is not equal to limit: 105",  
                None, 
                None,  
                "Sum value 201 in column 'b' is not equal to limit: 210", 
                None,  
            ],
        ],
        expected_schema,
    )


    assert_df_equality(actual, expected, ignore_nullable=True)

[mwojtyczka]

looks good, you need to add some more edge cases, e.g. value in the column is null, value is of type float (your dataframe only has int).

[davidwanner-8451]

looks good, you need to add some more edge cases, e.g. value in the column is null, value is of type float (your dataframe only has int).

Sounds good, given the current _apply_checks function expects just two columns I will add a few different dataframes to handle the float (rather than adding more columns) just to not have to modify that function

Below:

dqx/tests/integration/test_dataset_checks.py

Lines 619 to 639 in ae7b0d3

	def _apply_checks(
	test_df: DataFrame,
	checks: list[tuple[Column, Callable]],
	ref_dfs: dict[str, DataFrame] | None = None,
	spark: SparkSession | None = None,
	) -> DataFrame:
	df_checked = test_df
	kwargs: dict[str, Any] = {}
	if spark:
	kwargs["spark"] = spark
	if ref_dfs:
	kwargs["ref_dfs"] = ref_dfs
	for _, apply_closure in checks:
	df_checked = apply_closure(df_checked, **kwargs)
	# Now simply select the conditions directly without adding withColumn
	condition_columns = [condition for (condition, _) in checks]
	actual = df_checked.select("a", "b", *condition_columns)
	return actual