Fixed resolution of pii detection check function

src/databricks/labs/dqx/checks_resolver.py

@@ -7,6 +7,14 @@
 
 from databricks.labs.dqx import check_funcs
 from databricks.labs.dqx.geo import check_funcs as geo_check_funcs
+
+try:
+    from databricks.labs.dqx.pii import pii_detection_funcs as pii_check_funcs
+
+    PII_ENABLED = True
+except ImportError:
+    PII_ENABLED = False
+
 from databricks.labs.dqx.errors import InvalidCheckError
 
 logger = logging.getLogger(__name__)
@@ -32,8 +40,11 @@ def resolve_check_function(
     logger.debug(f"Resolving function: {function_name}")
     func = getattr(check_funcs, function_name, None)  # resolve using predefined checks first
     if not func:
-        # resolve using predefined geo checks, requires Databricks serverless or DBR >= 17.1
+        # try to resolve using predefined geo checks, requires Databricks serverless or DBR >= 17.1
         func = getattr(geo_check_funcs, function_name, None)
+    if not func and PII_ENABLED:
+        # try to resolve using predefined pii detection checks
+        func = getattr(pii_check_funcs, function_name, None)
     if not func and custom_check_functions:
         func = custom_check_functions.get(function_name)  # returns None if not found
     if fail_on_missing and not func:

tests/e2e/notebooks/pii_detection_notebook.py

@@ -33,8 +33,13 @@ def test_import_pii_module_fails_without_installation():
 
 # COMMAND ----------
 
+import yaml
+from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.engine import DQEngine
+from databricks.labs.dqx.rule import DQRowRule
 import pyspark.sql.functions as F
 from databricks.labs.dqx.pii.nlp_engine_config import NLPEngineConfig
+from databricks.labs.dqx.check_funcs import is_not_null
 from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii
 from chispa import assert_df_equality
 
@@ -222,3 +227,75 @@ def test_does_not_contain_pii_with_custom_nlp_config_dict():
     assert_df_equality(actual, expected, transforms=transforms)
 
 test_does_not_contain_pii_with_custom_nlp_config_dict()
+
+# COMMAND ----------
+# DBTITLE 1,test_does_not_contain_pii_apply_checks_by_metadata_and_split
+
+def test_does_not_contain_pii_apply_checks_by_metadata_and_split():
+    checks = yaml.safe_load("""
+    - criticality: error
+      check:
+        function: is_not_null
+        arguments:
+          column: val
+    - criticality: error
+      check:
+        function: does_not_contain_pii
+        arguments:
+          column: val
+    """)
+
+    # Initialize the DQX engine:
+    dq_engine = DQEngine(WorkspaceClient())
+
+    # Create some sample data:
+    data = [
+        ["My name is John Smith"],
+        ["The sky is blue, road runner"],
+        ["Jane Smith sent an email to sara@info.com"],
+        [None],
+    ]
+    df = spark.createDataFrame(data, "val string")
+
+    # Run the checks and display the output:
+    valid_df, invalid_df = dq_engine.apply_checks_by_metadata_and_split(df, checks)
+    assert valid_df.count() == 1
+    assert invalid_df.count() == 3
+
+test_does_not_contain_pii_apply_checks_by_metadata_and_split()
+
+# COMMAND ----------
+# DBTITLE 1,test_does_not_contain_pii_apply_checks_and_split
+
+def test_does_not_contain_pii_apply_checks_and_split():
+    checks = [
+        DQRowRule(
+            criticality="error",
+            check_func=is_not_null,
+            column="val",
+        ),
+        DQRowRule(
+            criticality="error",
+            check_func=does_not_contain_pii,
+            column="val",
+        )
+    ]
+
+    # Initialize the DQX engine:
+    dq_engine = DQEngine(WorkspaceClient())
+
+    # Create some sample data:
+    data = [
+        ["My name is John Smith"],
+        ["The sky is blue, road runner"],
+        ["Jane Smith sent an email to sara@info.com"],
+        [None],
+    ]
+    df = spark.createDataFrame(data, "val string")
+
+    # Run the checks and display the output:
+    valid_df, invalid_df = dq_engine.apply_checks_and_split(df, checks)
+    assert valid_df.count() == 1
+    assert invalid_df.count() == 3
+
+test_does_not_contain_pii_apply_checks_and_split()