feat: add filter attribute in rules generated from ODCS

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

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)

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

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

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: