From c28d041f8de83674e3ad276db6b96e73ebaf3c22 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 20 Mar 2026 18:21:27 -0400 Subject: [PATCH 01/21] feat: add message callable to DQRule for custom check messages Co-authored-by: Isaac --- src/databricks/labs/dqx/manager.py | 30 +++- src/databricks/labs/dqx/rule.py | 4 + tests/integration/test_custom_messages.py | 189 ++++++++++++++++++++++ tests/unit/test_custom_messages.py | 84 ++++++++++ 4 files changed, 306 insertions(+), 1 deletion(-) create mode 100644 tests/integration/test_custom_messages.py create mode 100644 tests/unit/test_custom_messages.py diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index 7ad14700b..6df958f58 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -1,3 +1,4 @@ +import inspect import logging from datetime import datetime from dataclasses import dataclass @@ -147,9 +148,11 @@ def _build_result_struct(self, condition: Column) -> Column: # or use literal run time if explicitly overridden run_time_expr = F.current_timestamp() if self.run_time_overwrite is None else F.lit(self.run_time_overwrite) + message_col = self._build_message_col(condition) + return F.struct( F.lit(self.check.name).alias("name"), - condition.alias("message"), + message_col.alias("message"), self.check.columns_as_string_expr.alias("columns"), F.lit(self.check.filter or None).cast("string").alias("filter"), F.lit(self.check.check_func.__name__).alias("function"), @@ -162,6 +165,31 @@ def _build_result_struct(self, condition: Column) -> Column: F.lit(self.rule_set_fingerprint).alias("rule_set_fingerprint"), ).cast(dq_result_item_schema) + def _build_message_col(self, condition: Column) -> Column: + """Build the message column, using a custom message callable if provided on the rule.""" + if self.check.message is None: + return condition + + # Build check_func_args dict from the check function's parameter names and stored args/kwargs + sig = inspect.signature(self.check.check_func) + param_names = list(sig.parameters.keys()) + check_func_args = dict(zip(param_names, self.check.check_func_args)) + check_func_args.update(self.check.check_func_kwargs) + + # Get the column value expression + column_value = F.col(self.check.column) if self.check.column else F.lit(None) + + # Call the user's message function + custom_message = self.check.message( + rule_name=self.check.name, + check_func_name=self.check.check_func.__name__, + check_func_args=check_func_args, + column_value=column_value, + ) + + # Keep the null/non-null semantics of the original condition but swap the message text + return F.when(condition.isNotNull(), custom_message).otherwise(F.lit(None).cast("string")) + def _get_invalid_cols_message(self) -> str: """ Returns invalid columns message containing info about invalid columns to check should be applied to or filter. diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 329ab815e..6ee034931 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -176,6 +176,7 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): check_func_args: list[Any] = field(default_factory=list) check_func_kwargs: dict[str, Any] = field(default_factory=dict) user_metadata: dict[str, str] | None = None + message: Callable | None = None def __post_init__(self): self._validate_rule_type(self.check_func) @@ -428,6 +429,7 @@ class DQForEachColRule(DQRuleTypeMixin): check_func_args: list[Any] = field(default_factory=list) check_func_kwargs: dict[str, Any] = field(default_factory=dict) user_metadata: dict[str, str] | None = None + message: Callable | None = None def get_rules(self) -> list[DQRule]: """Build a list of rules for a set of columns. @@ -453,6 +455,7 @@ def get_rules(self) -> list[DQRule]: criticality=self.criticality, filter=self.filter, user_metadata=self.user_metadata, + message=self.message, ) ) else: # default to row-level rule @@ -467,6 +470,7 @@ def get_rules(self) -> list[DQRule]: criticality=self.criticality, filter=self.filter, user_metadata=self.user_metadata, + message=self.message, ) ) return rules diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py new file mode 100644 index 000000000..97466184d --- /dev/null +++ b/tests/integration/test_custom_messages.py @@ -0,0 +1,189 @@ +"""Integration tests for custom message callable on DQRule.""" + +from typing import Any + +import pyspark.sql.functions as F +from pyspark.sql import Column + +from databricks.labs.dqx.engine import DQEngine +from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule +from databricks.labs.dqx import check_funcs +from tests.integration.conftest import ( + EXTRA_PARAMS, + assert_df_equality_ignore_fingerprints as assert_df_equality, + REPORTING_COLUMNS, + build_quality_violation, +) + + +SCHEMA = "a: int, b: int, c: int" +EXPECTED_SCHEMA = SCHEMA + REPORTING_COLUMNS + + +def _custom_message( + rule_name: str, + check_func_name: str, + check_func_args: dict[str, Any], # pylint: disable=unused-argument + column_value: Column, +) -> Column: + """Custom message that includes the rule name and check function name.""" + return F.concat(F.lit(f"Rule '{rule_name}' ({check_func_name}) failed for value: "), column_value.cast("string")) + + +def _static_custom_message( + rule_name: str, + check_func_name: str, # pylint: disable=unused-argument + check_func_args: dict[str, Any], # pylint: disable=unused-argument + column_value: Column, # pylint: disable=unused-argument +) -> Column: + """Custom message that returns a static string.""" + return F.lit(f"Custom error: {rule_name}") + + +def test_apply_checks_with_custom_message(ws, spark): + """Custom message callable should appear in the result DataFrame.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + message=_static_custom_message, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + + expected_errors = [ + [ + build_quality_violation( + name="c_not_null", + message="Custom error: c_not_null", + columns=["c"], + function="is_not_null", + ) + ] + ] + expected_df = spark.createDataFrame( + [[1, 3, None, expected_errors[0], None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_with_dynamic_custom_message(ws, spark): + """Custom message callable with column value should produce dynamic messages.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) + + rules = [ + DQRowRule( + name="b_not_null", + criticality="warn", + check_func=check_funcs.is_not_null, + column="b", + message=_custom_message, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + + expected_warnings = [ + [ + build_quality_violation( + name="b_not_null", + message="Rule 'b_not_null' (is_not_null) failed for value: null", + columns=["b"], + function="is_not_null", + ) + ] + ] + expected_df = spark.createDataFrame( + [[1, None, 3, None, expected_warnings[0]]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_without_custom_message_unchanged(ws, spark): + """Rules without custom message should produce the default message as before.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + + expected_errors = [ + [ + build_quality_violation( + name="c_not_null", + message="Column 'c' value is null", + columns=["c"], + function="is_not_null", + ) + ] + ] + expected_df = spark.createDataFrame( + [[1, 3, None, expected_errors[0], None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_passing_rows_have_no_custom_message(ws, spark): + """Rows that pass the check should not have a message even with a custom message callable.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 5]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + message=_static_custom_message, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + + expected_df = spark.createDataFrame( + [[1, 3, 5, None, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_for_each_col_rule_with_custom_message(ws, spark): + """DQForEachColRule with message should propagate to generated rules and appear in results.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[None, None, 3]], SCHEMA) + + rules = DQForEachColRule( + columns=["a", "b"], + check_func=check_funcs.is_not_null, + criticality="error", + message=_static_custom_message, + ).get_rules() + + checked_df = dq_engine.apply_checks(test_df, rules) + + # Both columns a and b are null, so both should have custom error messages + rows = checked_df.collect() + assert len(rows) == 1 + errors = rows[0]["_errors"] + assert len(errors) == 2 + messages = sorted([e["message"] for e in errors]) + # The auto-generated names include the column info + assert all(msg.startswith("Custom error:") for msg in messages) diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py new file mode 100644 index 000000000..e12b6a9d6 --- /dev/null +++ b/tests/unit/test_custom_messages.py @@ -0,0 +1,84 @@ +"""Tests for custom message callable on DQRule.""" + +import inspect +from typing import Any + +import pyspark.sql.functions as F +from pyspark.sql import Column + +from databricks.labs.dqx.check_funcs import is_not_null, is_not_null_and_not_empty +from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule + + +def _sample_message_callable( + rule_name: str, + check_func_name: str, # pylint: disable=unused-argument + check_func_args: dict[str, Any], # pylint: disable=unused-argument + column_value: Column, # pylint: disable=unused-argument +) -> Column: + return F.lit(f"Custom: {rule_name}") + + +def test_dq_row_rule_accepts_message_callable(): + """DQRowRule with message parameter should be creatable.""" + rule = DQRowRule( + check_func=is_not_null, + column="id", + message=_sample_message_callable, + ) + assert rule.message is _sample_message_callable + + +def test_dq_row_rule_message_defaults_to_none(): + """DQRowRule without message should default to None.""" + rule = DQRowRule(check_func=is_not_null, column="id") + assert rule.message is None + + +def test_dq_dataset_rule_accepts_message_callable(): + """DQDatasetRule without message should default to None.""" + rule = DQRowRule(check_func=is_not_null_and_not_empty, column="id") + assert rule.message is None + + +def test_dq_for_each_col_rule_propagates_message(): + """DQForEachColRule should propagate message to generated DQRowRules.""" + for_each_rule = DQForEachColRule( + columns=["a", "b"], + check_func=is_not_null, + message=_sample_message_callable, + ) + rules = for_each_rule.get_rules() + assert len(rules) == 2 + for rule in rules: + assert rule.message is _sample_message_callable + + +def test_dq_for_each_col_rule_message_defaults_to_none(): + """DQForEachColRule without message should produce rules with message=None.""" + for_each_rule = DQForEachColRule( + columns=["a", "b"], + check_func=is_not_null, + ) + rules = for_each_rule.get_rules() + for rule in rules: + assert rule.message is None + + +def test_message_callable_receives_correct_args(): + """The message callable signature should accept rule_name, check_func_name, check_func_args, column_value.""" + sig = inspect.signature(_sample_message_callable) + params = list(sig.parameters.keys()) + assert params == ["rule_name", "check_func_name", "check_func_args", "column_value"] + + +def test_dq_rule_to_dict_does_not_include_message(): + """to_dict() should not include the message callable since it is not serializable.""" + rule = DQRowRule( + check_func=is_not_null, + column="id", + name="id_not_null", + message=_sample_message_callable, + ) + rule_dict = rule.to_dict() + assert "message" not in rule_dict From 6721317978b88569a50a9f605bf4e5f7722c08bc Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 20 Mar 2026 21:40:30 -0400 Subject: [PATCH 02/21] Refactor --- src/databricks/labs/dqx/manager.py | 53 +++-- src/databricks/labs/dqx/rule.py | 8 +- tests/integration/test_custom_messages.py | 244 +++++++++++++++++----- 3 files changed, 234 insertions(+), 71 deletions(-) diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index 6df958f58..b216cf2b5 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -166,26 +166,47 @@ def _build_result_struct(self, condition: Column) -> Column: ).cast(dq_result_item_schema) def _build_message_col(self, condition: Column) -> Column: - """Build the message column, using a custom message callable if provided on the rule.""" + """ + Builds the message column, using the default message or a custom message callable if specified in the rule + definition. Allows use of the following arguments passed at execution time by the DQRuleManager: + + * *rule_name* is the name of the DQX rule + * *check_func_name* is the name of the DQX check function + * *check_func_args* is a dictionary of DQX check function arguments as key-value pairs + * *column_value* is the column value that failed the DQX check + + Args: + condition: Default DQX condition message returned by evaluating the DQX check function + + Returns: + The custom DQX condition message if specified in the rule definition, otherwise the default DQX condition + message + """ if self.check.message is None: return condition - # Build check_func_args dict from the check function's parameter names and stored args/kwargs - sig = inspect.signature(self.check.check_func) - param_names = list(sig.parameters.keys()) - check_func_args = dict(zip(param_names, self.check.check_func_args)) - check_func_args.update(self.check.check_func_kwargs) - - # Get the column value expression - column_value = F.col(self.check.column) if self.check.column else F.lit(None) - - # Call the user's message function - custom_message = self.check.message( - rule_name=self.check.name, - check_func_name=self.check.check_func.__name__, - check_func_args=check_func_args, - column_value=column_value, + check_func_signature = inspect.signature(self.check.check_func) + args, kwargs = self.check.prepare_check_func_args_and_kwargs() + bound = check_func_signature.bind_partial(*args, **kwargs) + check_func_args = dict(bound.arguments) + column_value = F.coalesce(F.col(self.check.column), F.lit("null")) if self.check.column else F.lit("null") + + default_message_func_args = { + "rule_name": self.check.name, + "check_func_name": self.check.check_func.__name__, + "check_func_args": check_func_args, + "column_value": column_value, + } + message_func_signature = inspect.signature(self.check.message) + message_func_accepts_keyword_args = any( + param.kind == inspect.Parameter.VAR_KEYWORD for param in message_func_signature.parameters.values() ) + message_func_args = { + arg: val + for arg, val in default_message_func_args.items() + if arg in list(message_func_signature.parameters.keys()) or message_func_accepts_keyword_args + } + custom_message = self.check.message(**message_func_args) # Keep the null/non-null semantics of the original condition but swap the message text return F.when(condition.isNotNull(), custom_message).otherwise(F.lit(None).cast("string")) diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 6ee034931..35fc7a4bb 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -165,6 +165,12 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): * *check_func_args* (optional) - Positional arguments for the check function (excluding *column*). * *check_func_kwargs* (optional) - Keyword arguments for the check function (excluding *column*). * *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. + * *custom_message_func* - A user-defined function that returns a message when the check fails. The function should + return a Spark Column and can optionally accept the following keyword arguments: + - *rule_name* is the name of the DQX rule + - *check_func_name* is the name of the DQX check function + - *check_func_args* is a dictionary of check function arguments as key-value pairs + - *column_value* is the column value that failed the DQX check """ check_func: Callable @@ -176,7 +182,7 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): check_func_args: list[Any] = field(default_factory=list) check_func_kwargs: dict[str, Any] = field(default_factory=dict) user_metadata: dict[str, str] | None = None - message: Callable | None = None + message: Callable[..., Column] | None = None def __post_init__(self): self._validate_rule_type(self.check_func) diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index 97466184d..24bf9497a 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -20,26 +20,6 @@ EXPECTED_SCHEMA = SCHEMA + REPORTING_COLUMNS -def _custom_message( - rule_name: str, - check_func_name: str, - check_func_args: dict[str, Any], # pylint: disable=unused-argument - column_value: Column, -) -> Column: - """Custom message that includes the rule name and check function name.""" - return F.concat(F.lit(f"Rule '{rule_name}' ({check_func_name}) failed for value: "), column_value.cast("string")) - - -def _static_custom_message( - rule_name: str, - check_func_name: str, # pylint: disable=unused-argument - check_func_args: dict[str, Any], # pylint: disable=unused-argument - column_value: Column, # pylint: disable=unused-argument -) -> Column: - """Custom message that returns a static string.""" - return F.lit(f"Custom error: {rule_name}") - - def test_apply_checks_with_custom_message(ws, spark): """Custom message callable should appear in the result DataFrame.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) @@ -56,19 +36,16 @@ def test_apply_checks_with_custom_message(ws, spark): ] checked_df = dq_engine.apply_checks(test_df, rules) - expected_errors = [ - [ - build_quality_violation( - name="c_not_null", - message="Custom error: c_not_null", - columns=["c"], - function="is_not_null", - ) - ] + build_quality_violation( + name="c_not_null", + message="Custom error: c_not_null", + columns=["c"], + function="is_not_null", + ) ] expected_df = spark.createDataFrame( - [[1, 3, None, expected_errors[0], None]], + [[1, 3, None, expected_errors, None]], EXPECTED_SCHEMA, ) assert_df_equality(checked_df, expected_df) @@ -90,19 +67,16 @@ def test_apply_checks_with_dynamic_custom_message(ws, spark): ] checked_df = dq_engine.apply_checks(test_df, rules) - expected_warnings = [ - [ - build_quality_violation( - name="b_not_null", - message="Rule 'b_not_null' (is_not_null) failed for value: null", - columns=["b"], - function="is_not_null", - ) - ] + build_quality_violation( + name="b_not_null", + message="Rule 'b_not_null' (is_not_null) failed for value: null", + columns=["b"], + function="is_not_null", + ) ] expected_df = spark.createDataFrame( - [[1, None, 3, None, expected_warnings[0]]], + [[1, None, 3, None, expected_warnings]], EXPECTED_SCHEMA, ) assert_df_equality(checked_df, expected_df) @@ -123,19 +97,16 @@ def test_apply_checks_without_custom_message_unchanged(ws, spark): ] checked_df = dq_engine.apply_checks(test_df, rules) - expected_errors = [ - [ - build_quality_violation( - name="c_not_null", - message="Column 'c' value is null", - columns=["c"], - function="is_not_null", - ) - ] + build_quality_violation( + name="c_not_null", + message="Column 'c' value is null", + columns=["c"], + function="is_not_null", + ) ] expected_df = spark.createDataFrame( - [[1, 3, None, expected_errors[0], None]], + [[1, 3, None, expected_errors, None]], EXPECTED_SCHEMA, ) assert_df_equality(checked_df, expected_df) @@ -157,7 +128,6 @@ def test_apply_checks_passing_rows_have_no_custom_message(ws, spark): ] checked_df = dq_engine.apply_checks(test_df, rules) - expected_df = spark.createDataFrame( [[1, 3, 5, None, None]], EXPECTED_SCHEMA, @@ -178,12 +148,178 @@ def test_for_each_col_rule_with_custom_message(ws, spark): ).get_rules() checked_df = dq_engine.apply_checks(test_df, rules) - - # Both columns a and b are null, so both should have custom error messages rows = checked_df.collect() assert len(rows) == 1 errors = rows[0]["_errors"] assert len(errors) == 2 messages = sorted([e["message"] for e in errors]) - # The auto-generated names include the column info assert all(msg.startswith("Custom error:") for msg in messages) + + +def test_apply_checks_with_kwargs_only_custom_message(ws, spark): + """Custom message callable with **kwargs should receive all supported context args.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) + + rules = [ + DQRowRule( + name="b_not_null", + criticality="warn", + check_func=check_funcs.is_not_null, + column="b", + message=_kwargs_only_custom_message, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_warnings = [ + build_quality_violation( + name="b_not_null", + message="kwargs keys: check_func_args,check_func_name,column_value,rule_name", + columns=["b"], + function="is_not_null", + ) + ] + expected_df = spark.createDataFrame( + [[1, None, 3, None, expected_warnings]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_with_rule_name_only_custom_message(ws, spark): + """Custom message callable with a single named argument should work.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) + + rules = [ + DQRowRule( + name="b_not_null", + criticality="warn", + check_func=check_funcs.is_not_null, + column="b", + message=_rule_name_only_custom_message, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_warnings = [ + build_quality_violation( + name="b_not_null", + message="Only rule: b_not_null", + columns=["b"], + function="is_not_null", + ) + ] + expected_df = spark.createDataFrame( + [[1, None, 3, None, expected_warnings]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_custom_message_check_func_args_include_column(ws, spark): + """check_func_args passed to custom message should include the rule 'column' argument.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) + + rules = [ + DQRowRule( + name="b_not_null", + criticality="warn", + check_func=check_funcs.is_not_null, + column="b", + message=_message_from_column_arg, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_warnings = [ + build_quality_violation( + name="b_not_null", + message="column arg: b", + columns=["b"], + function="is_not_null", + ) + ] + expected_df = spark.createDataFrame( + [[1, None, 3, None, expected_warnings]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_custom_message_check_func_args_include_columns(ws, spark): + """check_func_args passed to custom message should include the rule 'columns' argument.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[None, 2, 3]], SCHEMA) + + rules = [ + DQRowRule( + name="ab_any_null", + criticality="error", + check_func=_any_null_in_columns, + columns=["a", "b"], + message=_message_from_columns_arg, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_errors = [ + build_quality_violation( + name="ab_any_null", + message="columns arg: a,b", + columns=["a", "b"], + function="_any_null_in_columns", + ) + ] + expected_df = spark.createDataFrame( + [[None, 2, 3, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def _custom_message( + rule_name: str, + check_func_name: str, + column_value: Column, +) -> Column: + """Custom message that includes the rule name and check function name.""" + return F.concat(F.lit(f"Rule '{rule_name}' ({check_func_name}) failed for value: "), column_value.cast("string")) + + +def _static_custom_message( + rule_name: str, +) -> Column: + """Custom message that returns a static string.""" + return F.lit(f"Custom error: {rule_name}") + + +def _kwargs_only_custom_message(**kwargs: Any) -> Column: + """Custom message callable that accepts only **kwargs.""" + return F.lit(f"kwargs keys: {','.join(sorted(kwargs.keys()))}") + + +def _rule_name_only_custom_message(rule_name: str) -> Column: + """Custom message callable that accepts a single named argument.""" + return F.lit(f"Only rule: {rule_name}") + + +def _message_from_column_arg(check_func_args: dict[str, Any]) -> Column: + """Custom message that reads 'column' from check_func_args.""" + return F.lit(f"column arg: {check_func_args.get('column')}") + + +def _message_from_columns_arg(check_func_args: dict[str, Any]) -> Column: + """Custom message that reads 'columns' from check_func_args.""" + columns = check_func_args.get("columns", []) + return F.lit(f"columns arg: {','.join(columns)}") + + +def _any_null_in_columns(columns: list[str]) -> Column: + """Fails if any provided column value is null.""" + condition = F.col(columns[0]).isNull() + for col_name in columns[1:]: + condition = condition | F.col(col_name).isNull() + return check_funcs.make_condition(condition, "one or more columns are null", "any_null_in_columns") From ab9a4ff5ca15260cdcf6d879dd45ee287b64e653 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 23 Mar 2026 11:19:31 -0400 Subject: [PATCH 03/21] Update documentation and tests --- docs/dqx/docs/reference/quality_checks.mdx | 52 ++++++++++++++++++++++ src/databricks/labs/dqx/manager.py | 2 +- tests/integration/test_custom_messages.py | 3 +- 3 files changed, 55 insertions(+), 2 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index ba4cd8df9..05e97a6ce 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -3881,6 +3881,58 @@ Using DQX classes: When using dataset-level checks, the top-level `filter` condition is pushed down as `row_filter` to the check function and applied before aggregation, ensuring that the check operates only on the relevant subset of rows rather than on the aggregated results. +## Customizing check messages + +Messages created by DQX rules can be customized with user-defined message functions. + +First, define a message function that returns a string-valued Spark `Column`. Set the `message` attribute of any `DQRule` to use this function and override the check function's default message. + +At runtime, DQX passes the following optional keyword arguments to the message function: +- `rule_name`: name of the rule +- `check_func_name`: check function name +- `check_func_args`: effective check arguments as a dictionary (includes `column` or `columns` when used) +- `column_value`: Spark column expression for the failed row value (for single-column rules) + + +If your message concatenates `column_value`, use `coalesce` to avoid null messages. In Spark, `concat(..., null)` returns `null`. + + + + + ```python + import pyspark.sql.functions as F + from pyspark.sql import Column + from databricks.labs.dqx import check_funcs + from databricks.labs.dqx.rule import DQRowRule + + # define a custom message function + def custom_message( + rule_name: str, + check_func_name: str, + check_func_args: dict, + column_value: Column, + ) -> Column: + value_str = F.coalesce(column_value.cast("string"), F.lit("null")) + return F.concat( + F.lit(f"[{rule_name}] {check_func_name} failed for value="), + value_str, + F.lit(f", args={check_func_args}"), + ) + + # create a check with the custom message + checks = [ + DQRowRule( + name="email_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="email", + message=custom_message, + ) + ] + ``` + + + ## Converting checks between formats In DQX, checks can be defined either as Python classes or YAML declarations. When using YAML, the files are first parsed into dictionaries and then transformed into DQX class instances under the hood. Since both formats share the same internal structure, they are interchangeable and can be safely converted between one another. diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index b216cf2b5..6772ea002 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -189,7 +189,7 @@ def _build_message_col(self, condition: Column) -> Column: args, kwargs = self.check.prepare_check_func_args_and_kwargs() bound = check_func_signature.bind_partial(*args, **kwargs) check_func_args = dict(bound.arguments) - column_value = F.coalesce(F.col(self.check.column), F.lit("null")) if self.check.column else F.lit("null") + column_value = F.col(self.check.column).cast("string") if self.check.column else F.lit(None) default_message_func_args = { "rule_name": self.check.name, diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index 24bf9497a..b90c61afd 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -286,7 +286,8 @@ def _custom_message( column_value: Column, ) -> Column: """Custom message that includes the rule name and check function name.""" - return F.concat(F.lit(f"Rule '{rule_name}' ({check_func_name}) failed for value: "), column_value.cast("string")) + value_str = F.coalesce(column_value.cast("string"), F.lit("null")) + return F.concat(F.lit(f"Rule '{rule_name}' ({check_func_name}) failed for value: "), value_str) def _static_custom_message( From 05baafb14ed16af07e522cdf40996b189fa15498 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 23 Mar 2026 11:27:00 -0400 Subject: [PATCH 04/21] Fix tests --- tests/unit/test_custom_messages.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index e12b6a9d6..5a1af1b22 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -10,12 +10,8 @@ from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule -def _sample_message_callable( - rule_name: str, - check_func_name: str, # pylint: disable=unused-argument - check_func_args: dict[str, Any], # pylint: disable=unused-argument - column_value: Column, # pylint: disable=unused-argument -) -> Column: +def _sample_message_callable(rule_name: str) -> Column: + """Custom message function that returns a string-valued column.""" return F.lit(f"Custom: {rule_name}") From 468f229d167a128e4c4ffa1293ee6e0f8e723661 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 23 Mar 2026 11:43:35 -0400 Subject: [PATCH 05/21] Fix tests --- tests/unit/test_custom_messages.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index 5a1af1b22..30c305552 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -1,7 +1,6 @@ """Tests for custom message callable on DQRule.""" import inspect -from typing import Any import pyspark.sql.functions as F from pyspark.sql import Column @@ -65,7 +64,7 @@ def test_message_callable_receives_correct_args(): """The message callable signature should accept rule_name, check_func_name, check_func_args, column_value.""" sig = inspect.signature(_sample_message_callable) params = list(sig.parameters.keys()) - assert params == ["rule_name", "check_func_name", "check_func_args", "column_value"] + assert params == ["rule_name"] def test_dq_rule_to_dict_does_not_include_message(): From c08cf3197ff00e509d1eb3fa64ba70490df8e023 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 16 Apr 2026 12:33:06 -0400 Subject: [PATCH 06/21] Refactor to SQL expressions for custom check messages --- docs/dqx/docs/reference/quality_checks.mdx | 67 +++-- src/databricks/labs/dqx/checks_serializer.py | 4 + src/databricks/labs/dqx/manager.py | 47 ++-- src/databricks/labs/dqx/rule.py | 17 +- tests/integration/test_custom_messages.py | 265 ++++++++++--------- tests/unit/test_custom_messages.py | 48 ++-- 6 files changed, 238 insertions(+), 210 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 90bf7b7ad..f58b2f33e 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -3899,54 +3899,67 @@ When using dataset-level checks, the top-level `filter` condition is pushed down ## Customizing check messages -Messages created by DQX rules can be customized with user-defined message functions. +Users can provide a Spark SQL expression to customize messages created by DQX failures. Set the `message` attribute of any `DQRule` to a SQL expression that returns a string value. This overrides the check function's default message. -First, define a message function that returns a string-valued Spark `Column`. Set the `message` attribute of any `DQRule` to use this function and override the check function's default message. - -At runtime, DQX passes the following optional keyword arguments to the message function: -- `rule_name`: name of the rule -- `check_func_name`: check function name -- `check_func_args`: effective check arguments as a dictionary (includes `column` or `columns` when used) -- `column_value`: Spark column expression for the failed row value (for single-column rules) +Custom SQL expressions support the following placeholders, which are substituted before evaluation: +- `{rule_name}` is replaced with the name of the DQX rule +- `{check_func_name}` is replaced with the name of the DQX check function +- `{column_value}` is replaced with the observed column value (cast to string). This can only be used for single-column values. -If your message concatenates `column_value`, use `coalesce` to avoid null messages. In Spark, `concat(..., null)` returns `null`. +If your message references `{column_value}`, wrap it with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. ```python - import pyspark.sql.functions as F - from pyspark.sql import Column from databricks.labs.dqx import check_funcs from databricks.labs.dqx.rule import DQRowRule - # define a custom message function - def custom_message( - rule_name: str, - check_func_name: str, - check_func_args: dict, - column_value: Column, - ) -> Column: - value_str = F.coalesce(column_value.cast("string"), F.lit("null")) - return F.concat( - F.lit(f"[{rule_name}] {check_func_name} failed for value="), - value_str, - F.lit(f", args={check_func_args}"), - ) - - # create a check with the custom message + # simple static message checks = [ DQRowRule( name="email_not_null", criticality="error", check_func=check_funcs.is_not_null, column="email", - message=custom_message, + message="'Email must not be null'", + ) + ] + + # dynamic message with column value + checks = [ + DQRowRule( + name="age_positive", + criticality="error", + check_func=check_funcs.is_greater_than, + column="age", + check_func_kwargs={"limit": 0}, + message="concat('{rule_name}', ': value ', coalesce(cast({column_value} as string), 'null'), ' is not valid')", ) ] ``` + + ```yaml + - name: email_not_null + criticality: error + message: "'Email must not be null'" + check: + function: is_not_null + arguments: + column: email + + - name: age_positive + criticality: error + message: "concat('{rule_name}', ': value ', coalesce(cast({column_value} as string), 'null'), ' is not valid')" + check: + function: is_greater_than + arguments: + column: age + limit: 0 + ``` + ## Converting checks between formats diff --git a/src/databricks/labs/dqx/checks_serializer.py b/src/databricks/labs/dqx/checks_serializer.py index 8167817d4..a3536fe38 100644 --- a/src/databricks/labs/dqx/checks_serializer.py +++ b/src/databricks/labs/dqx/checks_serializer.py @@ -267,6 +267,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: criticality = check_def.get("criticality", "error") filter_str = check_def.get("filter") user_metadata = check_def.get("user_metadata") + message = check_def.get("message") # Exclude `column` and `columns` from check_func_kwargs # as these are always included in the check function call @@ -282,6 +283,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: filter=filter_str, check_func_kwargs=check_func_kwargs, user_metadata=user_metadata, + message=message, ).get_rules() else: rule_type = CHECK_FUNC_REGISTRY.get(func_name) @@ -296,6 +298,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: criticality=criticality, filter=filter_str, user_metadata=user_metadata, + message=message, ) ) else: # default to row-level rule @@ -309,6 +312,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: criticality=criticality, filter=filter_str, user_metadata=user_metadata, + message=message, ) ) diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index e67aa5e25..2eb2b788f 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -1,4 +1,3 @@ -import inspect import logging from datetime import datetime from dataclasses import dataclass @@ -172,13 +171,13 @@ def _build_result_struct(self, condition: Column, skipped: bool = False) -> Colu def _build_message_col(self, condition: Column) -> Column: """ - Builds the message column, using the default message or a custom message callable if specified in the rule - definition. Allows use of the following arguments passed at execution time by the DQRuleManager: + Builds the message column, using the default message or a custom Spark SQL expression if + specified in the rule definition. The SQL expression supports placeholders that are + substituted before evaluation: - * *rule_name* is the name of the DQX rule - * *check_func_name* is the name of the DQX check function - * *check_func_args* is a dictionary of DQX check function arguments as key-value pairs - * *column_value* is the column value that failed the DQX check + * '{rule_name}' is replaced with the quoted name of the DQX rule + * '{check_func_name}' is replaced with the quoted name of the DQX check function + * '{column_value}' is replaced with a SQL reference to the column value (cast to string) Args: condition: Default DQX condition message returned by evaluating the DQX check function @@ -190,30 +189,16 @@ def _build_message_col(self, condition: Column) -> Column: if self.check.message is None: return condition - check_func_signature = inspect.signature(self.check.check_func) - args, kwargs = self.check.prepare_check_func_args_and_kwargs() - bound = check_func_signature.bind_partial(*args, **kwargs) - check_func_args = dict(bound.arguments) - column_value = F.col(self.check.column).cast("string") if self.check.column else F.lit(None) - - default_message_func_args = { - "rule_name": self.check.name, - "check_func_name": self.check.check_func.__name__, - "check_func_args": check_func_args, - "column_value": column_value, - } - message_func_signature = inspect.signature(self.check.message) - message_func_accepts_keyword_args = any( - param.kind == inspect.Parameter.VAR_KEYWORD for param in message_func_signature.parameters.values() - ) - message_func_args = { - arg: val - for arg, val in default_message_func_args.items() - if arg in list(message_func_signature.parameters.keys()) or message_func_accepts_keyword_args - } - custom_message = self.check.message(**message_func_args) - - # Keep the null/non-null semantics of the original condition but swap the message text + column_value_sql = f"CAST(`{self.check.column}` AS STRING)" if self.check.column else "NULL" + escaped_check_name = self.check.name.replace("'", "\\'") + escaped_check_function = self.check.check_func.__name__.replace("'", "\\'") + + message_expr = self.check.message + message_expr = message_expr.replace("{rule_name}", escaped_check_name) + message_expr = message_expr.replace("{check_func_name}", escaped_check_function) + message_expr = message_expr.replace("{column_value}", column_value_sql) + + custom_message = F.expr(message_expr) return F.when(condition.isNotNull(), custom_message).otherwise(F.lit(None).cast("string")) def _get_invalid_cols_message(self) -> str: diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 35fc7a4bb..37eecc8fe 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -165,12 +165,11 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): * *check_func_args* (optional) - Positional arguments for the check function (excluding *column*). * *check_func_kwargs* (optional) - Keyword arguments for the check function (excluding *column*). * *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. - * *custom_message_func* - A user-defined function that returns a message when the check fails. The function should - return a Spark Column and can optionally accept the following keyword arguments: - - *rule_name* is the name of the DQX rule - - *check_func_name* is the name of the DQX check function - - *check_func_args* is a dictionary of check function arguments as key-value pairs - - *column_value* is the column value that failed the DQX check + * *message* (optional) - User-defined Spark SQL expression used as the check failure message. Supports the + following placeholders which are substituted when checks are evaluated: + - `{rule_name}` is replaced with the name of the DQX rule + - `{check_func_name}` is replaced with the name of the DQX check function + - `{column_value}` is replaced with a SQL expression referencing the column value (cast to string) """ check_func: Callable @@ -182,7 +181,7 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): check_func_args: list[Any] = field(default_factory=list) check_func_kwargs: dict[str, Any] = field(default_factory=dict) user_metadata: dict[str, str] | None = None - message: Callable[..., Column] | None = None + message: str | None = None def __post_init__(self): self._validate_rule_type(self.check_func) @@ -266,6 +265,8 @@ def to_dict(self) -> dict: if self.user_metadata: metadata["user_metadata"] = self.user_metadata + if self.message: + metadata["message"] = self.message return metadata def _initialize_column_if_missing(self): @@ -435,7 +436,7 @@ class DQForEachColRule(DQRuleTypeMixin): check_func_args: list[Any] = field(default_factory=list) check_func_kwargs: dict[str, Any] = field(default_factory=dict) user_metadata: dict[str, str] | None = None - message: Callable | None = None + message: str | None = None def get_rules(self) -> list[DQRule]: """Build a list of rules for a set of columns. diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index b90c61afd..7384f092a 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -1,9 +1,4 @@ -"""Integration tests for custom message callable on DQRule.""" - -from typing import Any - -import pyspark.sql.functions as F -from pyspark.sql import Column +"""Integration tests for custom message SQL expressions on DQRule.""" from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule @@ -15,13 +10,12 @@ build_quality_violation, ) - SCHEMA = "a: int, b: int, c: int" EXPECTED_SCHEMA = SCHEMA + REPORTING_COLUMNS -def test_apply_checks_with_custom_message(ws, spark): - """Custom message callable should appear in the result DataFrame.""" +def test_apply_checks_with_static_custom_message(ws, spark): + """A plain SQL literal message should appear in the result DataFrame.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) @@ -31,7 +25,7 @@ def test_apply_checks_with_custom_message(ws, spark): criticality="error", check_func=check_funcs.is_not_null, column="c", - message=_static_custom_message, + message="'Custom error: {rule_name}'", ), ] @@ -51,8 +45,8 @@ def test_apply_checks_with_custom_message(ws, spark): assert_df_equality(checked_df, expected_df) -def test_apply_checks_with_dynamic_custom_message(ws, spark): - """Custom message callable with column value should produce dynamic messages.""" +def test_apply_checks_with_dynamic_column_value_message(ws, spark): + """SQL expression referencing {column_value} should produce dynamic messages.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) @@ -62,7 +56,8 @@ def test_apply_checks_with_dynamic_custom_message(ws, spark): criticality="warn", check_func=check_funcs.is_not_null, column="b", - message=_custom_message, + message="concat('Rule ', '{rule_name}', ' ({check_func_name}) failed for value: '," + " coalesce(cast({column_value} as string), 'null'))", ), ] @@ -70,7 +65,7 @@ def test_apply_checks_with_dynamic_custom_message(ws, spark): expected_warnings = [ build_quality_violation( name="b_not_null", - message="Rule 'b_not_null' (is_not_null) failed for value: null", + message="Rule b_not_null (is_not_null) failed for value: null", columns=["b"], function="is_not_null", ) @@ -113,7 +108,7 @@ def test_apply_checks_without_custom_message_unchanged(ws, spark): def test_apply_checks_passing_rows_have_no_custom_message(ws, spark): - """Rows that pass the check should not have a message even with a custom message callable.""" + """Rows that pass the check should not have a message even with a custom message.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, 3, 5]], SCHEMA) @@ -123,7 +118,7 @@ def test_apply_checks_passing_rows_have_no_custom_message(ws, spark): criticality="error", check_func=check_funcs.is_not_null, column="c", - message=_static_custom_message, + message="'Custom error: {rule_name}'", ), ] @@ -144,100 +139,144 @@ def test_for_each_col_rule_with_custom_message(ws, spark): columns=["a", "b"], check_func=check_funcs.is_not_null, criticality="error", - message=_static_custom_message, + message="'Custom error: {rule_name}'", ).get_rules() checked_df = dq_engine.apply_checks(test_df, rules) - rows = checked_df.collect() - assert len(rows) == 1 - errors = rows[0]["_errors"] - assert len(errors) == 2 - messages = sorted([e["message"] for e in errors]) - assert all(msg.startswith("Custom error:") for msg in messages) + expected_errors = [ + build_quality_violation( + name="a_is_null", + message="Custom error: a_is_null", + columns=["a"], + function="is_not_null", + ), + build_quality_violation( + name="b_is_null", + message="Custom error: b_is_null", + columns=["b"], + function="is_not_null", + ), + ] + expected_df = spark.createDataFrame( + [[None, None, 3, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) -def test_apply_checks_with_kwargs_only_custom_message(ws, spark): - """Custom message callable with **kwargs should receive all supported context args.""" +def test_apply_checks_with_column_value_non_null(ws, spark): + """When a check fails with a non-null value, {column_value} should show the actual value.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) - test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) + test_df = spark.createDataFrame([[1, -5, 3]], SCHEMA) rules = [ DQRowRule( - name="b_not_null", - criticality="warn", - check_func=check_funcs.is_not_null, + name="b_positive", + criticality="error", + check_func=check_funcs.is_not_less_than, column="b", - message=_kwargs_only_custom_message, + check_func_kwargs={"limit": 0}, + message="concat('{rule_name}', ': value ', coalesce(cast({column_value} as string), 'null'), ' is not positive')", ), ] checked_df = dq_engine.apply_checks(test_df, rules) - expected_warnings = [ + expected_errors = [ build_quality_violation( - name="b_not_null", - message="kwargs keys: check_func_args,check_func_name,column_value,rule_name", + name="b_positive", + message="b_positive: value -5 is not positive", columns=["b"], - function="is_not_null", + function="is_not_less_than", ) ] expected_df = spark.createDataFrame( - [[1, None, 3, None, expected_warnings]], + [[1, -5, 3, expected_errors, None]], EXPECTED_SCHEMA, ) assert_df_equality(checked_df, expected_df) -def test_apply_checks_with_rule_name_only_custom_message(ws, spark): - """Custom message callable with a single named argument should work.""" +def test_apply_checks_simple_literal_message(ws, spark): + """A plain string literal (no placeholders) should work as a static message.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) - test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) rules = [ DQRowRule( - name="b_not_null", - criticality="warn", + name="c_not_null", + criticality="error", check_func=check_funcs.is_not_null, - column="b", - message=_rule_name_only_custom_message, + column="c", + message="'Column c must not be null'", ), ] checked_df = dq_engine.apply_checks(test_df, rules) - expected_warnings = [ + expected_errors = [ build_quality_violation( - name="b_not_null", - message="Only rule: b_not_null", - columns=["b"], + name="c_not_null", + message="Column c must not be null", + columns=["c"], function="is_not_null", ) ] expected_df = spark.createDataFrame( - [[1, None, 3, None, expected_warnings]], + [[1, 3, None, expected_errors, None]], EXPECTED_SCHEMA, ) assert_df_equality(checked_df, expected_df) -def test_apply_checks_custom_message_check_func_args_include_column(ws, spark): - """check_func_args passed to custom message should include the rule 'column' argument.""" +def test_metadata_static_custom_message(ws, spark): + """Static message defined in YAML-style metadata should appear in the result.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + + checks = [ + { + "name": "c_not_null", + "criticality": "error", + "message": "'Custom error: c_not_null'", + "check": {"function": "is_not_null", "arguments": {"column": "c"}}, + } + ] + + checked_df = dq_engine.apply_checks_by_metadata(test_df, checks) + expected_errors = [ + build_quality_violation( + name="c_not_null", + message="Custom error: c_not_null", + columns=["c"], + function="is_not_null", + ) + ] + expected_df = spark.createDataFrame( + [[1, 3, None, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_metadata_dynamic_column_value_message(ws, spark): + """Dynamic message with {column_value} placeholder from metadata should resolve correctly.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) - rules = [ - DQRowRule( - name="b_not_null", - criticality="warn", - check_func=check_funcs.is_not_null, - column="b", - message=_message_from_column_arg, - ), + checks = [ + { + "name": "b_not_null", + "criticality": "warn", + "message": "concat('Rule ', '{rule_name}', ' ({check_func_name}) failed for value: '," + " coalesce(cast({column_value} as string), 'null'))", + "check": {"function": "is_not_null", "arguments": {"column": "b"}}, + } ] - checked_df = dq_engine.apply_checks(test_df, rules) + checked_df = dq_engine.apply_checks_by_metadata(test_df, checks) expected_warnings = [ build_quality_violation( name="b_not_null", - message="column arg: b", + message="Rule b_not_null (is_not_null) failed for value: null", columns=["b"], function="is_not_null", ) @@ -249,78 +288,68 @@ def test_apply_checks_custom_message_check_func_args_include_column(ws, spark): assert_df_equality(checked_df, expected_df) -def test_apply_checks_custom_message_check_func_args_include_columns(ws, spark): - """check_func_args passed to custom message should include the rule 'columns' argument.""" +def test_metadata_without_message_uses_default(ws, spark): + """Metadata checks without a message field should produce the default message.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) - test_df = spark.createDataFrame([[None, 2, 3]], SCHEMA) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) - rules = [ - DQRowRule( - name="ab_any_null", - criticality="error", - check_func=_any_null_in_columns, - columns=["a", "b"], - message=_message_from_columns_arg, - ), + checks = [ + { + "name": "c_not_null", + "criticality": "error", + "check": {"function": "is_not_null", "arguments": {"column": "c"}}, + } ] - checked_df = dq_engine.apply_checks(test_df, rules) + checked_df = dq_engine.apply_checks_by_metadata(test_df, checks) expected_errors = [ build_quality_violation( - name="ab_any_null", - message="columns arg: a,b", - columns=["a", "b"], - function="_any_null_in_columns", + name="c_not_null", + message="Column 'c' value is null", + columns=["c"], + function="is_not_null", ) ] expected_df = spark.createDataFrame( - [[None, 2, 3, expected_errors, None]], + [[1, 3, None, expected_errors, None]], EXPECTED_SCHEMA, ) assert_df_equality(checked_df, expected_df) -def _custom_message( - rule_name: str, - check_func_name: str, - column_value: Column, -) -> Column: - """Custom message that includes the rule name and check function name.""" - value_str = F.coalesce(column_value.cast("string"), F.lit("null")) - return F.concat(F.lit(f"Rule '{rule_name}' ({check_func_name}) failed for value: "), value_str) - - -def _static_custom_message( - rule_name: str, -) -> Column: - """Custom message that returns a static string.""" - return F.lit(f"Custom error: {rule_name}") - - -def _kwargs_only_custom_message(**kwargs: Any) -> Column: - """Custom message callable that accepts only **kwargs.""" - return F.lit(f"kwargs keys: {','.join(sorted(kwargs.keys()))}") - - -def _rule_name_only_custom_message(rule_name: str) -> Column: - """Custom message callable that accepts a single named argument.""" - return F.lit(f"Only rule: {rule_name}") - - -def _message_from_column_arg(check_func_args: dict[str, Any]) -> Column: - """Custom message that reads 'column' from check_func_args.""" - return F.lit(f"column arg: {check_func_args.get('column')}") - - -def _message_from_columns_arg(check_func_args: dict[str, Any]) -> Column: - """Custom message that reads 'columns' from check_func_args.""" - columns = check_func_args.get("columns", []) - return F.lit(f"columns arg: {','.join(columns)}") +def test_metadata_for_each_column_with_custom_message(ws, spark): + """for_each_column in metadata with message should propagate to all generated rules.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[None, None, 3]], SCHEMA) + checks = [ + { + "criticality": "error", + "message": "'Custom error: {rule_name}'", + "check": { + "function": "is_not_null", + "for_each_column": ["a", "b"], + }, + } + ] -def _any_null_in_columns(columns: list[str]) -> Column: - """Fails if any provided column value is null.""" - condition = F.col(columns[0]).isNull() - for col_name in columns[1:]: - condition = condition | F.col(col_name).isNull() - return check_funcs.make_condition(condition, "one or more columns are null", "any_null_in_columns") + checked_df = dq_engine.apply_checks_by_metadata(test_df, checks) + expected_errors = [ + build_quality_violation( + name="a_is_null", + message="Custom error: a_is_null", + columns=["a"], + function="is_not_null", + ), + build_quality_violation( + name="b_is_null", + message="Custom error: b_is_null", + columns=["b"], + function="is_not_null", + ), + ] + expected_df = spark.createDataFrame( + [[None, None, 3, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index 30c305552..92a852c94 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -1,27 +1,17 @@ -"""Tests for custom message callable on DQRule.""" - -import inspect - -import pyspark.sql.functions as F -from pyspark.sql import Column +"""Tests for custom message SQL expressions on DQRule.""" from databricks.labs.dqx.check_funcs import is_not_null, is_not_null_and_not_empty from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule -def _sample_message_callable(rule_name: str) -> Column: - """Custom message function that returns a string-valued column.""" - return F.lit(f"Custom: {rule_name}") - - -def test_dq_row_rule_accepts_message_callable(): - """DQRowRule with message parameter should be creatable.""" +def test_dq_row_rule_accepts_message_string(): + """DQRowRule with message parameter should accept a SQL expression string.""" rule = DQRowRule( check_func=is_not_null, column="id", - message=_sample_message_callable, + message="concat('Failed: ', '{rule_name}')", ) - assert rule.message is _sample_message_callable + assert rule.message == "concat('Failed: ', '{rule_name}')" def test_dq_row_rule_message_defaults_to_none(): @@ -30,7 +20,7 @@ def test_dq_row_rule_message_defaults_to_none(): assert rule.message is None -def test_dq_dataset_rule_accepts_message_callable(): +def test_dq_dataset_rule_message_defaults_to_none(): """DQDatasetRule without message should default to None.""" rule = DQRowRule(check_func=is_not_null_and_not_empty, column="id") assert rule.message is None @@ -38,15 +28,16 @@ def test_dq_dataset_rule_accepts_message_callable(): def test_dq_for_each_col_rule_propagates_message(): """DQForEachColRule should propagate message to generated DQRowRules.""" + msg = "concat('Check failed for ', '{rule_name}')" for_each_rule = DQForEachColRule( columns=["a", "b"], check_func=is_not_null, - message=_sample_message_callable, + message=msg, ) rules = for_each_rule.get_rules() assert len(rules) == 2 for rule in rules: - assert rule.message is _sample_message_callable + assert rule.message == msg def test_dq_for_each_col_rule_message_defaults_to_none(): @@ -60,20 +51,25 @@ def test_dq_for_each_col_rule_message_defaults_to_none(): assert rule.message is None -def test_message_callable_receives_correct_args(): - """The message callable signature should accept rule_name, check_func_name, check_func_args, column_value.""" - sig = inspect.signature(_sample_message_callable) - params = list(sig.parameters.keys()) - assert params == ["rule_name"] +def test_dq_rule_to_dict_includes_message(): + """to_dict() should include the message string since it is serializable.""" + msg = "'Custom error for ' || '{rule_name}'" + rule = DQRowRule( + check_func=is_not_null, + column="id", + name="id_not_null", + message=msg, + ) + rule_dict = rule.to_dict() + assert rule_dict["message"] == msg -def test_dq_rule_to_dict_does_not_include_message(): - """to_dict() should not include the message callable since it is not serializable.""" +def test_dq_rule_to_dict_omits_message_when_none(): + """to_dict() should not include message key when it is None.""" rule = DQRowRule( check_func=is_not_null, column="id", name="id_not_null", - message=_sample_message_callable, ) rule_dict = rule.to_dict() assert "message" not in rule_dict From 7ed09c7bcd7975e0f2d8b0c5d76aab28bc4b0b8e Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 3 May 2026 13:44:45 -0400 Subject: [PATCH 07/21] Rename DQRule.message to message_expr; accept str | Column; drop placeholders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The custom-message field is renamed to message_expr and now accepts either a Spark SQL expression string or a Spark Column object. The expression is evaluated as-is — DQX no longer substitutes {rule_name}, {check_func_name}, or {column_value} placeholders. Callers reference columns and identifiers directly inside the expression. - DQRule / DQForEachColRule: rename message -> message_expr (str | Column | None) - to_dict serialises message_expr only when it is a string (Column is in-process) - DQRuleManager._build_message_col simplified: F.expr(str) or use Column directly - checks_serializer: read message_expr key from metadata, pass to constructors - Unit tests rewritten for the new attribute name and Column variant - Integration tests rewritten without placeholders; the existing test_apply_checks_without_custom_message_unchanged is dropped (the Engine's no-message default path is already covered by surrounding tests) - Docs (quality_checks.mdx) rewritten to drop the placeholder section, add a Column-based Python example, and annotate each example with the rendered message string Co-authored-by: Isaac --- docs/dqx/docs/reference/quality_checks.mdx | 44 +++++-- src/databricks/labs/dqx/checks_serializer.py | 8 +- src/databricks/labs/dqx/manager.py | 30 ++--- src/databricks/labs/dqx/rule.py | 24 ++-- tests/integration/test_custom_messages.py | 125 +++++++++++-------- tests/unit/test_custom_messages.py | 80 ++++++++---- 6 files changed, 184 insertions(+), 127 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 690072f6d..d6f94c6c3 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -4082,35 +4082,37 @@ When using dataset-level checks, the top-level `filter` condition is pushed down ## Customizing check messages -Users can provide a Spark SQL expression to customize messages created by DQX failures. Set the `message` attribute of any `DQRule` to a SQL expression that returns a string value. This overrides the check function's default message. +Users can override the default failure message of any `DQRule` by setting its `message_expr` attribute. The value is evaluated as-is — DQX does not substitute placeholders, so any column references, casts, or rule-identifying literals must be supplied directly by the caller. -Custom SQL expressions support the following placeholders, which are substituted before evaluation: -- `{rule_name}` is replaced with the name of the DQX rule -- `{check_func_name}` is replaced with the name of the DQX check function -- `{column_value}` is replaced with the observed column value (cast to string). This can only be used for single-column values. +`message_expr` accepts either: +- a Spark SQL expression string (evaluated via `F.expr(...)`), or +- a Spark `Column` object (used directly). + +In the YAML/JSON metadata form, only the string variant is supported, because Spark `Column` objects have no canonical serialised representation. -If your message references `{column_value}`, wrap it with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. +When your expression references a column, wrap it with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. ```python + import pyspark.sql.functions as F from databricks.labs.dqx import check_funcs from databricks.labs.dqx.rule import DQRowRule - # simple static message + # static message (string expression): "Email must not be null" checks = [ DQRowRule( name="email_not_null", criticality="error", check_func=check_funcs.is_not_null, column="email", - message="'Email must not be null'", + message_expr="'Email must not be null'", ) ] - # dynamic message with column value + # dynamic message with column value (string expression): "age_positive: value is not valid" checks = [ DQRowRule( name="age_positive", @@ -4118,24 +4120,42 @@ If your message references `{column_value}`, wrap it with `coalesce` to avoid nu check_func=check_funcs.is_greater_than, column="age", check_func_kwargs={"limit": 0}, - message="concat('{rule_name}', ': value ', coalesce(cast({column_value} as string), 'null'), ' is not valid')", + message_expr="concat('age_positive: value ', coalesce(cast(age as string), 'null'), ' is not valid')", + ) + ] + + # dynamic message with column value (Column expression): "age_positive: value is not valid" + checks = [ + DQRowRule( + name="age_positive", + criticality="error", + check_func=check_funcs.is_greater_than, + column="age", + check_func_kwargs={"limit": 0}, + message_expr=F.concat( + F.lit("age_positive: value "), + F.coalesce(F.col("age").cast("string"), F.lit("null")), + F.lit(" is not valid"), + ), ) ] ``` ```yaml + # static message: "Email must not be null" - name: email_not_null criticality: error - message: "'Email must not be null'" + message_expr: "'Email must not be null'" check: function: is_not_null arguments: column: email + # dynamic message with column value: "age_positive: value is not valid" - name: age_positive criticality: error - message: "concat('{rule_name}', ': value ', coalesce(cast({column_value} as string), 'null'), ' is not valid')" + message_expr: "concat('age_positive: value ', coalesce(cast(age as string), 'null'), ' is not valid')" check: function: is_greater_than arguments: diff --git a/src/databricks/labs/dqx/checks_serializer.py b/src/databricks/labs/dqx/checks_serializer.py index a3536fe38..df489dd26 100644 --- a/src/databricks/labs/dqx/checks_serializer.py +++ b/src/databricks/labs/dqx/checks_serializer.py @@ -267,7 +267,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: criticality = check_def.get("criticality", "error") filter_str = check_def.get("filter") user_metadata = check_def.get("user_metadata") - message = check_def.get("message") + message_expr = check_def.get("message_expr") # Exclude `column` and `columns` from check_func_kwargs # as these are always included in the check function call @@ -283,7 +283,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: filter=filter_str, check_func_kwargs=check_func_kwargs, user_metadata=user_metadata, - message=message, + message_expr=message_expr, ).get_rules() else: rule_type = CHECK_FUNC_REGISTRY.get(func_name) @@ -298,7 +298,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: criticality=criticality, filter=filter_str, user_metadata=user_metadata, - message=message, + message_expr=message_expr, ) ) else: # default to row-level rule @@ -312,7 +312,7 @@ def deserialize(self, checks: list[dict]) -> list[DQRule]: criticality=criticality, filter=filter_str, user_metadata=user_metadata, - message=message, + message_expr=message_expr, ) ) diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index 2eb2b788f..0b4a0d4ec 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -171,34 +171,24 @@ def _build_result_struct(self, condition: Column, skipped: bool = False) -> Colu def _build_message_col(self, condition: Column) -> Column: """ - Builds the message column, using the default message or a custom Spark SQL expression if - specified in the rule definition. The SQL expression supports placeholders that are - substituted before evaluation: - - * '{rule_name}' is replaced with the quoted name of the DQX rule - * '{check_func_name}' is replaced with the quoted name of the DQX check function - * '{column_value}' is replaced with a SQL reference to the column value (cast to string) + Builds the message column, using the default message or the user-supplied + ``message_expr`` from the rule definition. The expression is evaluated as-is — DQX + does not substitute placeholders. Accepts either a Spark SQL expression string or a + Spark Column. Args: condition: Default DQX condition message returned by evaluating the DQX check function Returns: - The custom DQX condition message if specified in the rule definition, otherwise the default DQX condition - message + The custom DQX condition message if ``message_expr`` is set on the rule, otherwise the + default DQX condition message. """ - if self.check.message is None: + if self.check.message_expr is None: return condition - column_value_sql = f"CAST(`{self.check.column}` AS STRING)" if self.check.column else "NULL" - escaped_check_name = self.check.name.replace("'", "\\'") - escaped_check_function = self.check.check_func.__name__.replace("'", "\\'") - - message_expr = self.check.message - message_expr = message_expr.replace("{rule_name}", escaped_check_name) - message_expr = message_expr.replace("{check_func_name}", escaped_check_function) - message_expr = message_expr.replace("{column_value}", column_value_sql) - - custom_message = F.expr(message_expr) + custom_message = ( + self.check.message_expr if isinstance(self.check.message_expr, Column) else F.expr(self.check.message_expr) + ) return F.when(condition.isNotNull(), custom_message).otherwise(F.lit(None).cast("string")) def _get_invalid_cols_message(self) -> str: diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 37eecc8fe..83c8e349c 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -165,11 +165,11 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): * *check_func_args* (optional) - Positional arguments for the check function (excluding *column*). * *check_func_kwargs* (optional) - Keyword arguments for the check function (excluding *column*). * *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. - * *message* (optional) - User-defined Spark SQL expression used as the check failure message. Supports the - following placeholders which are substituted when checks are evaluated: - - `{rule_name}` is replaced with the name of the DQX rule - - `{check_func_name}` is replaced with the name of the DQX check function - - `{column_value}` is replaced with a SQL expression referencing the column value (cast to string) + * *message_expr* (optional) - User-defined expression used as the check failure message. Accepts either a + Spark SQL expression string or a Spark *Column*. The expression is evaluated as-is — DQX does not + substitute placeholders, so any column references, casts, or rule-identifying literals must be + supplied directly by the caller (e.g., ``F.concat(F.lit('age_positive: value '), F.col('age').cast('string'))`` + or ``"concat('age_positive: value ', cast(age as string))"``). """ check_func: Callable @@ -181,7 +181,7 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): check_func_args: list[Any] = field(default_factory=list) check_func_kwargs: dict[str, Any] = field(default_factory=dict) user_metadata: dict[str, str] | None = None - message: str | None = None + message_expr: str | Column | None = None def __post_init__(self): self._validate_rule_type(self.check_func) @@ -265,8 +265,10 @@ def to_dict(self) -> dict: if self.user_metadata: metadata["user_metadata"] = self.user_metadata - if self.message: - metadata["message"] = self.message + # Only string expressions can be round-tripped through metadata; Column objects are + # in-process Spark expressions with no canonical YAML/JSON representation. + if isinstance(self.message_expr, str) and self.message_expr: + metadata["message_expr"] = self.message_expr return metadata def _initialize_column_if_missing(self): @@ -436,7 +438,7 @@ class DQForEachColRule(DQRuleTypeMixin): check_func_args: list[Any] = field(default_factory=list) check_func_kwargs: dict[str, Any] = field(default_factory=dict) user_metadata: dict[str, str] | None = None - message: str | None = None + message_expr: str | Column | None = None def get_rules(self) -> list[DQRule]: """Build a list of rules for a set of columns. @@ -462,7 +464,7 @@ def get_rules(self) -> list[DQRule]: criticality=self.criticality, filter=self.filter, user_metadata=self.user_metadata, - message=self.message, + message_expr=self.message_expr, ) ) else: # default to row-level rule @@ -477,7 +479,7 @@ def get_rules(self) -> list[DQRule]: criticality=self.criticality, filter=self.filter, user_metadata=self.user_metadata, - message=self.message, + message_expr=self.message_expr, ) ) return rules diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index 7384f092a..587e75f7b 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -1,4 +1,6 @@ -"""Integration tests for custom message SQL expressions on DQRule.""" +"""Integration tests for custom message expressions on DQRule.""" + +import pyspark.sql.functions as F from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule @@ -25,7 +27,7 @@ def test_apply_checks_with_static_custom_message(ws, spark): criticality="error", check_func=check_funcs.is_not_null, column="c", - message="'Custom error: {rule_name}'", + message_expr="'Custom error: c_not_null'", ), ] @@ -46,7 +48,7 @@ def test_apply_checks_with_static_custom_message(ws, spark): def test_apply_checks_with_dynamic_column_value_message(ws, spark): - """SQL expression referencing {column_value} should produce dynamic messages.""" + """SQL expression referencing the actual column should produce dynamic messages.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) @@ -56,8 +58,10 @@ def test_apply_checks_with_dynamic_column_value_message(ws, spark): criticality="warn", check_func=check_funcs.is_not_null, column="b", - message="concat('Rule ', '{rule_name}', ' ({check_func_name}) failed for value: '," - " coalesce(cast({column_value} as string), 'null'))", + message_expr=( + "concat('Rule b_not_null (is_not_null) failed for value: '," + " coalesce(cast(b as string), 'null'))" + ), ), ] @@ -77,36 +81,6 @@ def test_apply_checks_with_dynamic_column_value_message(ws, spark): assert_df_equality(checked_df, expected_df) -def test_apply_checks_without_custom_message_unchanged(ws, spark): - """Rules without custom message should produce the default message as before.""" - dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) - test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) - - rules = [ - DQRowRule( - name="c_not_null", - criticality="error", - check_func=check_funcs.is_not_null, - column="c", - ), - ] - - checked_df = dq_engine.apply_checks(test_df, rules) - expected_errors = [ - build_quality_violation( - name="c_not_null", - message="Column 'c' value is null", - columns=["c"], - function="is_not_null", - ) - ] - expected_df = spark.createDataFrame( - [[1, 3, None, expected_errors, None]], - EXPECTED_SCHEMA, - ) - assert_df_equality(checked_df, expected_df) - - def test_apply_checks_passing_rows_have_no_custom_message(ws, spark): """Rows that pass the check should not have a message even with a custom message.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) @@ -118,7 +92,7 @@ def test_apply_checks_passing_rows_have_no_custom_message(ws, spark): criticality="error", check_func=check_funcs.is_not_null, column="c", - message="'Custom error: {rule_name}'", + message_expr="'Custom error: c_not_null'", ), ] @@ -131,7 +105,12 @@ def test_apply_checks_passing_rows_have_no_custom_message(ws, spark): def test_for_each_col_rule_with_custom_message(ws, spark): - """DQForEachColRule with message should propagate to generated rules and appear in results.""" + """DQForEachColRule with message_expr should propagate to generated rules. + + The same expression is used for every generated rule. To produce a per-column + message, reference each column inline (e.g. concat with a per-column literal) or + construct rules individually. + """ dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[None, None, 3]], SCHEMA) @@ -139,20 +118,20 @@ def test_for_each_col_rule_with_custom_message(ws, spark): columns=["a", "b"], check_func=check_funcs.is_not_null, criticality="error", - message="'Custom error: {rule_name}'", + message_expr="'Custom error: column missing'", ).get_rules() checked_df = dq_engine.apply_checks(test_df, rules) expected_errors = [ build_quality_violation( name="a_is_null", - message="Custom error: a_is_null", + message="Custom error: column missing", columns=["a"], function="is_not_null", ), build_quality_violation( name="b_is_null", - message="Custom error: b_is_null", + message="Custom error: column missing", columns=["b"], function="is_not_null", ), @@ -165,7 +144,7 @@ def test_for_each_col_rule_with_custom_message(ws, spark): def test_apply_checks_with_column_value_non_null(ws, spark): - """When a check fails with a non-null value, {column_value} should show the actual value.""" + """When a check fails with a non-null value, the message can include that value.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, -5, 3]], SCHEMA) @@ -176,7 +155,10 @@ def test_apply_checks_with_column_value_non_null(ws, spark): check_func=check_funcs.is_not_less_than, column="b", check_func_kwargs={"limit": 0}, - message="concat('{rule_name}', ': value ', coalesce(cast({column_value} as string), 'null'), ' is not positive')", + message_expr=( + "concat('b_positive: value ', coalesce(cast(b as string), 'null')," + " ' is not positive')" + ), ), ] @@ -197,7 +179,7 @@ def test_apply_checks_with_column_value_non_null(ws, spark): def test_apply_checks_simple_literal_message(ws, spark): - """A plain string literal (no placeholders) should work as a static message.""" + """A plain string literal (no expression) should work as a static message.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) @@ -207,7 +189,7 @@ def test_apply_checks_simple_literal_message(ws, spark): criticality="error", check_func=check_funcs.is_not_null, column="c", - message="'Column c must not be null'", + message_expr="'Column c must not be null'", ), ] @@ -227,6 +209,41 @@ def test_apply_checks_simple_literal_message(ws, spark): assert_df_equality(checked_df, expected_df) +def test_apply_checks_with_column_message_expr(ws, spark): + """A Spark Column passed as message_expr should be used directly without conversion.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + message_expr=F.concat( + F.lit("Custom error: c_not_null (value="), + F.coalesce(F.col("c").cast("string"), F.lit("null")), + F.lit(")"), + ), + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_errors = [ + build_quality_violation( + name="c_not_null", + message="Custom error: c_not_null (value=null)", + columns=["c"], + function="is_not_null", + ) + ] + expected_df = spark.createDataFrame( + [[1, 3, None, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + def test_metadata_static_custom_message(ws, spark): """Static message defined in YAML-style metadata should appear in the result.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) @@ -236,7 +253,7 @@ def test_metadata_static_custom_message(ws, spark): { "name": "c_not_null", "criticality": "error", - "message": "'Custom error: c_not_null'", + "message_expr": "'Custom error: c_not_null'", "check": {"function": "is_not_null", "arguments": {"column": "c"}}, } ] @@ -258,7 +275,7 @@ def test_metadata_static_custom_message(ws, spark): def test_metadata_dynamic_column_value_message(ws, spark): - """Dynamic message with {column_value} placeholder from metadata should resolve correctly.""" + """Dynamic message expression referencing the column from metadata should resolve correctly.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, None, 3]], SCHEMA) @@ -266,8 +283,10 @@ def test_metadata_dynamic_column_value_message(ws, spark): { "name": "b_not_null", "criticality": "warn", - "message": "concat('Rule ', '{rule_name}', ' ({check_func_name}) failed for value: '," - " coalesce(cast({column_value} as string), 'null'))", + "message_expr": ( + "concat('Rule b_not_null (is_not_null) failed for value: '," + " coalesce(cast(b as string), 'null'))" + ), "check": {"function": "is_not_null", "arguments": {"column": "b"}}, } ] @@ -289,7 +308,7 @@ def test_metadata_dynamic_column_value_message(ws, spark): def test_metadata_without_message_uses_default(ws, spark): - """Metadata checks without a message field should produce the default message.""" + """Metadata checks without a message_expr field should produce the default message.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) @@ -318,14 +337,14 @@ def test_metadata_without_message_uses_default(ws, spark): def test_metadata_for_each_column_with_custom_message(ws, spark): - """for_each_column in metadata with message should propagate to all generated rules.""" + """for_each_column in metadata with message_expr should propagate to all generated rules.""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) test_df = spark.createDataFrame([[None, None, 3]], SCHEMA) checks = [ { "criticality": "error", - "message": "'Custom error: {rule_name}'", + "message_expr": "'Custom error: column missing'", "check": { "function": "is_not_null", "for_each_column": ["a", "b"], @@ -337,13 +356,13 @@ def test_metadata_for_each_column_with_custom_message(ws, spark): expected_errors = [ build_quality_violation( name="a_is_null", - message="Custom error: a_is_null", + message="Custom error: column missing", columns=["a"], function="is_not_null", ), build_quality_violation( name="b_is_null", - message="Custom error: b_is_null", + message="Custom error: column missing", columns=["b"], function="is_not_null", ), diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index 92a852c94..d34f89d5b 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -1,75 +1,101 @@ -"""Tests for custom message SQL expressions on DQRule.""" +"""Tests for custom message expressions on DQRule.""" + +import pyspark.sql.functions as F +from pyspark.sql import Column from databricks.labs.dqx.check_funcs import is_not_null, is_not_null_and_not_empty from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule -def test_dq_row_rule_accepts_message_string(): - """DQRowRule with message parameter should accept a SQL expression string.""" +def test_dq_row_rule_accepts_message_expr_string(): + """DQRowRule with message_expr should accept a SQL expression string.""" + rule = DQRowRule( + check_func=is_not_null, + column="id", + message_expr="concat('Failed: ', 'id_not_null')", + ) + assert rule.message_expr == "concat('Failed: ', 'id_not_null')" + + +def test_dq_row_rule_accepts_message_expr_column(): + """DQRowRule with message_expr should accept a Spark Column object.""" + column_expr = F.concat(F.lit("Failed: "), F.lit("id_not_null")) rule = DQRowRule( check_func=is_not_null, column="id", - message="concat('Failed: ', '{rule_name}')", + message_expr=column_expr, ) - assert rule.message == "concat('Failed: ', '{rule_name}')" + assert isinstance(rule.message_expr, Column) -def test_dq_row_rule_message_defaults_to_none(): - """DQRowRule without message should default to None.""" +def test_dq_row_rule_message_expr_defaults_to_none(): + """DQRowRule without message_expr should default to None.""" rule = DQRowRule(check_func=is_not_null, column="id") - assert rule.message is None + assert rule.message_expr is None -def test_dq_dataset_rule_message_defaults_to_none(): - """DQDatasetRule without message should default to None.""" +def test_dq_dataset_rule_message_expr_defaults_to_none(): + """DQDatasetRule without message_expr should default to None.""" rule = DQRowRule(check_func=is_not_null_and_not_empty, column="id") - assert rule.message is None + assert rule.message_expr is None -def test_dq_for_each_col_rule_propagates_message(): - """DQForEachColRule should propagate message to generated DQRowRules.""" - msg = "concat('Check failed for ', '{rule_name}')" +def test_dq_for_each_col_rule_propagates_message_expr(): + """DQForEachColRule should propagate message_expr to generated DQRowRules.""" + msg = "concat('Check failed for ', 'rule_name')" for_each_rule = DQForEachColRule( columns=["a", "b"], check_func=is_not_null, - message=msg, + message_expr=msg, ) rules = for_each_rule.get_rules() assert len(rules) == 2 for rule in rules: - assert rule.message == msg + assert rule.message_expr == msg -def test_dq_for_each_col_rule_message_defaults_to_none(): - """DQForEachColRule without message should produce rules with message=None.""" +def test_dq_for_each_col_rule_message_expr_defaults_to_none(): + """DQForEachColRule without message_expr should produce rules with message_expr=None.""" for_each_rule = DQForEachColRule( columns=["a", "b"], check_func=is_not_null, ) rules = for_each_rule.get_rules() for rule in rules: - assert rule.message is None + assert rule.message_expr is None + + +def test_dq_rule_to_dict_includes_message_expr_when_string(): + """to_dict() should include the message_expr key when set to a SQL string.""" + msg = "'Custom error for id_not_null'" + rule = DQRowRule( + check_func=is_not_null, + column="id", + name="id_not_null", + message_expr=msg, + ) + rule_dict = rule.to_dict() + assert rule_dict["message_expr"] == msg -def test_dq_rule_to_dict_includes_message(): - """to_dict() should include the message string since it is serializable.""" - msg = "'Custom error for ' || '{rule_name}'" +def test_dq_rule_to_dict_omits_message_expr_when_none(): + """to_dict() should not include message_expr key when it is None.""" rule = DQRowRule( check_func=is_not_null, column="id", name="id_not_null", - message=msg, ) rule_dict = rule.to_dict() - assert rule_dict["message"] == msg + assert "message_expr" not in rule_dict -def test_dq_rule_to_dict_omits_message_when_none(): - """to_dict() should not include message key when it is None.""" +def test_dq_rule_to_dict_omits_message_expr_when_column(): + """Column message_expr is in-process only and is excluded from serialised metadata.""" rule = DQRowRule( check_func=is_not_null, column="id", name="id_not_null", + message_expr=F.lit("Custom error"), ) rule_dict = rule.to_dict() - assert "message" not in rule_dict + assert "message_expr" not in rule_dict From 52cd2df8ef2afcb445654cd3e160c9a0bf429cb1 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 3 May 2026 13:57:22 -0400 Subject: [PATCH 08/21] Update docs --- docs/dqx/docs/reference/quality_checks.mdx | 28 +++++++++------------- src/databricks/labs/dqx/rule.py | 10 ++++---- 2 files changed, 16 insertions(+), 22 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index d6f94c6c3..a23432dd9 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -4082,13 +4082,7 @@ When using dataset-level checks, the top-level `filter` condition is pushed down ## Customizing check messages -Users can override the default failure message of any `DQRule` by setting its `message_expr` attribute. The value is evaluated as-is — DQX does not substitute placeholders, so any column references, casts, or rule-identifying literals must be supplied directly by the caller. - -`message_expr` accepts either: -- a Spark SQL expression string (evaluated via `F.expr(...)`), or -- a Spark `Column` object (used directly). - -In the YAML/JSON metadata form, only the string variant is supported, because Spark `Column` objects have no canonical serialised representation. +Users can override the default failure message of any `DQRule` by specifying a custom message expression. Set `message_expr` to either a Spark SQL expression string or a Spark `Column` expression that returns a string-valued message when a check fails. When your expression references a column, wrap it with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. @@ -4101,7 +4095,7 @@ When your expression references a column, wrap it with `coalesce` to avoid null from databricks.labs.dqx import check_funcs from databricks.labs.dqx.rule import DQRowRule - # static message (string expression): "Email must not be null" + # static message: "Email must not be null" checks = [ DQRowRule( name="email_not_null", @@ -4112,28 +4106,28 @@ When your expression references a column, wrap it with `coalesce` to avoid null ) ] - # dynamic message with column value (string expression): "age_positive: value is not valid" + # dynamic message using a SQL expression string: "age_positive: age is not valid" checks = [ DQRowRule( name="age_positive", criticality="error", - check_func=check_funcs.is_greater_than, + check_func=check_funcs.is_not_less_than, column="age", check_func_kwargs={"limit": 0}, - message_expr="concat('age_positive: value ', coalesce(cast(age as string), 'null'), ' is not valid')", + message_expr="concat('age_positive: age ', coalesce(cast(age as string), 'null'), ' is not valid')", ) ] - # dynamic message with column value (Column expression): "age_positive: value is not valid" + # dynamic message using a Spark Column expression: "age_positive: age is not valid" checks = [ DQRowRule( name="age_positive", criticality="error", - check_func=check_funcs.is_greater_than, + check_func=check_funcs.is_not_less_than, column="age", check_func_kwargs={"limit": 0}, message_expr=F.concat( - F.lit("age_positive: value "), + F.lit("age_positive: age "), F.coalesce(F.col("age").cast("string"), F.lit("null")), F.lit(" is not valid"), ), @@ -4152,12 +4146,12 @@ When your expression references a column, wrap it with `coalesce` to avoid null arguments: column: email - # dynamic message with column value: "age_positive: value is not valid" + # dynamic message using a SQL expression string: "age_positive: age is not valid" - name: age_positive criticality: error - message_expr: "concat('age_positive: value ', coalesce(cast(age as string), 'null'), ' is not valid')" + message_expr: "concat('age_positive: age ', coalesce(cast(age as string), 'null'), ' is not valid')" check: - function: is_greater_than + function: is_not_less_than arguments: column: age limit: 0 diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 83c8e349c..4a0b10fdc 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -165,11 +165,11 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): * *check_func_args* (optional) - Positional arguments for the check function (excluding *column*). * *check_func_kwargs* (optional) - Keyword arguments for the check function (excluding *column*). * *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check. - * *message_expr* (optional) - User-defined expression used as the check failure message. Accepts either a - Spark SQL expression string or a Spark *Column*. The expression is evaluated as-is — DQX does not - substitute placeholders, so any column references, casts, or rule-identifying literals must be - supplied directly by the caller (e.g., ``F.concat(F.lit('age_positive: value '), F.col('age').cast('string'))`` - or ``"concat('age_positive: value ', cast(age as string))"``). + * *message_expr* (optional) - User-defined expression used as the check failure message. Accepts either + a Spark SQL expression string or a Spark *Column* expression. The expression is evaluated as-is. + Any column references, casts, or rule-identifying literals must be supplied directly by the caller + (e.g., ``F.concat(F.lit('age_positive: value '), F.col('age').cast('string'))`` or + ``"concat('age_positive: value ', cast(age as string))"``). """ check_func: Callable From 49c9a59cb2b95cd397735578a3573f8369e9f22f Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 3 Jun 2026 18:53:28 -0400 Subject: [PATCH 09/21] Format tests --- tests/integration/test_custom_messages.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index 587e75f7b..176b674ad 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -59,8 +59,7 @@ def test_apply_checks_with_dynamic_column_value_message(ws, spark): check_func=check_funcs.is_not_null, column="b", message_expr=( - "concat('Rule b_not_null (is_not_null) failed for value: '," - " coalesce(cast(b as string), 'null'))" + "concat('Rule b_not_null (is_not_null) failed for value: '," " coalesce(cast(b as string), 'null'))" ), ), ] @@ -155,10 +154,7 @@ def test_apply_checks_with_column_value_non_null(ws, spark): check_func=check_funcs.is_not_less_than, column="b", check_func_kwargs={"limit": 0}, - message_expr=( - "concat('b_positive: value ', coalesce(cast(b as string), 'null')," - " ' is not positive')" - ), + message_expr=("concat('b_positive: value ', coalesce(cast(b as string), 'null')," " ' is not positive')"), ), ] @@ -284,8 +280,7 @@ def test_metadata_dynamic_column_value_message(ws, spark): "name": "b_not_null", "criticality": "warn", "message_expr": ( - "concat('Rule b_not_null (is_not_null) failed for value: '," - " coalesce(cast(b as string), 'null'))" + "concat('Rule b_not_null (is_not_null) failed for value: '," " coalesce(cast(b as string), 'null'))" ), "check": {"function": "is_not_null", "arguments": {"column": "b"}}, } From 237a572eeb76a3bc02faf5ada1c72e579945ec48 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 4 Jun 2026 00:05:13 -0400 Subject: [PATCH 10/21] Fix unit tests --- tests/unit/test_custom_messages.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index d34f89d5b..26bffc2cb 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -25,7 +25,7 @@ def test_dq_row_rule_accepts_message_expr_column(): column="id", message_expr=column_expr, ) - assert isinstance(rule.message_expr, Column) + assert rule.message_expr is not None def test_dq_row_rule_message_expr_defaults_to_none(): From 7d2ec21d54ea7bb9e5cdd67ca3606aac2a4281cd Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 4 Jun 2026 17:21:01 -0400 Subject: [PATCH 11/21] Format --- tests/unit/test_custom_messages.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index 26bffc2cb..9b1cee404 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -1,7 +1,6 @@ """Tests for custom message expressions on DQRule.""" import pyspark.sql.functions as F -from pyspark.sql import Column from databricks.labs.dqx.check_funcs import is_not_null, is_not_null_and_not_empty from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule From 71727650f97edae15f00e4ad8894e77b9717a650 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 4 Jun 2026 19:52:59 -0400 Subject: [PATCH 12/21] Add validation safety --- src/databricks/labs/dqx/manager.py | 20 +++++++++++++++----- src/databricks/labs/dqx/rule.py | 28 +++++++++++++++++++++++++--- 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index 0b4a0d4ec..59dd82f23 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -8,12 +8,13 @@ from pyspark.sql import DataFrame, Column, SparkSession from databricks.labs.dqx import check_funcs +from databricks.labs.dqx.errors import UnsafeSqlQueryError from databricks.labs.dqx.executor import DQCheckResult, DQRuleExecutorFactory from databricks.labs.dqx.rule import ( DQRule, ) from databricks.labs.dqx.schema.dq_result_schema import dq_result_item_schema -from databricks.labs.dqx.utils import get_column_name_or_alias +from databricks.labs.dqx.utils import get_column_name_or_alias, is_sql_query_safe logger = logging.getLogger(__name__) @@ -183,13 +184,22 @@ def _build_message_col(self, condition: Column) -> Column: The custom DQX condition message if ``message_expr`` is set on the rule, otherwise the default DQX condition message. """ + _max_message_length = 500 if self.check.message_expr is None: return condition - custom_message = ( - self.check.message_expr if isinstance(self.check.message_expr, Column) else F.expr(self.check.message_expr) - ) - return F.when(condition.isNotNull(), custom_message).otherwise(F.lit(None).cast("string")) + if isinstance(self.check.message_expr, str): + if not is_sql_query_safe(self.check.message_expr): + raise UnsafeSqlQueryError( + "Provided message expression is not safe for execution. Please ensure it does not contain any unsafe operations." + ) + return F.when( + condition.isNotNull(), F.substr(F.expr(self.check.message_expr)), F.lit(1), F.lit(_max_message_length) + ).otherwise(F.lit(None).cast("string")) + + return F.when( + condition.isNotNull(), F.substr(self.check.message_expr, F.lit(1), F.lit(_max_message_length)) + ).otherwise(F.lit(None).cast("string")) def _get_invalid_cols_message(self) -> str: """ diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 4a0b10fdc..e0783c5f0 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -12,7 +12,7 @@ from pyspark.sql import Column import pyspark.sql.functions as F from databricks.labs.dqx.utils import get_column_name_or_alias, normalize_bound_args -from databricks.labs.dqx.errors import InvalidCheckError +from databricks.labs.dqx.errors import InvalidCheckError, InvalidParameterError logger = logging.getLogger(__name__) @@ -190,6 +190,8 @@ def __post_init__(self): self._validate_attributes() check_condition = self.get_check_condition() self._initialize_name_if_missing(check_condition) + if self.message_expr and isinstance(self.message_expr, str): + self._validate_message_expression(self.message_expr) @abc.abstractmethod def get_check_condition(self) -> Column: @@ -267,8 +269,10 @@ def to_dict(self) -> dict: metadata["user_metadata"] = self.user_metadata # Only string expressions can be round-tripped through metadata; Column objects are # in-process Spark expressions with no canonical YAML/JSON representation. - if isinstance(self.message_expr, str) and self.message_expr: - metadata["message_expr"] = self.message_expr + if self.message_expr: + if isinstance(self.message_expr, str): + metadata["message_expr"] = self.message_expr + logger.warning("Message expressions of type 'Column' cannot be serialized; falling back to default message") return metadata def _initialize_column_if_missing(self): @@ -351,6 +355,24 @@ def _is_optional_argument(self, signature: inspect.Signature, arg_name: str): return None # Argument not present return param.default is not inspect.Parameter.empty + @staticmethod + def _validate_message_expression(message_expr: str) -> None: + """ + Checks that the message expression is a logically valid Spark SQL expression. + + Args: + message_expr: Message expression + + Raises: + InvalidParameterError: If the expression is not a logically valid Spark SQL expression. + """ + try: + F.expr(message_expr) + except Exception as exc: + raise InvalidParameterError( + f"Custom message expression '{message_expr}' is not a valid Spark SQL expression." + ) from exc + @dataclass(frozen=True) class DQRowRule(DQRule): From 2891bf0159188cf70270c35535221dc6b26a55ae Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 4 Jun 2026 19:53:39 -0400 Subject: [PATCH 13/21] Fix expression --- src/databricks/labs/dqx/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index 59dd82f23..c6745fe1a 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -194,7 +194,7 @@ def _build_message_col(self, condition: Column) -> Column: "Provided message expression is not safe for execution. Please ensure it does not contain any unsafe operations." ) return F.when( - condition.isNotNull(), F.substr(F.expr(self.check.message_expr)), F.lit(1), F.lit(_max_message_length) + condition.isNotNull(), F.substr(F.expr(self.check.message_expr), F.lit(1), F.lit(_max_message_length)) ).otherwise(F.lit(None).cast("string")) return F.when( From 80bcd2e67651056326f8443fb8d1183538de391e Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 4 Jun 2026 20:16:59 -0400 Subject: [PATCH 14/21] Update edge case handling and widen tests --- src/databricks/labs/dqx/rule.py | 8 ++--- tests/integration/test_custom_messages.py | 39 ++++++++++++++++++++ tests/unit/test_custom_messages.py | 43 +++++++++++++++++++++++ 3 files changed, 86 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index e0783c5f0..d9352229c 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -190,7 +190,7 @@ def __post_init__(self): self._validate_attributes() check_condition = self.get_check_condition() self._initialize_name_if_missing(check_condition) - if self.message_expr and isinstance(self.message_expr, str): + if isinstance(self.message_expr, str): self._validate_message_expression(self.message_expr) @abc.abstractmethod @@ -269,9 +269,9 @@ def to_dict(self) -> dict: metadata["user_metadata"] = self.user_metadata # Only string expressions can be round-tripped through metadata; Column objects are # in-process Spark expressions with no canonical YAML/JSON representation. - if self.message_expr: - if isinstance(self.message_expr, str): - metadata["message_expr"] = self.message_expr + if isinstance(self.message_expr, str): + metadata["message_expr"] = self.message_expr + elif self.message_expr is not None: logger.warning("Message expressions of type 'Column' cannot be serialized; falling back to default message") return metadata diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index 176b674ad..b7e047ac8 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -1,8 +1,10 @@ """Integration tests for custom message expressions on DQRule.""" import pyspark.sql.functions as F +import pytest from databricks.labs.dqx.engine import DQEngine +from databricks.labs.dqx.errors import UnsafeSqlQueryError from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule from databricks.labs.dqx import check_funcs from tests.integration.conftest import ( @@ -367,3 +369,40 @@ def test_metadata_for_each_column_with_custom_message(ws, spark): EXPECTED_SCHEMA, ) assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_with_unsafe_message_expr_raises(ws, spark): + """A string message_expr containing an unsafe SQL keyword must be rejected when applied.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + message_expr="concat('msg ', (select x from t where y = (drop table t)))", + ), + ] + + with pytest.raises(UnsafeSqlQueryError, match="not safe for execution"): + dq_engine.apply_checks(test_df, rules) + + +def test_apply_checks_by_metadata_with_unsafe_message_expr_raises(ws, spark): + """An unsafe message_expr supplied via metadata must also be rejected when applied.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + + checks = [ + { + "name": "c_not_null", + "criticality": "error", + "check": {"function": "is_not_null", "arguments": {"column": "c"}}, + "message_expr": "delete from audit", + } + ] + + with pytest.raises(UnsafeSqlQueryError, match="not safe for execution"): + dq_engine.apply_checks_by_metadata(test_df, checks) diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index 9b1cee404..0f938f643 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -1,5 +1,7 @@ """Tests for custom message expressions on DQRule.""" +import logging + import pyspark.sql.functions as F from databricks.labs.dqx.check_funcs import is_not_null, is_not_null_and_not_empty @@ -98,3 +100,44 @@ def test_dq_rule_to_dict_omits_message_expr_when_column(): ) rule_dict = rule.to_dict() assert "message_expr" not in rule_dict + + +def test_dq_rule_to_dict_warns_when_message_expr_is_column(caplog): + """Serialising a Column message_expr should warn that it cannot be round-tripped.""" + rule = DQRowRule( + check_func=is_not_null, + column="id", + name="id_not_null", + message_expr=F.lit("Custom error"), + ) + with caplog.at_level(logging.WARNING, logger="databricks.labs.dqx.rule"): + rule.to_dict() + assert "cannot be serialized" in caplog.text + + +def test_dq_rule_to_dict_does_not_warn_when_message_expr_is_string(caplog): + """A string message_expr is serializable, so no serialization warning should be emitted.""" + rule = DQRowRule( + check_func=is_not_null, + column="id", + name="id_not_null", + message_expr="'Custom error for id_not_null'", + ) + with caplog.at_level(logging.WARNING, logger="databricks.labs.dqx.rule"): + rule.to_dict() + assert "cannot be serialized" not in caplog.text + + +def test_dq_for_each_col_rule_to_dict_warns_for_each_column_when_message_expr_is_column(caplog): + """Column message_expr propagated by DQForEachColRule should warn once per generated rule.""" + for_each_rule = DQForEachColRule( + columns=["a", "b"], + check_func=is_not_null, + message_expr=F.lit("Custom error"), + ) + rules = for_each_rule.get_rules() + with caplog.at_level(logging.WARNING, logger="databricks.labs.dqx.rule"): + for rule in rules: + rule_dict = rule.to_dict() + assert "message_expr" not in rule_dict + assert caplog.text.count("cannot be serialized") == len(rules) From ec8dfd68bd430482b9c51e4028552c61cbba0dda Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 5 Jun 2026 09:37:04 -0400 Subject: [PATCH 15/21] Add documentation of edge cases and best practices --- docs/dqx/docs/reference/quality_checks.mdx | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 31dcb2b56..3f54fef58 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -4099,8 +4099,12 @@ When using dataset-level checks, the top-level `filter` condition is pushed down Users can override the default failure message of any `DQRule` by specifying a custom message expression. Set `message_expr` to either a Spark SQL expression string or a Spark `Column` expression that returns a string-valued message when a check fails. - -When your expression references a column, wrap it with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. + +When using custom message expressions: + +* Wrap column references with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. +* Pass literal string values with quotes (e.g. `'Email must not be null'`) to ensure they are not evaluated as SQL expressions. +* Avoid long messages or referencing unbounded string columns. Message text is **truncated at 500 characters**. From ebdf61e11b8f8acdfa7e44133852802daa54bcd1 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 5 Jun 2026 09:51:20 -0400 Subject: [PATCH 16/21] Improve docstring and fix tests --- src/databricks/labs/dqx/rule.py | 3 ++- tests/unit/test_custom_messages.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index d9352229c..68e9ce9f3 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -169,7 +169,8 @@ class DQRule(abc.ABC, DQRuleTypeMixin, SingleColumnMixin, MultipleColumnsMixin): a Spark SQL expression string or a Spark *Column* expression. The expression is evaluated as-is. Any column references, casts, or rule-identifying literals must be supplied directly by the caller (e.g., ``F.concat(F.lit('age_positive: value '), F.col('age').cast('string'))`` or - ``"concat('age_positive: value ', cast(age as string))"``). + ``"concat('age_positive: value ', cast(age as string))"``). The same message is shared across all + rules generated from a ``DQForEachColRule``. """ check_func: Callable diff --git a/tests/unit/test_custom_messages.py b/tests/unit/test_custom_messages.py index 0f938f643..5e25d3e76 100644 --- a/tests/unit/test_custom_messages.py +++ b/tests/unit/test_custom_messages.py @@ -4,8 +4,8 @@ import pyspark.sql.functions as F -from databricks.labs.dqx.check_funcs import is_not_null, is_not_null_and_not_empty -from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule +from databricks.labs.dqx.check_funcs import is_not_null, is_unique +from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule, DQDatasetRule def test_dq_row_rule_accepts_message_expr_string(): @@ -37,7 +37,7 @@ def test_dq_row_rule_message_expr_defaults_to_none(): def test_dq_dataset_rule_message_expr_defaults_to_none(): """DQDatasetRule without message_expr should default to None.""" - rule = DQRowRule(check_func=is_not_null_and_not_empty, column="id") + rule = DQDatasetRule(check_func=is_unique, columns=["id"]) assert rule.message_expr is None From a78bd20a4dbb2a245fc26100778dbba146ec07f3 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 5 Jun 2026 10:00:33 -0400 Subject: [PATCH 17/21] Note validation behavior in documentation --- docs/dqx/docs/reference/quality_checks.mdx | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 3f54fef58..f1c27dc80 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -4102,6 +4102,7 @@ Users can override the default failure message of any `DQRule` by specifying a c When using custom message expressions: +* Messages defined as SQL expressions are only validated when checks are run (e.g. with `apply_checks_by_metadata(...)`). * Wrap column references with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. * Pass literal string values with quotes (e.g. `'Email must not be null'`) to ensure they are not evaluated as SQL expressions. * Avoid long messages or referencing unbounded string columns. Message text is **truncated at 500 characters**. From 7170d09407d3d531c0602df64f27719fceeea46b Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 7 Jun 2026 10:05:35 -0400 Subject: [PATCH 18/21] Remove usage of is_sql_query_safe and add documentation --- docs/dqx/docs/reference/quality_checks.mdx | 3 +- src/databricks/labs/dqx/manager.py | 38 ++++++++++------------ 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index f1c27dc80..f20e60661 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -4097,7 +4097,7 @@ When using dataset-level checks, the top-level `filter` condition is pushed down ## Customizing check messages -Users can override the default failure message of any `DQRule` by specifying a custom message expression. Set `message_expr` to either a Spark SQL expression string or a Spark `Column` expression that returns a string-valued message when a check fails. +Users can override the default failure message of any `DQRule` by specifying a custom message expression. Set `message_expr` to either a Spark SQL expression string or a Spark `Column` expression that returns a string-valued message when data fails a check. When using custom message expressions: @@ -4106,6 +4106,7 @@ When using custom message expressions: * Wrap column references with `coalesce` to avoid null messages. In Spark SQL, `concat(..., null)` returns `null`. * Pass literal string values with quotes (e.g. `'Email must not be null'`) to ensure they are not evaluated as SQL expressions. * Avoid long messages or referencing unbounded string columns. Message text is **truncated at 500 characters**. +* Avoid unquoted SQL keywords that may be executed naively (e.g. `DELETE` or `TRUNCATE`). diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index c6745fe1a..e3a8ed7fc 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -152,7 +152,7 @@ def _build_result_struct(self, condition: Column, skipped: bool = False) -> Colu # or use literal run time if explicitly overridden run_time_expr = F.current_timestamp() if self.run_time_overwrite is None else F.lit(self.run_time_overwrite) - message_col = self._build_message_col(condition) + message_col = self._build_message_col(condition, skipped=skipped) return F.struct( F.lit(self.check.name).alias("name"), @@ -170,36 +170,34 @@ def _build_result_struct(self, condition: Column, skipped: bool = False) -> Colu F.lit(skipped or None).alias("skipped"), ).cast(dq_result_item_schema) - def _build_message_col(self, condition: Column) -> Column: + def _build_message_col(self, condition: Column, skipped: bool = False) -> Column: """ - Builds the message column, using the default message or the user-supplied - ``message_expr`` from the rule definition. The expression is evaluated as-is — DQX - does not substitute placeholders. Accepts either a Spark SQL expression string or a + Builds the message column, using the default message or the user-supplied ``message_expr`` from the + rule definition. The expression is evaluated as-is. Accepts either a Spark SQL expression string or a Spark Column. Args: condition: Default DQX condition message returned by evaluating the DQX check function + skipped: Whether the check was skipped (default False) Returns: - The custom DQX condition message if ``message_expr`` is set on the rule, otherwise the - default DQX condition message. + The custom DQX condition message if ``message_expr`` is set on the rule and the check was not skipped, + otherwise the default DQX condition message. """ - _max_message_length = 500 + if skipped: + return condition + if self.check.message_expr is None: return condition - if isinstance(self.check.message_expr, str): - if not is_sql_query_safe(self.check.message_expr): - raise UnsafeSqlQueryError( - "Provided message expression is not safe for execution. Please ensure it does not contain any unsafe operations." - ) - return F.when( - condition.isNotNull(), F.substr(F.expr(self.check.message_expr), F.lit(1), F.lit(_max_message_length)) - ).otherwise(F.lit(None).cast("string")) - - return F.when( - condition.isNotNull(), F.substr(self.check.message_expr, F.lit(1), F.lit(_max_message_length)) - ).otherwise(F.lit(None).cast("string")) + _max_message_length = 500 + message_expr = F.substr( + F.expr(self.check.message_expr) if isinstance(self.check.message_expr, str) else self.check.message_expr, + F.lit(1), + F.lit(_max_message_length) + ) + + return F.when(condition.isNotNull(), message_expr).otherwise(F.lit(None).cast("string")) def _get_invalid_cols_message(self) -> str: """ From a950d737f70e1207aeced2bebfcd6958b2eb373d Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 7 Jun 2026 11:23:25 -0400 Subject: [PATCH 19/21] Format, update docs, and update tests --- docs/dqx/docs/reference/quality_checks.mdx | 1 + src/databricks/labs/dqx/manager.py | 5 +-- tests/integration/test_custom_messages.py | 52 +++++++++++----------- 3 files changed, 29 insertions(+), 29 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index f20e60661..e101bfdd3 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -4098,6 +4098,7 @@ When using dataset-level checks, the top-level `filter` condition is pushed down ## Customizing check messages Users can override the default failure message of any `DQRule` by specifying a custom message expression. Set `message_expr` to either a Spark SQL expression string or a Spark `Column` expression that returns a string-valued message when data fails a check. +If a check cannot be evaluated (for example, because it references invalid columns or uses an invalid SQL expression), the results report the default skip message instead of the custom message. When using custom message expressions: diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index e3a8ed7fc..aee861581 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -8,13 +8,12 @@ from pyspark.sql import DataFrame, Column, SparkSession from databricks.labs.dqx import check_funcs -from databricks.labs.dqx.errors import UnsafeSqlQueryError from databricks.labs.dqx.executor import DQCheckResult, DQRuleExecutorFactory from databricks.labs.dqx.rule import ( DQRule, ) from databricks.labs.dqx.schema.dq_result_schema import dq_result_item_schema -from databricks.labs.dqx.utils import get_column_name_or_alias, is_sql_query_safe +from databricks.labs.dqx.utils import get_column_name_or_alias logger = logging.getLogger(__name__) @@ -194,7 +193,7 @@ def _build_message_col(self, condition: Column, skipped: bool = False) -> Column message_expr = F.substr( F.expr(self.check.message_expr) if isinstance(self.check.message_expr, str) else self.check.message_expr, F.lit(1), - F.lit(_max_message_length) + F.lit(_max_message_length), ) return F.when(condition.isNotNull(), message_expr).otherwise(F.lit(None).cast("string")) diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index b7e047ac8..e783355d4 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -1,16 +1,16 @@ """Integration tests for custom message expressions on DQRule.""" import pyspark.sql.functions as F -import pytest from databricks.labs.dqx.engine import DQEngine -from databricks.labs.dqx.errors import UnsafeSqlQueryError from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule from databricks.labs.dqx import check_funcs from tests.integration.conftest import ( EXTRA_PARAMS, assert_df_equality_ignore_fingerprints as assert_df_equality, REPORTING_COLUMNS, + RUN_ID, + RUN_TIME, build_quality_violation, ) @@ -371,38 +371,38 @@ def test_metadata_for_each_column_with_custom_message(ws, spark): assert_df_equality(checked_df, expected_df) -def test_apply_checks_with_unsafe_message_expr_raises(ws, spark): - """A string message_expr containing an unsafe SQL keyword must be rejected when applied.""" +def test_apply_checks_skip_message_overrides_custom_message_for_invalid_sql_expression(ws, spark): + """test that message_expr is overridden when a sql_expression check is skipped due to an invalid expression""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) - test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + test_df = spark.createDataFrame([[1, 3, 5]], SCHEMA) rules = [ DQRowRule( - name="c_not_null", + name="invalid_sql_expression", criticality="error", - check_func=check_funcs.is_not_null, - column="c", - message_expr="concat('msg ', (select x from t where y = (drop table t)))", + check_func=check_funcs.sql_expression, + check_func_kwargs={"expression": "missing_col > 0"}, + message_expr="'Custom error: this message must be overridden by the skip message'", + user_metadata={"should_be": "IGNORED"}, ), ] - with pytest.raises(UnsafeSqlQueryError, match="not safe for execution"): - dq_engine.apply_checks(test_df, rules) - - -def test_apply_checks_by_metadata_with_unsafe_message_expr_raises(ws, spark): - """An unsafe message_expr supplied via metadata must also be rejected when applied.""" - dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) - test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) - - checks = [ + checked_df = dq_engine.apply_checks(test_df, rules) + expected_errors = [ { - "name": "c_not_null", - "criticality": "error", - "check": {"function": "is_not_null", "arguments": {"column": "c"}}, - "message_expr": "delete from audit", + "name": "invalid_sql_expression", + "message": "Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", + "columns": None, + "filter": None, + "function": "sql_expression", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {"should_be": "IGNORED"}, + "skipped": True, } ] - - with pytest.raises(UnsafeSqlQueryError, match="not safe for execution"): - dq_engine.apply_checks_by_metadata(test_df, checks) + expected_df = spark.createDataFrame( + [[1, 3, 5, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) From dbf5137a2718995e555d945cd9099dc71d21e8ac Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 8 Jun 2026 09:05:13 -0400 Subject: [PATCH 20/21] Add handling for invalid message expressions --- src/databricks/labs/dqx/manager.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index aee861581..f79c4f792 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -100,6 +100,13 @@ def has_invalid_filter(self) -> bool: """ return self._is_invalid_column(self.filter_condition) + @cached_property + def has_invalid_custom_message(self) -> bool: + """ + Returns a boolean indicating whether the custom message expression is invalid in the input DataFrame. + """ + return self._is_invalid_column(self.check.message_expr) + @cached_property def invalid_sql_expression(self) -> str | None: """ @@ -189,6 +196,9 @@ def _build_message_col(self, condition: Column, skipped: bool = False) -> Column if self.check.message_expr is None: return condition + if self.has_invalid_custom_message: + return F.lit(self._get_invalid_cols_message()) + _max_message_length = 500 message_expr = F.substr( F.expr(self.check.message_expr) if isinstance(self.check.message_expr, str) else self.check.message_expr, @@ -218,6 +228,12 @@ def _get_invalid_cols_message(self) -> str: f"Check evaluation skipped due to invalid check filter: '{self.check.filter}'" ) + if self.has_invalid_custom_message: + logger.warning(f"Skipping check '{self.check.name}' due to invalid custom message expression: '{self.check.message_expr}'") + invalid_cols_message_parts.append( + f"Check evaluation skipped due to invalid custom message expression: '{self.check.message_expr}'" + ) + if self.invalid_sql_expression: logger.warning( f"Skipping check '{self.check.name}' due to invalid sql expression: '{self.invalid_sql_expression}'" From 243002ca7954771bf3b2f143a4c4d4d14bdbbada Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 8 Jun 2026 10:11:07 -0400 Subject: [PATCH 21/21] Handle invalid message expressions --- src/databricks/labs/dqx/manager.py | 15 +- tests/integration/conftest.py | 24 +++ tests/integration/test_apply_checks.py | 247 ++++++++-------------- tests/integration/test_custom_messages.py | 177 ++++++++++++++-- 4 files changed, 290 insertions(+), 173 deletions(-) diff --git a/src/databricks/labs/dqx/manager.py b/src/databricks/labs/dqx/manager.py index f79c4f792..bfb76bc21 100644 --- a/src/databricks/labs/dqx/manager.py +++ b/src/databricks/labs/dqx/manager.py @@ -105,6 +105,8 @@ def has_invalid_custom_message(self) -> bool: """ Returns a boolean indicating whether the custom message expression is invalid in the input DataFrame. """ + if self.check.message_expr is None: + return False return self._is_invalid_column(self.check.message_expr) @cached_property @@ -196,9 +198,6 @@ def _build_message_col(self, condition: Column, skipped: bool = False) -> Column if self.check.message_expr is None: return condition - if self.has_invalid_custom_message: - return F.lit(self._get_invalid_cols_message()) - _max_message_length = 500 message_expr = F.substr( F.expr(self.check.message_expr) if isinstance(self.check.message_expr, str) else self.check.message_expr, @@ -229,9 +228,15 @@ def _get_invalid_cols_message(self) -> str: ) if self.has_invalid_custom_message: - logger.warning(f"Skipping check '{self.check.name}' due to invalid custom message expression: '{self.check.message_expr}'") + if isinstance(self.check.message_expr, str): + custom_message_detail = f": '{self.check.message_expr}'" + else: + custom_message_detail = "" + logger.warning( + f"Skipping check '{self.check.name}' due to invalid custom message expression{custom_message_detail}" + ) invalid_cols_message_parts.append( - f"Check evaluation skipped due to invalid custom message expression: '{self.check.message_expr}'" + f"Check evaluation skipped due to invalid custom message expression{custom_message_detail}" ) if self.invalid_sql_expression: diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 5b8231914..2b88f2b3c 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -121,6 +121,30 @@ def build_quality_violation( } +def build_skipped_violation( + name: str, + message: str, + columns: list[str] | None, + *, + function: str = "is_not_null", + filter_expr: str | None = None, + user_metadata: dict | None = None, +) -> dict[str, Any]: + """Helper for constructing expected entries for checks that were skipped during evaluation.""" + + return { + "name": name, + "message": message, + "columns": columns, + "filter": filter_expr, + "function": function, + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": user_metadata or {}, + "skipped": True, + } + + def assert_check_and_split_results( checked: DataFrame, good_df: DataFrame, diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 57a567d48..d7a48ccd0 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -28,6 +28,7 @@ EXTRA_PARAMS, RUN_ID, build_quality_violation, + build_skipped_violation, assert_check_and_split_results, assert_df_equality_ignore_fingerprints as assert_df_equality, generate_checks_with_rule_and_set_fingerprint_from_rules, @@ -10014,29 +10015,19 @@ def test_apply_checks_with_has_valid_schema_extra_columns_in_params(ws, spark): expected_schema = schema + REPORTING_COLUMNS - expected_skip_strict = { - "name": "has_valid_schema_strict", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']", - "columns": ["id", "v1", "missing_col"], - "filter": None, - "function": "has_valid_schema", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - } - - expected_skip_permissive = { - "name": "has_valid_schema_permissive", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']", - "columns": ["id", "v1", "missing_col"], - "filter": None, - "function": "has_valid_schema", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - } + expected_skip_strict = build_skipped_violation( + name="has_valid_schema_strict", + message="Check evaluation skipped due to invalid check columns: ['missing_col']", + columns=["id", "v1", "missing_col"], + function="has_valid_schema", + ) + + expected_skip_permissive = build_skipped_violation( + name="has_valid_schema_permissive", + message="Check evaluation skipped due to invalid check columns: ['missing_col']", + columns=["id", "v1", "missing_col"], + function="has_valid_schema", + ) expected = spark.createDataFrame( [ @@ -10163,76 +10154,49 @@ def test_apply_checks_skip_checks_with_missing_columns(ws, spark): {"key1": 1}, {"field1": 1}, [ - { - "name": "b_is_null_or_empty", - "message": "Check evaluation skipped due to invalid check filter: 'missing_col > 0'", - "columns": ["b"], - "filter": "missing_col > 0", - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "missing_col_is_null", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']", - "columns": ["missing_col"], - "filter": None, - "function": "is_not_null", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "missing_col_sql_expression", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']; " + build_skipped_violation( + name="b_is_null_or_empty", + message="Check evaluation skipped due to invalid check filter: 'missing_col > 0'", + columns=["b"], + function="is_not_null_and_not_empty", + filter_expr="missing_col > 0", + ), + build_skipped_violation( + name="missing_col_is_null", + message="Check evaluation skipped due to invalid check columns: ['missing_col']", + columns=["missing_col"], + ), + build_skipped_violation( + name="missing_col_sql_expression", + message="Check evaluation skipped due to invalid check columns: ['missing_col']; " "Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", - "columns": ["missing_col"], - "filter": None, - "function": "sql_expression", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "missing_col_is_unique", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']; " + columns=["missing_col"], + function="sql_expression", + ), + build_skipped_violation( + name="missing_col_is_unique", + message="Check evaluation skipped due to invalid check columns: ['missing_col']; " "Check evaluation skipped due to invalid check filter: 'missing_col > 0'", - "columns": ["missing_col"], - "filter": "missing_col > 0", - "function": "is_unique", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "invalid_col_sql_expression", - "message": "Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", - "columns": None, - "filter": None, - "function": "sql_expression", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, + columns=["missing_col"], + function="is_unique", + filter_expr="missing_col > 0", + ), + build_skipped_violation( + name="invalid_col_sql_expression", + message="Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", + columns=None, + function="sql_expression", + ), ], [ - { - "name": "missing_col_is_null_or_empty", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']", - "columns": ["missing_col"], - "filter": "a > 0", - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {"tag1": "value1", "tag2": "value2"}, - "skipped": True, - }, + build_skipped_violation( + name="missing_col_is_null_or_empty", + message="Check evaluation skipped due to invalid check columns: ['missing_col']", + columns=["missing_col"], + function="is_not_null_and_not_empty", + filter_expr="a > 0", + user_metadata={"tag1": "value1", "tag2": "value2"}, + ), ], ] ], @@ -10347,76 +10311,49 @@ def test_apply_checks_by_metadata_skip_checks_with_missing_columns(ws, spark): {"key1": 1}, {"field1": 1}, [ - { - "name": "b_is_null_or_empty", - "message": "Check evaluation skipped due to invalid check filter: 'missing_col > 0'", - "columns": ["b"], - "filter": "missing_col > 0", - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "missing_col_is_null", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']", - "columns": ["missing_col"], - "filter": None, - "function": "is_not_null", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "missing_col_sql_expression", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']; " + build_skipped_violation( + name="b_is_null_or_empty", + message="Check evaluation skipped due to invalid check filter: 'missing_col > 0'", + columns=["b"], + function="is_not_null_and_not_empty", + filter_expr="missing_col > 0", + ), + build_skipped_violation( + name="missing_col_is_null", + message="Check evaluation skipped due to invalid check columns: ['missing_col']", + columns=["missing_col"], + ), + build_skipped_violation( + name="missing_col_sql_expression", + message="Check evaluation skipped due to invalid check columns: ['missing_col']; " "Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", - "columns": ["missing_col"], - "filter": None, - "function": "sql_expression", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "missing_col_is_unique", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']; " + columns=["missing_col"], + function="sql_expression", + ), + build_skipped_violation( + name="missing_col_is_unique", + message="Check evaluation skipped due to invalid check columns: ['missing_col']; " "Check evaluation skipped due to invalid check filter: 'missing_col > 0'", - "columns": ["missing_col"], - "filter": "missing_col > 0", - "function": "is_unique", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, - { - "name": "invalid_col_sql_expression", - "message": "Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", - "columns": None, - "filter": None, - "function": "sql_expression", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - "skipped": True, - }, + columns=["missing_col"], + function="is_unique", + filter_expr="missing_col > 0", + ), + build_skipped_violation( + name="invalid_col_sql_expression", + message="Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", + columns=None, + function="sql_expression", + ), ], [ - { - "name": "missing_col_is_null_or_empty", - "message": "Check evaluation skipped due to invalid check columns: ['missing_col']", - "columns": ["missing_col"], - "filter": "a > 0", - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {"tag1": "value1", "tag2": "value2"}, - "skipped": True, - }, + build_skipped_violation( + name="missing_col_is_null_or_empty", + message="Check evaluation skipped due to invalid check columns: ['missing_col']", + columns=["missing_col"], + function="is_not_null_and_not_empty", + filter_expr="a > 0", + user_metadata={"tag1": "value1", "tag2": "value2"}, + ), ], ] ], diff --git a/tests/integration/test_custom_messages.py b/tests/integration/test_custom_messages.py index e783355d4..6e6d65ffa 100644 --- a/tests/integration/test_custom_messages.py +++ b/tests/integration/test_custom_messages.py @@ -9,9 +9,8 @@ EXTRA_PARAMS, assert_df_equality_ignore_fingerprints as assert_df_equality, REPORTING_COLUMNS, - RUN_ID, - RUN_TIME, build_quality_violation, + build_skipped_violation, ) SCHEMA = "a: int, b: int, c: int" @@ -371,6 +370,162 @@ def test_metadata_for_each_column_with_custom_message(ws, spark): assert_df_equality(checked_df, expected_df) +def test_apply_checks_skip_message_for_invalid_custom_message(ws, spark): + """A check with a valid evaluation but an unresolvable custom message expression should be skipped.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 5]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + message_expr="concat('Custom error for value: ', missing_col)", + user_metadata={"should_be": "IGNORED"}, + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_errors = [ + build_skipped_violation( + name="c_not_null", + message="Check evaluation skipped due to invalid custom message expression: " + "'concat('Custom error for value: ', missing_col)'", + columns=["c"], + user_metadata={"should_be": "IGNORED"}, + ) + ] + expected_df = spark.createDataFrame( + [[1, 3, 5, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_metadata_skip_message_for_invalid_custom_message(ws, spark): + """A metadata-defined check with an unresolvable custom message expression should be skipped.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 5]], SCHEMA) + + checks = [ + { + "name": "c_not_null", + "criticality": "error", + "message_expr": "concat('Custom error for value: ', missing_col)", + "check": {"function": "is_not_null", "arguments": {"column": "c"}}, + } + ] + + checked_df = dq_engine.apply_checks_by_metadata(test_df, checks) + expected_errors = [ + build_skipped_violation( + name="c_not_null", + message="Check evaluation skipped due to invalid custom message expression: " + "'concat('Custom error for value: ', missing_col)'", + columns=["c"], + ) + ] + expected_df = spark.createDataFrame( + [[1, 3, 5, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_skip_message_for_invalid_column_message_expr(ws, spark): + """A Column message_expr that cannot be resolved should be skipped without injecting the Column repr.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 5]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + message_expr=F.col("missing_col"), + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_errors = [ + build_skipped_violation( + name="c_not_null", + message="Check evaluation skipped due to invalid custom message expression", + columns=["c"], + ) + ] + expected_df = spark.createDataFrame( + [[1, 3, 5, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_skip_message_combines_invalid_columns_and_custom_message(ws, spark): + """When both the check column and the custom message are invalid, both reasons should be reported.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 5]], SCHEMA) + + rules = [ + DQRowRule( + name="missing_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="missing_col", + message_expr="concat('Custom error for value: ', another_missing_col)", + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_errors = [ + build_skipped_violation( + name="missing_not_null", + message="Check evaluation skipped due to invalid check columns: ['missing_col']; " + "Check evaluation skipped due to invalid custom message expression: " + "'concat('Custom error for value: ', another_missing_col)'", + columns=["missing_col"], + ) + ] + expected_df = spark.createDataFrame( + [[1, 3, 5, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + +def test_apply_checks_valid_custom_message_not_skipped(ws, spark): + """A check without invalid fields and with a valid custom message must not be marked skipped.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, None]], SCHEMA) + + rules = [ + DQRowRule( + name="c_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="c", + message_expr="concat('Custom error for c: ', coalesce(cast(c as string), 'null'))", + ), + ] + + checked_df = dq_engine.apply_checks(test_df, rules) + expected_errors = [ + build_quality_violation( + name="c_not_null", + message="Custom error for c: null", + columns=["c"], + function="is_not_null", + ) + ] + expected_df = spark.createDataFrame( + [[1, 3, None, expected_errors, None]], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked_df, expected_df) + + def test_apply_checks_skip_message_overrides_custom_message_for_invalid_sql_expression(ws, spark): """test that message_expr is overridden when a sql_expression check is skipped due to an invalid expression""" dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) @@ -389,17 +544,13 @@ def test_apply_checks_skip_message_overrides_custom_message_for_invalid_sql_expr checked_df = dq_engine.apply_checks(test_df, rules) expected_errors = [ - { - "name": "invalid_sql_expression", - "message": "Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", - "columns": None, - "filter": None, - "function": "sql_expression", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {"should_be": "IGNORED"}, - "skipped": True, - } + build_skipped_violation( + name="invalid_sql_expression", + message="Check evaluation skipped due to invalid sql expression: 'missing_col > 0'", + columns=None, + function="sql_expression", + user_metadata={"should_be": "IGNORED"}, + ) ] expected_df = spark.createDataFrame( [[1, 3, 5, expected_errors, None]],