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()

tests/unit/test_resolve_check_function.py

@@ -1,6 +1,9 @@
 import textwrap
+import sys
+import importlib
 
 import pytest
+from databricks.labs import dqx
 from databricks.labs.dqx.checks_resolver import resolve_check_function, resolve_custom_check_functions_from_path
 from databricks.labs.dqx.errors import InvalidCheckError
 
@@ -101,3 +104,31 @@ def test_resolve_custom_check_functions_from_path_with_dependency(tmp_path):
 
     func = resolve_custom_check_functions_from_path({"main_func": str(main_path)})["main_func"]
     assert func() == "dependency ok"
+
+
+def test_pii_module_import_failure(monkeypatch):
+    """Test that the code handles gracefully when PII module is not available."""
+    # Save the original module state
+    original_pii_module = sys.modules.get('databricks.labs.dqx.pii')
+
+    # Remove the PII module from sys.modules to simulate ImportError
+    monkeypatch.setitem(sys.modules, 'databricks.labs.dqx.pii', None)
+
+    # Force module reload to trigger the import logic
+    importlib.reload(dqx.checks_resolver)
+
+    # Verify that PII_ENABLED is False when import fails
+    assert dqx.checks_resolver.PII_ENABLED is False
+
+    # Verify that standard check functions still work
+    func = dqx.checks_resolver.resolve_check_function('is_not_null')
+    assert func is not None
+
+    # Verify that a non-existent function (like PII) cannot be resolved when PII is disabled
+    func = dqx.checks_resolver.resolve_check_function('some_missing_func', fail_on_missing=False)
+    assert func is None
+
+    # Restore the module to its original state
+    if original_pii_module is not None:
+        sys.modules['databricks.labs.dqx.pii'] = original_pii_module
+    importlib.reload(dqx.checks_resolver)