From f762341c6c50593be147cf703bdf8d5e519ec9ae Mon Sep 17 00:00:00 2001 From: Pierre Monnet Date: Wed, 24 Dec 2025 17:53:56 +0100 Subject: [PATCH 1/3] feat: add filter attribute in rules generated from ODCS --- .../datacontract/contract_rules_generator.py | 10 +++++++--- .../test_datacontract_integration.py | 2 +- tests/resources/sample_datacontract.yaml | 14 ++++++++++++++ tests/unit/test_datacontract_generator.py | 17 +++++++++++++++-- 4 files changed, 37 insertions(+), 6 deletions(-) diff --git a/src/databricks/labs/dqx/datacontract/contract_rules_generator.py b/src/databricks/labs/dqx/datacontract/contract_rules_generator.py index ebe5397f3..69902ce94 100644 --- a/src/databricks/labs/dqx/datacontract/contract_rules_generator.py +++ b/src/databricks/labs/dqx/datacontract/contract_rules_generator.py @@ -991,12 +991,12 @@ def _build_explicit_rule_from_implementation( Raises ODCSContractError if the implementation structure is invalid. """ try: - check, name, criticality = self._extract_impl_attributes(impl, default_criticality) + check, name, criticality, filter_rule = self._extract_impl_attributes(impl, default_criticality) if check is None: logger.warning("Implementation missing 'check' attribute, skipping rule") return None - return self._build_rule_dict(check, name, criticality, schema_name, property_name, odcs) + return self._build_rule_dict(check, name, criticality, filter_rule, schema_name, property_name, odcs) except (AttributeError, KeyError, TypeError) as e: # Malformed contract structure - fail fast raise ODCSContractError( @@ -1016,13 +1016,15 @@ def _extract_impl_attributes(self, impl: str | dict[str, Any] | None, default_cr check = impl.get("check") name = impl.get("name", "unnamed_rule") criticality = impl.get("criticality", default_criticality) - return check, name, criticality + filter_rule: str | None = impl.get("filter") + return check, name, criticality, filter_rule def _build_rule_dict( self, check_dict: dict, name: str, criticality: str, + filter_rule: str | None, schema_name: str, property_name: str | None, odcs: OpenDataContractStandard, @@ -1044,4 +1046,6 @@ def _build_rule_dict( "criticality": criticality, "user_metadata": user_metadata, } + if filter_rule: + rule["filter"] = filter_rule return rule diff --git a/tests/integration/test_datacontract_integration.py b/tests/integration/test_datacontract_integration.py index b1664413d..10a0aa79d 100644 --- a/tests/integration/test_datacontract_integration.py +++ b/tests/integration/test_datacontract_integration.py @@ -22,7 +22,7 @@ def test_generate_rules_with_text_processing(self, ws, spark, sample_contract_pa ) # Verify rules were generated - assert len(rules) > 35 # More than predefined + explicit rules due to text-based rules + assert len(rules) > 36 # More than predefined + explicit rules due to text-based rules # Verify that text-based rules were processed by LLM # The sample contract has a text expectation about duplicate sensor readings diff --git a/tests/resources/sample_datacontract.yaml b/tests/resources/sample_datacontract.yaml index 01c8ac557..0ecae0bfd 100644 --- a/tests/resources/sample_datacontract.yaml +++ b/tests/resources/sample_datacontract.yaml @@ -310,3 +310,17 @@ schema: description: | The dataset should not contain duplicate sensor readings for the same sensor_id and reading_timestamp combination within a 1-second window. + + # Check quality for subsets of data (filter condition)) + - type: custom + engine: dqx + description: Ensure technician_id and alert_email are present for high alert levels + implementation: + criticality: warn + filter: alert_level in ('high', 'critical') + name: technician_and_email_are_mandatory_for_high_alerts + check: + function: is_not_null_and_not_empty + for_each_column: + - technician_id + - alert_email \ No newline at end of file diff --git a/tests/unit/test_datacontract_generator.py b/tests/unit/test_datacontract_generator.py index 325434d46..a81db3d4d 100644 --- a/tests/unit/test_datacontract_generator.py +++ b/tests/unit/test_datacontract_generator.py @@ -610,6 +610,19 @@ def get_expected_contract_rules(self): 'rule_type': 'explicit', }, }, + { + 'check': {'function': 'is_not_null_and_not_empty', 'for_each_column': ['technician_id', 'alert_email']}, + 'name': 'technician_and_email_are_mandatory_for_high_alerts', + 'criticality': 'warn', + 'user_metadata': { + 'contract_id': 'urn:datacontract:sensors:iot_sensor_data', + 'contract_version': '2.1.0', + 'odcs_version': 'v3.0.2', + 'schema': 'sensor_readings', + 'rule_type': 'explicit', + }, + 'filter': 'alert_level in (\'high\', \'critical\')', + }, ] def create_basic_contract( @@ -937,8 +950,8 @@ def test_generate_rules_and_skip_predefined_rules(self, generator, sample_contra ) # Should only have explicit rules (no predefined rules) - # Sample ODCS v3.x contract has 7 explicit DQX rules (5 property-level + 2 schema-level) - assert len(rules) == 7 + # Sample ODCS v3.x contract has 8 explicit DQX rules (5 property-level + 3 schema-level) + assert len(rules) == 8 # Verify all rules are explicit for rule in rules: From 26db37e8bf96531e4dccaa4b128e344d0d86d4e3 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 28 Dec 2025 11:28:56 +0100 Subject: [PATCH 2/3] Update tests/resources/sample_datacontract.yaml Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/resources/sample_datacontract.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/resources/sample_datacontract.yaml b/tests/resources/sample_datacontract.yaml index 0ecae0bfd..70660fa17 100644 --- a/tests/resources/sample_datacontract.yaml +++ b/tests/resources/sample_datacontract.yaml @@ -311,7 +311,7 @@ schema: The dataset should not contain duplicate sensor readings for the same sensor_id and reading_timestamp combination within a 1-second window. - # Check quality for subsets of data (filter condition)) + # Check quality for subsets of data (filter condition) - type: custom engine: dqx description: Ensure technician_id and alert_email are present for high alert levels From b989d95560fb4adc5d459961317144f70f335e7d Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 28 Dec 2025 14:52:54 +0100 Subject: [PATCH 3/3] improve test --- tests/integration/test_ai_rules_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_ai_rules_generator.py b/tests/integration/test_ai_rules_generator.py index 3c3c3ff06..fe41b61ec 100644 --- a/tests/integration/test_ai_rules_generator.py +++ b/tests/integration/test_ai_rules_generator.py @@ -94,7 +94,7 @@ def not_ends_with_suffix(column: str, suffix: str): custom_check_functions = {"not_ends_with_suffix": not_ends_with_suffix} - user_input = USER_INPUT + "\nEmail address must not end with '@gmail.com'." + user_input = USER_INPUT + "\nEmail address must not end with '@gmail.com'. Use not_ends_with_suffix function." generator = DQGenerator(ws, spark, custom_check_functions=custom_check_functions) actual_checks = generator.generate_dq_rules_ai_assisted(user_input=user_input)