diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 84bbc80b1..53f7b3fab 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -38,7 +38,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install hatch run: pip install hatch==1.9.4 @@ -90,7 +90,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install hatch run: pip install hatch==1.9.4 @@ -122,7 +122,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install hatch run: pip install hatch==1.9.4 @@ -156,7 +156,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install hatch run: pip install hatch==1.9.4 diff --git a/.github/workflows/downstreams.yml b/.github/workflows/downstreams.yml index 1de68d7ad..e796115e9 100644 --- a/.github/workflows/downstreams.yml +++ b/.github/workflows/downstreams.yml @@ -39,7 +39,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install toolchain run: | diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 0d1b8c58a..77ff9dd4d 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -29,7 +29,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install hatch run: pip install hatch==1.9.4 @@ -78,7 +78,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install hatch run: pip install hatch==1.9.4 @@ -111,7 +111,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Install hatch run: pip install hatch==1.9.4 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 3f1b14724..c989b8edd 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -23,7 +23,7 @@ jobs: with: cache: 'pip' cache-dependency-path: '**/pyproject.toml' - python-version: '3.10' + python-version: '3.11' - name: Build wheels run: | diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index 065467f97..a8ad7ce17 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -1,30 +1,27 @@ # Databricks notebook source # MAGIC %md # MAGIC # Using DQX for PII Detection -# MAGIC Increased regulation makes Databricks customers responsible for any Personally Identifiable Information (PII) stored in Unity Catalog. While [Lakehouse Monitoring](https://docs.databricks.com/aws/en/lakehouse-monitoring/data-classification#discover-sensitive-data) can identify sensitive data in-place, many customers need to proactively quarantine or anonymize PII before writing the data to Delta. +# MAGIC Increased regulation makes Databricks customers responsible for any Personally Identifiable Information (PII) stored in Unity Catalog. Companies need to be able to perform PII detection for data at-rest and in-transit to proactively quarantine or anonymize PII before persisting the data. # MAGIC -# MAGIC [Databricks Labs' DQX project](https://databrickslabs.github.io/dqx/) provides in-flight data quality monitoring for Spark `DataFrames`. Customers can apply checks, get row-level metadata, and quarantine failing records. Workloads can use DQX's built-in checks or custom user-defined functions. -# MAGIC -# MAGIC In this notebook, we'll use DQX with a custom function to detect PII in JSON strings. +# MAGIC DQX provides in-flight data quality monitoring for Spark `DataFrames`. You can apply checks, get row-level metadata, and quarantine failing records. Workloads can also use DQX's built-in functions to check `DataFrames` for PII. # COMMAND ---------- # MAGIC %md -# MAGIC ## Prerequisites -# MAGIC This notebook uses [Presidio](https://microsoft.github.io/presidio/) to detect PII in strings. To run this notebook: -# MAGIC - Use DBR 15.4LTS -# MAGIC - Install [SpaCy](https://spacy.io/usage/models#download) as a cluster-scoped library +# MAGIC # Install DQX with PII extras +# MAGIC +# MAGIC To enable PII detection quality checking, DQX has to be installed with `pii` extras: +# MAGIC +# MAGIC `%pip install databricks-labs-dqx[pii]` # COMMAND ---------- dbutils.widgets.text("test_library_ref", "", "Test Library Ref") if dbutils.widgets.get("test_library_ref") != "": - %pip install '{dbutils.widgets.get("test_library_ref")}' + %pip install 'databricks-labs-dqx[pii] @ {dbutils.widgets.get("test_library_ref")}' else: - %pip install databricks-labs-dqx - -%pip install presidio_analyzer numpy==1.23.5 + %pip install databricks-labs-dqx[pii] # COMMAND ---------- @@ -32,101 +29,145 @@ # COMMAND ---------- -import json -import pandas as pd - -from pyspark.sql.functions import concat_ws, col, lit, pandas_udf -from pyspark.sql import Column -from presidio_analyzer import AnalyzerEngine from databricks.sdk import WorkspaceClient from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.rule import DQRowRule -from databricks.labs.dqx.check_funcs import make_condition +from databricks.labs.dqx.pii.nlp_engine_config import NLPEngineConfig +from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii # COMMAND ---------- # MAGIC %md -# MAGIC ## Creating the Presidio analyzer -# MAGIC First, we'll use Presidio's `AnalyzerEngine` to define a function that checks values for PII. For any PII detected, the `entity_mapping` will contain the type of PII identified and a confidence score. +# MAGIC ## Detecting PII with DQX +# MAGIC DQX supports built-in PII detection using Presidio's `AnalyzerEngine` to define a function that checks values for PII. For any PII detected, the `entity_mapping` will contain the type of PII identified and a confidence score. # COMMAND ---------- -# Create the Presidio analyzer: -analyzer = AnalyzerEngine() - -# Get the list of entities to download the model: -entities = analyzer.get_supported_entities() - -# Create a wrapper function to generate the entity mapping results: -def get_entity_mapping(data: str) -> str | None: - if data: - # Run the Presidio analyzer to detect PII in the string: - results = analyzer.analyze( - text=data, - entities=["PERSON", "EMAIL_ADDRESS"], - language='en', - score_threshold=0.5, - ) - if results != []: - output = [] - # Validate and return the results: - for result in results: - # Ignore if the result is low confidence: - if result.score < 0.5: - continue - # Append the result to the output: - output.append({ - "entity_type": result.entity_type, - "start": int(result.start), - "end": int(result.end), - "score": float(result.score), - }) - if output != []: - # Return the results as JSON: - return json.dumps(output) - return None +# Define the DQX rule: +checks = [ + DQRowRule( + criticality="error", + check_func=does_not_contain_pii, + column="val", + name="does_not_contain_pii", + ) +] + +# 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: +checked_df = dq_engine.apply_checks(df, checks) +display(checked_df) # COMMAND ---------- # MAGIC %md -# MAGIC ## Creating a Pandas UDF -# MAGIC We can call `get_entity_mapping` on DataFrame rows using a Pandas user-defined function. This provides good performance with batched execution over the arriving records. +# MAGIC ## Configuring the PII detection settings +# MAGIC DQX supports several configurable settings which control PII detection: +# MAGIC - `threshold` controls the specificity of the PII detection (higher values give more specificity with less sensitivity) +# MAGIC - `entities` specifies which [entity types](https://microsoft.github.io/presidio/supported_entities/) are marked as PII +# MAGIC - `language` sets the detection language +# MAGIC - `nlp_engine_config` sets various properties of the Presidio analyzer's [named entity recognition model](https://microsoft.github.io/presidio/samples/python/ner_model_configuration/) # COMMAND ---------- -# Register a pandas UDF to run the analyzer: -@pandas_udf('string') -def contains_pii(batch: pd.Series) -> pd.Series: - # Apply `get_entity_mapping` to each value: - return batch.map(get_entity_mapping) +# Use a built-in NLP configuration for detecting PII: +nlp_engine_config = NLPEngineConfig.SPACY_MEDIUM + +checks = [ + # Define a PII check with a lower threshold (more sensitivity): + DQRowRule( + criticality="error", + check_func=does_not_contain_pii, + check_func_kwargs={"threshold": 0.5}, + column="val", + name="does_not_contain_pii_lower_threshold", + ), + # Define a PII check with a subset of named entities: + DQRowRule( + criticality="error", + check_func=does_not_contain_pii, + check_func_kwargs={ + "entities": ["EMAIL_ADDRESS"], + }, + column="val", + name="contains_email_address_data", + ), + # Define a PII check with a built-in named-entity recognizer (SpaCy medium): + DQRowRule( + criticality="error", + check_func=does_not_contain_pii, + check_func_kwargs={ + "entities": ["PERSON", "LOCATION"], + "nlp_engine_config": NLPEngineConfig.SPACY_MEDIUM + }, + column="val", + name="contains_person_or_address_data", + ), + # Define a PII check with a built-in named-entity recognizer (SpaCy medium): + DQRowRule( + criticality="error", + check_func=does_not_contain_pii, + check_func_kwargs={ + "entities": ["PERSON", "LOCATION"], + "nlp_engine_config": NLPEngineConfig.SPACY_MEDIUM + }, + column="val", + name="contains_person_or_address_data", + ), +] + +# Initialize the DQX engine: +dq_engine = DQEngine(WorkspaceClient()) + +# Create some sample data: +data = [ + ["My name is John Smith and I live at 123 Main St New York, NY 07008"], + ["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: +checked_df = dq_engine.apply_checks(df, checks) +display(checked_df) # COMMAND ---------- # MAGIC %md -# MAGIC ## Making and applying a DQX condition -# MAGIC Once our Presidio algorithm can be called as a Spark UDF, we can use DQX's `make_condition` to implement a custom check that identifies PII and generates row-level metadata about the `json` keys and types of PII identified. +# MAGIC ## Using a custom named entity recognizer +# MAGIC DQX supports custom named entity recognizers passed as Python dictionaries. All dependencies must be pre-loaded for use with DQX's built-in PII detection checks. +# MAGIC +# MAGIC ***WARNING:** Using custom named entity recognizers can significantly degrade performance of quality checking at scale. Sample data or use smaller models for best performance. Run checks on non-serverless compute when using large named entity recognizers.* # COMMAND ---------- -def does_not_contain_pii(column: str) -> Column: - # Define a PII detection expression calling the pandas UDF: - pii_info = contains_pii(col(column)) - - # Return the DQX condition that uses the PII detection expression: - return make_condition( - pii_info.isNotNull(), - concat_ws( - ' ', - lit(column), - lit('contains pii with the following info:'), - pii_info - ), - f'{column}_contains_pii' - ) +# Define the NLP configuration: +nlp_engine_config = { + "nlp_engine_name": "spacy", + "models": [{"lang_code": "en", "model_name": "en_core_web_md"}] +} -# Define the DQX rule: checks = [ - DQRowRule(criticality='error', check_func=does_not_contain_pii, column='val') + # Define a PII check with a custom named-entity recognizer (Stanford De-Identifier Base): + DQRowRule( + criticality="error", + check_func=does_not_contain_pii, + check_func_kwargs={"nlp_engine_config": nlp_engine_config}, + column="val", + name="contains_pii_custom_recognizer", + ), ] # Initialize the DQX engine: @@ -134,11 +175,12 @@ def does_not_contain_pii(column: str) -> Column: # 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'] + ["My name is John Smith and I live at 123 Main St New York, NY 07008"], + ["The sky is blue, road runner"], + ["Jane Smith sent an email to sara@info.com"], + [None], ] -df = spark.createDataFrame(data, 'val string') +df = spark.createDataFrame(data, "val string") # Run the checks and display the output: checked_df = dq_engine.apply_checks(df, checks) diff --git a/demos/dqx_streaming_demo_diy.py b/demos/dqx_streaming_demo_diy.py index 862123bee..dbf99df3f 100644 --- a/demos/dqx_streaming_demo_diy.py +++ b/demos/dqx_streaming_demo_diy.py @@ -180,7 +180,6 @@ from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient -from pyspark.sql import DataFrame dq_engine = DQEngine(WorkspaceClient()) diff --git a/demos/dqx_streaming_demo_native.py b/demos/dqx_streaming_demo_native.py index 5c7771323..1b61bd8f6 100644 --- a/demos/dqx_streaming_demo_native.py +++ b/demos/dqx_streaming_demo_native.py @@ -1,5 +1,4 @@ # Databricks notebook source -# MAGIC # MAGIC %md # MAGIC ### Environment Setup for DQX with Structured Streaming # MAGIC @@ -167,13 +166,22 @@ silver_checkpoint = dbutils.widgets.get("silver_checkpoint") quarantine_checkpoint = dbutils.widgets.get("quarantine_checkpoint") + uuid = uuid4() +bronze_volume_name = f"bronze_{uuid}".replace("-", "_") +spark.sql(f"CREATE VOLUME IF NOT EXISTS {catalog}.{schema}.{bronze_volume_name}") +bronze_path = f"/Volumes/{catalog}/{schema}/{bronze_volume_name}" silver_table = f"{catalog}.{schema}.`silver_{uuid}`" quarantine_table = f"{catalog}.{schema}.`quarantine_{uuid}`" +print(f"Demo Bronze Path: {bronze_path}") print(f"Demo Silver Table: {silver_table}") print(f"Demo Quarantine Table: {quarantine_table}") +# prepare sample test data +bronze_df = spark.read.load("/databricks-datasets/delta-sharing/samples/nyctaxi_2019") +bronze_df.write.mode("overwrite").format("parquet").save(f"{bronze_path}/data") + # COMMAND ---------- # MAGIC %md @@ -190,17 +198,11 @@ from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.config import InputConfig, OutputConfig from databricks.sdk import WorkspaceClient -from pyspark.sql import DataFrame - -# prepare sample test data -bronze_loc = "/tmp/dq/bronze" -bronze_df = spark.read.load("/databricks-datasets/delta-sharing/samples/nyctaxi_2019") -bronze_df.write.mode("overwrite").format("parquet").save(f"{bronze_loc}/data") input_config=InputConfig( - location=f"{bronze_loc}/data", + location=f"{bronze_path}/data", format="cloudFiles", - options={"cloudFiles.format": "parquet", "cloudFiles.schemaLocation": f"{bronze_loc}/schema"}, + options={"cloudFiles.format": "parquet", "cloudFiles.schemaLocation": f"{bronze_path}/schema"}, is_streaming=True ) @@ -246,8 +248,8 @@ # COMMAND ---------- -dbutils.fs.rm(bronze_loc, True) dbutils.fs.rm(silver_checkpoint, True) dbutils.fs.rm(quarantine_checkpoint, True) +spark.sql(f"DROP VOLUME IF EXISTS {catalog}.{schema}.{bronze_volume_name}") spark.sql(f"DROP TABLE IF EXISTS {silver_table}") spark.sql(f"DROP TABLE IF EXISTS {quarantine_table}") \ No newline at end of file diff --git a/docs/dqx/docs/reference/quality_rules.mdx b/docs/dqx/docs/reference/quality_rules.mdx index 494a8b424..c938e9d07 100644 --- a/docs/dqx/docs/reference/quality_rules.mdx +++ b/docs/dqx/docs/reference/quality_rules.mdx @@ -15,30 +15,30 @@ You can also define your own custom checks (see [Creating custom checks](#creati
**Available row-level checks** -| Check | Description | Arguments | -| ---------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `is_not_null` | Checks whether the values in the input column are not null. | `column`: column to check (can be a string column name or a column expression) | -| `is_not_empty` | Checks whether the values in the input column are not empty (but may be null). | `column`: column to check (can be a string column name or a column expression) | -| `is_not_null_and_not_empty` | Checks whether the values in the input column are not null and not empty. | `column`: column to check (can be a string column name or a column expression); `trim_strings`: optional boolean flag to trim spaces from strings | -| `is_in_list` | Checks whether the values in the input column are present in the list of allowed values (null values are allowed). This check is not suited for large lists of allowed values. In such cases, it’s recommended to use the `foreign_key` dataset-level check instead. | `column`: column to check (can be a string column name or a column expression); `allowed`: list of allowed values | -| `is_not_null_and_is_in_list` | Checks whether the values in the input column are not null and present in the list of allowed values. This check is not suited for large lists of allowed values. In such cases, it’s recommended to use the `foreign_key` dataset-level check instead. | `column`: column to check (can be a string column name or a column expression); `allowed`: list of allowed values | -| `is_not_null_and_not_empty_array` | Checks whether the values in the array input column are not null and not empty. | `column`: column to check (can be a string column name or a column expression) | -| `is_in_range` | Checks whether the values in the input column are in the provided range (inclusive of both boundaries). | `column`: column to check (can be a string column name or a column expression); `min_limit`: min limit as number, date, timestamp, column name or sql expression; `max_limit`: max limit as number, date, timestamp, column name or sql expression | -| `is_not_in_range` | Checks whether the values in the input column are outside the provided range (inclusive of both boundaries). | `column`: column to check (can be a string column name or a column expression); `min_limit`: min limit as number, date, timestamp, column name or sql expression; `max_limit`: max limit as number, date, timestamp, column name or sql expression | -| `is_not_less_than` | Checks whether the values in the input column are not less than the provided limit. | `column`: column to check (can be a string column name or a column expression); `limit`: limit as number, date, timestamp, column name or sql expression | -| `is_not_greater_than` | Checks whether the values in the input column are not greater than the provided limit. | `column`: column to check (can be a string column name or a column expression); `limit`: limit as number, date, timestamp, column name or sql expression | -| `is_valid_date` | Checks whether the values in the input column have valid date formats. | `column`: column to check (can be a string column name or a column expression); `date_format`: optional date format (e.g. 'yyyy-mm-dd') | -| `is_valid_timestamp` | Checks whether the values in the input column have valid timestamp formats. | `column`: column to check (can be a string column name or a column expression); `timestamp_format`: optional timestamp format (e.g. 'yyyy-mm-dd HH:mm:ss') | -| `is_not_in_future` | Checks whether the values in the input column contain a timestamp that is not in the future, where 'future' is defined as current_timestamp + offset (in seconds). | `column`: column to check (can be a string column name or a column expression); `offset`: offset to use; `curr_timestamp`: current timestamp, if not provided current_timestamp() function is used | -| `is_not_in_near_future` | Checks whether the values in the input column contain a timestamp that is not in the near future, where 'near future' is defined as greater than the current timestamp but less than the current_timestamp + offset (in seconds). | `column`: column to check (can be a string column name or a column expression); `offset`: offset to use; `curr_timestamp`: current timestamp, if not provided current_timestamp() function is used | -| `is_older_than_n_days` | Checks whether the values in one input column are at least N days older than the values in another column. | `column`: column to check (can be a string column name or a column expression); `days`: number of days; `curr_date`: current date, if not provided current_date() function is used; `negate`: if the condition should be negated | -| `is_older_than_col2_for_n_days` | Checks whether the values in one input column are at least N days older than the values in another column. | `column1`: first column to check (can be a string column name or a column expression); `column2`: second column to check (can be a string column name or a column expression); `days`: number of days; `negate`: if the condition should be negated | -| `regex_match` | Checks whether the values in the input column match a given regex. | `column`: column to check (can be a string column name or a column expression); regex: regex to check; `negate`: if the condition should be negated (true) or not | -| `is_valid_ipv4_address` | Checks whether the values in the input column have valid IPv4 address format. | `column` to check (can be a string column name or a column expression) | -| `is_ipv4_address_in_cidr` | Checks whether the values in the input column have valid IPv4 address format and fall within the given CIDR block. | `column`: column to check (can be a string column name or a column expression); `cidr_block`: CIDR block string | -| `sql_expression` | Checks whether the values meet the condition provided as an SQL expression, e.g. `a = 'str1' and a > b`. SQL expressions are evaluated at runtime, so ensure that the expression is safe and that functions used within it (e.g. h3_ischildof, division) do not throw exceptions. You can achieve this by validating input arguments or columns beforehand using guards such as CASE WHEN, IS NOT NULL, RLIKE, or type try casts. | `expression`: sql expression to check on a DataFrame (fail the check if expression evaluates to True, pass if it evaluates to False); `msg`: optional message to output; `name`: optional name of the resulting column (it can be overwritten by `name` specified at the check level); `negate`: if the condition should be negated; `columns`: optional list of columns to be used for reporting and as name prefix if name not provided, unused in the actual logic | -| PII detection | Checks whether the values in the input column contain Personally Identifiable Information (PII). This check is not included in DQX built-in rules to avoid introducing 3rd-party dependencies. An example implementation can be found [here](/docs/reference/quality_rules#detecting-personally-identifiable-information-pii-using-a-python-function). | | -| `is_data_fresh` | Checks whether the values in the input timestamp column are not older than the specified number of minutes from the base timestamp column. This is useful for identifying stale data due to delayed pipelines and helps catch upstream issues early. | `column`: column of type timestamp/date to check (can be a string column name or a column expression); `max_age_minutes`: maximum age in minutes before data is considered stale; `base_timestamp`: optional base timestamp column from which the stale check is calculated. This can be a string, column expression, datetime value or literal value ex:F.lit(datetime(2024,1,1)). If not provided current_timestamp() function is used | +| Check | Description | Arguments | +| ---------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `is_not_null` | Checks whether the values in the input column are not null. | `column`: column to check (can be a string column name or a column expression) | +| `is_not_empty` | Checks whether the values in the input column are not empty (but may be null). | `column`: column to check (can be a string column name or a column expression) | +| `is_not_null_and_not_empty` | Checks whether the values in the input column are not null and not empty. | `column`: column to check (can be a string column name or a column expression); `trim_strings`: optional boolean flag to trim spaces from strings | +| `is_in_list` | Checks whether the values in the input column are present in the list of allowed values (null values are allowed). This check is not suited for large lists of allowed values. In such cases, it’s recommended to use the `foreign_key` dataset-level check instead. | `column`: column to check (can be a string column name or a column expression); `allowed`: list of allowed values | +| `is_not_null_and_is_in_list` | Checks whether the values in the input column are not null and present in the list of allowed values. This check is not suited for large lists of allowed values. In such cases, it’s recommended to use the `foreign_key` dataset-level check instead. | `column`: column to check (can be a string column name or a column expression); `allowed`: list of allowed values | +| `is_not_null_and_not_empty_array` | Checks whether the values in the array input column are not null and not empty. | `column`: column to check (can be a string column name or a column expression) | +| `is_in_range` | Checks whether the values in the input column are in the provided range (inclusive of both boundaries). | `column`: column to check (can be a string column name or a column expression); `min_limit`: min limit as number, date, timestamp, column name or sql expression; `max_limit`: max limit as number, date, timestamp, column name or sql expression | +| `is_not_in_range` | Checks whether the values in the input column are outside the provided range (inclusive of both boundaries). | `column`: column to check (can be a string column name or a column expression); `min_limit`: min limit as number, date, timestamp, column name or sql expression; `max_limit`: max limit as number, date, timestamp, column name or sql expression | +| `is_not_less_than` | Checks whether the values in the input column are not less than the provided limit. | `column`: column to check (can be a string column name or a column expression); `limit`: limit as number, date, timestamp, column name or sql expression | +| `is_not_greater_than` | Checks whether the values in the input column are not greater than the provided limit. | `column`: column to check (can be a string column name or a column expression); `limit`: limit as number, date, timestamp, column name or sql expression | +| `is_valid_date` | Checks whether the values in the input column have valid date formats. | `column`: column to check (can be a string column name or a column expression); `date_format`: optional date format (e.g. 'yyyy-mm-dd') | +| `is_valid_timestamp` | Checks whether the values in the input column have valid timestamp formats. | `column`: column to check (can be a string column name or a column expression); `timestamp_format`: optional timestamp format (e.g. 'yyyy-mm-dd HH:mm:ss') | +| `is_not_in_future` | Checks whether the values in the input column contain a timestamp that is not in the future, where 'future' is defined as current_timestamp + offset (in seconds). | `column`: column to check (can be a string column name or a column expression); `offset`: offset to use; `curr_timestamp`: current timestamp, if not provided current_timestamp() function is used | +| `is_not_in_near_future` | Checks whether the values in the input column contain a timestamp that is not in the near future, where 'near future' is defined as greater than the current timestamp but less than the current_timestamp + offset (in seconds). | `column`: column to check (can be a string column name or a column expression); `offset`: offset to use; `curr_timestamp`: current timestamp, if not provided current_timestamp() function is used | +| `is_older_than_n_days` | Checks whether the values in one input column are at least N days older than the values in another column. | `column`: column to check (can be a string column name or a column expression); `days`: number of days; `curr_date`: current date, if not provided current_date() function is used; `negate`: if the condition should be negated | +| `is_older_than_col2_for_n_days` | Checks whether the values in one input column are at least N days older than the values in another column. | `column1`: first column to check (can be a string column name or a column expression); `column2`: second column to check (can be a string column name or a column expression); `days`: number of days; `negate`: if the condition should be negated | +| `regex_match` | Checks whether the values in the input column match a given regex. | `column`: column to check (can be a string column name or a column expression); regex: regex to check; `negate`: if the condition should be negated (true) or not | +| `is_valid_ipv4_address` | Checks whether the values in the input column have valid IPv4 address format. | `column` to check (can be a string column name or a column expression) | +| `is_ipv4_address_in_cidr` | Checks whether the values in the input column have valid IPv4 address format and fall within the given CIDR block. | `column`: column to check (can be a string column name or a column expression); `cidr_block`: CIDR block string | +| `sql_expression` | Checks whether the values meet the condition provided as an SQL expression, e.g. `a = 'str1' and a > b`. SQL expressions are evaluated at runtime, so ensure that the expression is safe and that functions used within it (e.g. h3_ischildof, division) do not throw exceptions. You can achieve this by validating input arguments or columns beforehand using guards such as CASE WHEN, IS NOT NULL, RLIKE, or type try casts. | `expression`: sql expression to check on a DataFrame (fail the check if expression evaluates to True, pass if it evaluates to False); `msg`: optional message to output; `name`: optional name of the resulting column (it can be overwritten by `name` specified at the check level); `negate`: if the condition should be negated; `columns`: optional list of columns to be used for reporting and as name prefix if name not provided, unused in the actual logic | +| `is_data_fresh` | Checks whether the values in the input timestamp column are not older than the specified number of minutes from the base timestamp column. This is useful for identifying stale data due to delayed pipelines and helps catch upstream issues early. | `column`: column of type timestamp/date to check (can be a string column name or a column expression); `max_age_minutes`: maximum age in minutes before data is considered stale; `base_timestamp`: optional base timestamp column from which the stale check is calculated. This can be a string, column expression, datetime value or literal value ex:F.lit(datetime(2024,1,1)). If not provided current_timestamp() function is used | +| `does_not_contain_pii` | Checks whether the values in the input column contain Personally Identifiable Information (PII). Uses Microsoft Presidio to detect various named entities (e.g. PERSON, ADDRESS, EMAIL_ADDRESS). Requires installation of PII detection extras: `pip install 'databricks-labs-dqx[pii-detection]'`. See more details [here](#detecting-personally-identifiable-information-pii). | `column`: column to check (can be a string column name or a column expression); `threshold`: confidence threshold for PII detection (0.0 to 1.0, default: 0.7); `language`: optional language of the text (default: 'en'); `entities`: optional list of entities to detect; `nlp_engine_config`: optional dictionary configuring the NLP engine used for PII detection, see the [Presidio documentation](https://microsoft.github.io/presidio/analyzer/customizing_nlp_models/) for more information |
@@ -340,6 +340,32 @@ For brevity, the `name` field in the examples is omitted and it will be auto-gen - col2 - col3 +# does_not_contain_pii check +- criticality: error + check: + function: does_not_contain_pii + arguments: + column: col1 + threshold: 0.7 + language: en + +# does_not_contain_pii check with custom entities and NLP engine config +- criticality: error + check: + function: does_not_contain_pii + arguments: + column: col1 + threshold: 0.8 + entities: + - PERSON + - EMAIL_ADDRESS + - PHONE_NUMBER + nlp_engine_config: + nlp_engine_name: spacy + models: + - lang_code: en + model_name: en_core_web_md + # apply check to multiple columns - criticality: error check: @@ -381,6 +407,7 @@ For brevity, the `name` field in the examples is omitted and it will be auto-gen ```python from databricks.labs.dqx.rule import DQRowRule, DQForEachColRule from databricks.labs.dqx import check_funcs +from databricks.labs.dqx.pii import pii_detection_funcs from datetime import datetime checks = [ @@ -677,6 +704,32 @@ checks = [ } ), + # does_not_contain_pii check + DQRowRule( + criticality="error", + check_func=pii_detection_funcs.does_not_contain_pii, + column="col1", + check_func_kwargs={ + "threshold": 0.7, + "language": "en" + } + ), + + # does_not_contain_pii check with custom entities and NLP engine config + DQRowRule( + criticality="error", + check_func=pii_detection_funcs.does_not_contain_pii, + column="col1", + check_func_kwargs={ + "threshold": 0.8, + "entities": ["PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER"], + "nlp_engine_config": { + "nlp_engine_name": "spacy", + "models": [{"lang_code": "en", "model_name": "en_core_web_md"}] + } + } + ), + # column can be provided as attribute DQRowRule( criticality="error", @@ -2424,64 +2477,174 @@ Only non-complex column expressions (e.g., simple strings or direct col("x") ref This is because the metadata format relies on string representations that can't capture the full structure of Spark Column objects or Python functions. -## Detecting Personally Identifiable Information (PII) using a Python function +## Detecting Personally Identifiable Information (PII) + +DQX provides both built-in PII detection and support for custom PII detection implementations. PII detection checks can be installed and run as extras: + +```bash + pip install databricks-labs-dqx[pii] + ``` + + +PII detection using natural language models is non-deterministic. DQX cannot guarantee detection of all PII in an input dataset. While DQX makes a best-effort to detect PII in your data, a broader system should be designed to ensure data is cleansed of PII. + -You can write custom check to detect Personally Identifiable Information (PII) in the input data. To use PII-detection libraries (e.g. Presidio), -install any required dependencies, define a Spark UDF to detect PII, and use `make_condition` to call the UDF as a DQX check. See [DQX for PII Detection Notebook](https://github.com/databrickslabs/dqx/blob/v0.8.0/demos/dqx_demo_pii_detection.py) -for a full example. +### Using DQX's Built-in PII Detection + +The PII detection extras include a built-in `does_not_contain_pii` check that uses Microsoft Presidio to detect various types of personally identifiable information. + + + + ```yaml + # Basic PII detection check + - criticality: error + check: + function: does_not_contain_pii + arguments: + column: description + + # PII detection check with custom threshold and named entities + - criticality: error + check: + function: does_not_contain_pii + arguments: + column: description + threshold: 0.8 + entities: + - PERSON + - EMAIL_ADDRESS + ``` + + + ```python + from databricks.labs.dqx.rule import DQRowRule + from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii + + checks = [ + # Basic PII detection check + DQRowRule( + name="col_description_contains_any_pii", + criticality="error", + column="description", + check_func=does_not_contain_pii + ), + # PII detection check with custom threshold and named entities + DQRowRule( + name="col_description_contains_person_or_email_address", + criticality="error", + column="description", + check_func=does_not_contain_pii, + check_func_kwargs={"threshold": 0.8, "entities": ["PERSON", "EMAIL_ADDRESS"]} + ), + ] + ``` + + + +### Configuring Built-in PII Detection +The built-in check supports several configurable parameters: +- `threshold` tunes the sensitivity of the detection model (0.0 to 1.0, default: 0.7) +- `language` sets the language used by the model (default: 'en') +- `entities` sets the named entities (e.g. `PERSON`, `ADDRESS`) to be detected by the model; If none are provided, the model will detect all available named entities +- `nlp_engine_config` configures the NLP model used to detect named entities + +### NLP Engine Configuration +`does_not_contain_pii` uses NLP models to detect named entities in the input text. The choice of model, model parameters, and named entity mapping can be configured. Either built-in or custom configuration profiles can be used. + +#### Using Built-in NLP Engine Configurations + +DQX provides several built-in NLP engine configurations: + +- `SPACY_SMALL`: Uses the [en_core_web_sm](https://spacy.io/models/en#en_core_web_sm) model +- `SPACY_MEDIUM`: Uses the [en_core_web_md](https://spacy.io/models/en#en_core_web_md) model +- `SPACY_LARGE`: Uses the [en_core_web_lg](https://spacy.io/models/en#en_core_web_lg) model + +These can be loaded using `NLPEngineConfig`: - ```python - from pyspark.sql.functions import concat_ws, col, lit, pandas_udf - from pyspark.sql import Column - from databricks.labs.dqx.row_checks import make_condition - from presidio_analyzer import AnalyzerEngine - - # Create the Presidio analyzer: - analyzer = AnalyzerEngine() - - # Create a wrapper function to generate the entity mapping results: - def get_entity_mapping(data: str) -> str | None: - # Run the Presidio analyzer to detect PII in the string: - results = analyzer.analyze( - text=data, - entities=["PERSON", "EMAIL_ADDRESS"], - language='en', - score_threshold=0.5, - ) + ```python + from databricks.labs.dqx.rule import DQRowRule + from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii + from databricks.labs.dqx.pii.nlp_engine_config import NLPEngineConfig + + checks = [ + # PII detection check using spacy as a named entity recognizer + DQRowRule( + name="col_description_contains_pii", + criticality="error", + column="description", + check_func=does_not_contain_pii, + check_func_kwargs={"nlp_engine_config": NLPEngineConfig.SPACY_MEDIUM} + ), + ] + ``` + + - # Validate and return the results: - results = [ - result.to_dict() - for result in results.entity_mapping() - if result.score >= 0.5 - ] - if results != []: - return json.dumps(results) - return None - - # Register a pandas UDF to run the analyzer: - @pandas_udf('string') - def contains_pii(batch: pd.Series) -> pd.Series: - # Apply `get_entity_mapping` to each value: - return batch.map(get_entity_mapping) - - def does_not_contain_pii(column: str) -> Column: - # Define a PII detection expression calling the pandas UDF: - pii_info = contains_pii(col(column)) - - # Return the DQX condition that uses the PII detection expression: - return make_condition( - pii_info.isNotNull(), - concat_ws( - ' ', - lit(column), - lit('contains pii with the following info:'), - pii_info - ), - f'{column}_contains_pii' - ) - ``` +#### Using a Custom NLP Engine Configuration +Custom profiles can be provided to configure the PII detection model with advanced settings: + + +Using custom models for named-entity recognition may require you to install these models on your Databricks cluster using `pip` commands or as [compute-scoped libraries](https://docs.databricks.com/aws/en/libraries/cluster-libraries). + + + + + ```python + from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii + from databricks.labs.dqx.rule import DQRowRule + from databricks.labs.dqx.engine import DQEngine + from databricks.sdk import WorkspaceClient + + nlp_engine_config = { + 'nlp_engine_name': 'transformers_stanford_deidentifier_base', + 'models': [ + { + 'lang_code': 'en', + 'model_name': {'spacy': 'en_core_web_sm', 'transformers': 'StanfordAIMI/stanford-deidentifier-base'}, + } + ], + 'ner_model_configuration': { + 'labels_to_ignore': ['O'], + 'aggregation_strategy': 'max', + 'stride': 16, + 'alignment_mode': 'expand', + 'model_to_presidio_entity_mapping': { + 'PER': 'PERSON', + 'LOC': 'LOCATION', + 'ORG': 'ORGANIZATION', + 'AGE': 'AGE', + 'ID': 'ID', + 'EMAIL': 'EMAIL', + 'PATIENT': 'PERSON', + 'STAFF': 'PERSON', + 'HOSP': 'ORGANIZATION', + 'PATORG': 'ORGANIZATION', + 'DATE': 'DATE_TIME', + 'PHONE': 'PHONE_NUMBER', + 'HCW': 'PERSON', + 'HOSPITAL': 'LOCATION', + 'VENDOR': 'ORGANIZATION', + }, + 'low_confidence_score_multiplier': 0.4, + 'low_score_entity_names': ['ID'], + }, + } + checks = [ + # PII detection check using a custom named entity recognizer configuration + DQRowRule( + name="col_description_contains_pii", + criticality="error", + column="description", + check_func=does_not_contain_pii, + check_func_kwargs={"nlp_engine_config": nlp_engine_config}, + ), + ] + + dq_engine = DQEngine(WorkspaceClient()) + df = spark.read.table("main.default.table") + valid_df, quarantine_df = dq_engine.apply_checks_and_split(df, checks) + ``` diff --git a/pyproject.toml b/pyproject.toml index d4e5c5a51..9205f4959 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,14 @@ dependencies = ["databricks-labs-blueprint>=0.9.1,<0.10", [project.optional-dependencies] cli = ["pyspark~=3.5.0"] +pii = [ + "numpy<2.0,>=1.20", + "presidio-analyzer~=2.2.359", + "spacy~=3.8", + "en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl", + "en_core_web_md @ https://github.com/explosion/spacy-models/releases/download/en_core_web_md-3.8.0/en_core_web_md-3.8.0-py3-none-any.whl", + "en_core_web_lg @ https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-3.8.0/en_core_web_lg-3.8.0-py3-none-any.whl", +] [project.entry-points.databricks] runtime = "databricks.labs.dqx.runtime:main" @@ -54,6 +62,7 @@ include = ["src"] path = "src/databricks/labs/dqx/__about__.py" [tool.hatch.envs.default] +features = ["pii"] dependencies = [ "black~=24.3.0", "chispa~=0.10.1", @@ -62,6 +71,7 @@ dependencies = [ "databricks-labs-pytester~=0.7.2", "mypy~=1.9.0", "pylint~=3.3.1", + "pylint-per-file-ignores~=1.3", "pylint-pytest==2.0.0a0", "pytest~=8.3.3", "pytest-cov~=4.1.0", @@ -72,7 +82,7 @@ dependencies = [ "types-PyYAML~=6.0.12", "types-requests~=2.31.0", "databricks-connect~=15.4", - "pyspark~=3.5.0", + "pyspark~=3.5.0" ] python="3.10" @@ -100,7 +110,7 @@ extract_yaml_checks_examples = "python src/databricks/labs/dqx/llm/extract_yaml_ profile = "black" [tool.mypy] -exclude = ['venv', '.venv', 'demos/*'] +exclude = ['venv', '.venv', 'demos/*', 'tests/e2e/*'] [tool.pytest.ini_options] addopts = "--no-header" @@ -111,13 +121,13 @@ filterwarnings = ["ignore::DeprecationWarning"] target-version = ["py310"] line-length = 120 skip-string-normalization = true -extend-exclude = 'demos/' +extend-exclude = 'demos/|tests/e2e/notebooks' [tool.ruff] cache-dir = ".venv/ruff-cache" target-version = "py310" line-length = 120 -exclude = ["demos/*"] +exclude = ["demos/*", "tests/e2e/notebooks"] [tool.ruff.lint.isort] known-first-party = ["databricks.labs.dqx"] @@ -218,6 +228,7 @@ load-plugins = [ "pylint.extensions.redefined_variable_type", "pylint.extensions.set_membership", "pylint.extensions.typing", + "pylint_per_file_ignores", ] # Pickle collected data for later comparisons. @@ -743,3 +754,6 @@ ignored-argument-names = "_.*|^ignored_|^unused_" # List of qualified module names which can have objects that can redefine # builtins. redefining-builtins-modules = ["six.moves", "past.builtins", "future.builtins", "builtins", "io"] + +[tool.hatch.metadata] +allow-direct-references = true diff --git a/src/databricks/labs/dqx/llm/resources/yaml_checks_examples.yml b/src/databricks/labs/dqx/llm/resources/yaml_checks_examples.yml index c999b729c..02bd64ee0 100644 --- a/src/databricks/labs/dqx/llm/resources/yaml_checks_examples.yml +++ b/src/databricks/labs/dqx/llm/resources/yaml_checks_examples.yml @@ -225,6 +225,28 @@ columns: - col2 - col3 +- criticality: error + check: + function: does_not_contain_pii + arguments: + column: col1 + threshold: 0.7 + language: en +- criticality: error + check: + function: does_not_contain_pii + arguments: + column: col1 + threshold: 0.8 + entities: + - PERSON + - EMAIL_ADDRESS + - PHONE_NUMBER + nlp_engine_config: + nlp_engine_name: spacy + models: + - lang_code: en + model_name: en_core_web_md - criticality: error check: function: is_not_null @@ -545,6 +567,20 @@ arguments: columns: - col1 +- criticality: error + check: + function: does_not_contain_pii + arguments: + column: description +- criticality: error + check: + function: does_not_contain_pii + arguments: + column: description + threshold: 0.8 + entities: + - PERSON + - EMAIL_ADDRESS - criticality: warn check: function: is_not_null_and_not_empty diff --git a/src/databricks/labs/dqx/pii/__init__.py b/src/databricks/labs/dqx/pii/__init__.py new file mode 100644 index 000000000..9e076be60 --- /dev/null +++ b/src/databricks/labs/dqx/pii/__init__.py @@ -0,0 +1,16 @@ +from importlib.util import find_spec + +required_specs = [ + "presidio_analyzer", + "spacy", + "en_core_web_sm", + "en_core_web_md", + "en_core_web_lg", +] + +# Check that PII detection modules are installed +if not all(find_spec(spec) for spec in required_specs): + raise ImportError( + "PII detection extras not installed; Install additional " + "dependencies by running `pip install databricks-labs-dqx[pii]" + ) diff --git a/src/databricks/labs/dqx/pii/nlp_engine_config.py b/src/databricks/labs/dqx/pii/nlp_engine_config.py new file mode 100644 index 000000000..650051233 --- /dev/null +++ b/src/databricks/labs/dqx/pii/nlp_engine_config.py @@ -0,0 +1,20 @@ +from enum import Enum + + +class NLPEngineConfig(Enum): + """Enum class defining various NLP engine configurations for PII detection.""" + + SPACY_SMALL = { + "nlp_engine_name": "spacy", + "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}], + } + + SPACY_MEDIUM = { + "nlp_engine_name": "spacy", + "models": [{"lang_code": "en", "model_name": "en_core_web_md"}], + } + + SPACY_LARGE = { + "nlp_engine_name": "spacy", + "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}], + } diff --git a/src/databricks/labs/dqx/pii/pii_detection_funcs.py b/src/databricks/labs/dqx/pii/pii_detection_funcs.py new file mode 100644 index 000000000..997fb4fac --- /dev/null +++ b/src/databricks/labs/dqx/pii/pii_detection_funcs.py @@ -0,0 +1,153 @@ +import logging +import json +import re +import warnings +from collections.abc import Callable +import pandas as pd # type: ignore[import-untyped] + +import pyspark.sql.connect.session +from presidio_analyzer import AnalyzerEngine +from presidio_analyzer.nlp_engine import NlpEngineProvider + +from pyspark.sql import Column +from pyspark.sql.functions import concat_ws, lit, pandas_udf + +from databricks.labs.dqx.rule import register_rule +from databricks.labs.dqx.check_funcs import make_condition, _get_normalized_column_and_expr +from databricks.labs.dqx.pii.nlp_engine_config import NLPEngineConfig + +logging.getLogger("presidio_analyzer").setLevel(logging.ERROR) +logger = logging.getLogger(__name__) + +_default_nlp_engine_config = NLPEngineConfig.SPACY_SMALL + + +@register_rule("row") +def does_not_contain_pii( + column: str | Column, + language: str = "en", + threshold: float = 0.7, + entities: list[str] | None = None, + nlp_engine_config: NLPEngineConfig | dict | None = None, +) -> Column: + """ + Check if a column contains personally-identifying information (PII). Uses Microsoft Presidio to detect various named + entities (e.g. PERSON, ADDRESS, EMAIL_ADDRESS). If PII is detected, the message includes a JSON string with the + entity types, location within the string, and confidence score from the model. + + :param column: Column to check; can be a string column name or a column expression + :param language: Optional language of the text (default: 'en') + :param threshold: Confidence threshold for PII detection (0.0 to 1.0, default: 0.7) + Higher values = less sensitive, fewer false positives + Lower values = more sensitive, more potential false positives + :param entities: Optional list of entities to detect + :param nlp_engine_config: Optional NLP engine configuration used for PII detection; Can be `dict` or `NLPEngineConfiguration` + :return: Column object for condition that fails when PII is detected + """ + warnings.warn( + "PII detection uses pandas user-defined functions which may degrade performance. " + "Sample or limit large datasets when running PII detection.", + ) + + if threshold < 0.0 or threshold > 1.0: + raise ValueError(f"Provided threshold {threshold} must be between 0.0 and 1.0") + + if not nlp_engine_config: + nlp_engine_config = _default_nlp_engine_config + + if not isinstance(nlp_engine_config, dict | NLPEngineConfig): + raise ValueError(f"Invalid type provided for 'nlp_engine_config': {type(nlp_engine_config)}") + + _validate_environment() + config_dict = nlp_engine_config if isinstance(nlp_engine_config, dict) else nlp_engine_config.value + entity_detection_udf = _build_detection_udf(config_dict, language, threshold, entities) + col_str_norm, _, col_expr = _get_normalized_column_and_expr(column) + entity_info = entity_detection_udf(col_expr) + condition = entity_info.isNotNull() + message = concat_ws(" ", lit(f"Column '{col_str_norm}' contains PII:"), entity_info) + + return make_condition(condition=condition, message=message, alias=f"{col_str_norm}_contains_pii") + + +def _validate_environment() -> None: + """ + Validates that the environment can run PII detection checks which use Python dependencies. + + As of Databricks Connect 17.1, strict limits are imposed on the size of dependencies for + user-defined functions. UDFs will fail with out-of-memory errors if these limits are exceeded. + + Because of this limitation, we limit use of PII detection checks to local Spark or a Databricks + workspace. Databricks Connect uses a `pyspark.sql.connect.session.SparkSession` with an external + host (e.g. 'https://hostname.cloud.databricks.com'). To raise a clear error message, we check + the session and intentionally fail if `does_not_contain_pii` is called using Databricks Connect. + """ + connect_session_pattern = re.compile(r"127.0.0.1|.*grpc.sock") + session = pyspark.sql.SparkSession.builder.getOrCreate() + if isinstance(session, pyspark.sql.connect.session.SparkSession) and not connect_session_pattern.search( + session.client.host + ): + raise ImportError("'does_not_contain_pii' is not supported when running checks with Databricks Connect") + + +def _build_detection_udf( + nlp_engine_config: dict, language: str, threshold: float, entities: list[str] | None +) -> Callable: + """ + Builds a UDF with the provided threshold, entities, language, and analyzer. + + :param nlp_engine_config: Dictionary configuring the NLP engine used for PII detection + :param language: Language of the text + :param threshold: Confidence threshold for named entity detection (0.0 to 1.0) + :param entities: List of entities to detect + :return: PySpark UDF which can be called to detect PII with the given configuration + """ + + @pandas_udf("string") # type: ignore[call-overload] + def handler(batch: pd.Series) -> pd.Series: + def _get_analyzer() -> AnalyzerEngine: + """ + Gets an `AnalyzerEngine` for use with PII detection checks. + + :return: Presidio `AnalyzerEngine` + """ + provider = NlpEngineProvider(nlp_configuration=nlp_engine_config) + nlp_engine = provider.create_engine() + return AnalyzerEngine(nlp_engine=nlp_engine) + + analyzer = _get_analyzer() + + def _detect_named_entities(text: str) -> str | None: + """ + Detects named entities in the input text using a Presidio analyzer. + + :param text: Input text to analyze for named entities + :return: JSON string with detected entities, or `None` if no entities are found + """ + if not text: + return None + + results = analyzer.analyze( + text=text, + entities=entities, + language=language, + score_threshold=threshold, + ) + + qualified_results = [result for result in results if result.score >= threshold] + if not qualified_results: + return None + + return json.dumps( + [ + { + "entity_type": result.entity_type, + "score": float(result.score), + "text": text[result.start : result.end], + } + for result in qualified_results + ] + ) + + return batch.map(_detect_named_entities) + + return handler diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py new file mode 100644 index 000000000..c2f5a51d9 --- /dev/null +++ b/tests/e2e/conftest.py @@ -0,0 +1,16 @@ +import logging +import os +import pytest + + +logging.getLogger("tests").setLevel("DEBUG") +logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") +logger = logging.getLogger(__name__) + + +@pytest.fixture +def library_ref() -> str: + test_library_ref = "git+https://github.com/databrickslabs/dqx" + if os.getenv("REF_NAME"): + test_library_ref = f"{test_library_ref}.git@refs/pull/{os.getenv('REF_NAME')}" + return test_library_ref diff --git a/tests/e2e/notebooks/pii_detection_notebook.py b/tests/e2e/notebooks/pii_detection_notebook.py new file mode 100644 index 000000000..38aff24ed --- /dev/null +++ b/tests/e2e/notebooks/pii_detection_notebook.py @@ -0,0 +1,219 @@ +# Databricks notebook source + +dbutils.widgets.text("test_library_ref", "", "Test Library Ref") +%pip install 'databricks-labs-dqx @ {dbutils.widgets.get("test_library_ref")}' pytest + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +import pytest +from importlib import import_module + +def test_import_pii_module_fails_without_installation(): + with pytest.raises(ImportError): + import_module('databricks.labs.dqx.pii') + +test_import_pii_module_fails_without_installation() + +# COMMAND ---------- + +%pip install 'databricks-labs-dqx[pii] @ {dbutils.widgets.get("test_library_ref")}' chispa==0.10.1 + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +import pyspark.sql.functions as F +from databricks.labs.dqx.pii.nlp_engine_config import NLPEngineConfig +from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii +from chispa import assert_df_equality + +# COMMAND ---------- +# DBTITLE 1,test_does_not_contain_pii_basic + +def test_does_not_contain_pii_basic(): + schema_pii = "col1: string, col2: string, col3: string" + test_df = spark.createDataFrame( + [ + ["Hello world", "John Doe", "john.doe@example.com"], + ["No sensitive data here", "Not a person", "not-an-email"], + ["", "", ""], + [None, None, None], + ], + schema_pii, + ) + + actual = test_df.select( + does_not_contain_pii("col1"), + does_not_contain_pii("col2"), + does_not_contain_pii("col3"), + ) + + checked_schema = "col1_contains_pii: string, col2_contains_pii: string, col3_contains_pii: string" + expected = spark.createDataFrame( + [ + [ + None, + """Column 'col2' contains PII: [{"entity_type": "PERSON", "score": 1.0, "text": "John Doe"}]""", + """Column 'col3' contains PII: [{"entity_type": "EMAIL_ADDRESS", "score": 1.0, "text": "john.doe@example.com"}]""", + ], + [ + None, + None, + None, + ], + [ + None, + None, + None, + ], + [ + None, + None, + None, + ], + ], + checked_schema, + ) + transforms = [ + lambda df: df.select( + F.ilike("col1_contains_pii", F.lit("Column 'col1' contains PII: %")).alias("col1_contains_pii"), + F.ilike("col2_contains_pii", F.lit("Column 'col2' contains PII: %")).alias("col2_contains_pii"), + F.ilike("col3_contains_pii", F.lit("Column 'col3' contains PII: %")).alias("col3_contains_pii"), + ) + ] + assert_df_equality(actual, expected, transforms=transforms) + +test_does_not_contain_pii_basic() + +# COMMAND ---------- +# DBTITLE 1,test_does_not_contain_pii_with_entities_list + +def test_does_not_contain_pii_with_entities_list(): + schema_pii = "col1: string, col2: string" + test_df = spark.createDataFrame( + [ + ["John Doe", "John Doe lives at 123 Main St and can be reached at john@example.com"], + ["test@email.com", "Just an email here"], + ["No PII content", "Nothing sensitive"], + [None, None], + ], + schema_pii, + ) + + actual = test_df.select( + does_not_contain_pii("col1", entities=["PERSON", "EMAIL_ADDRESS"]), + does_not_contain_pii("col2", entities=["PERSON"]), + ) + + checked_schema = "col1_contains_pii: string, col2_contains_pii: string" + expected = spark.createDataFrame( + [ + [ + """Column 'col1' contains PII: [{"entity_type": "PERSON", "score": 1.0, "text": "John Doe"}]""", + """Column 'col2' contains PII: [{"entity_type": "PERSON", "score": 1.0, "text": "John Doe"},{"entity_type": "EMAIL_ADDRESS", "score": 1.0, "text": "john@example.com"}]""", + ], + [ + """Column 'col1' contains PII: [{"entity_type": "EMAIL_ADDRESS", "score": 1.0, "text": "test@email.com"}]""", + None, + ], + [ + None, + None, + ], + [ + None, + None, + ], + ], + checked_schema, + ) + transforms = [ + lambda df: df.select( + F.ilike("col1_contains_pii", F.lit("Column 'col1' contains PII: %")).alias("col1_contains_pii"), + F.ilike("col2_contains_pii", F.lit("Column 'col2' contains PII: %")).alias("col2_contains_pii"), + ) + ] + assert_df_equality(actual, expected, transforms=transforms) + +test_does_not_contain_pii_with_entities_list() + +# COMMAND ---------- +# DBTITLE 1,test_does_not_contain_pii_with_builtin_nlp_engine_config + +def test_does_not_contain_pii_with_builtin_nlp_engine_config(): + schema_pii = "col1: string" + test_df = spark.createDataFrame( + [ + ["Dr. Jane Smith works at Memorial Hospital"], + ["Patient ID: 12345, DOB: 1990-01-01"], + ["Regular text without PII"], + [None], + ], + schema_pii, + ) + + actual = test_df.select(does_not_contain_pii("col1", entities=["PERSON", "DATE_TIME"], nlp_engine_config=NLPEngineConfig.SPACY_MEDIUM)) + + checked_schema = "col1_contains_pii: string" + expected = spark.createDataFrame( + [ + ["""Column 'col1' contains PII: [{"entity_type": "PERSON", "score": 1.0, "text": "Jane Smith"}]"""], + ["""Column 'col1' contains PII: [{"entity_type": "DATE_TIME", "score": 1.0, "text": "1990-01-01"}]"""], + [None], + [None], + ], + checked_schema, + ) + transforms = [ + lambda df: df.select( + F.ilike("col1_contains_pii", F.lit("Column 'col1' contains PII: %")).alias("col1_contains_pii"), + ) + ] + assert_df_equality(actual, expected, transforms=transforms) + +test_does_not_contain_pii_with_builtin_nlp_engine_config() + +# COMMAND ---------- +# DBTITLE 1,test_does_not_contain_pii_with_custom_nlp_config_dict + +def test_does_not_contain_pii_with_custom_nlp_config_dict(): + schema_pii = "col1: string" + test_df = spark.createDataFrame( + [ + ["Dr. Jane Smith treated patient John Doe at City Hospital"], + ["Lorem ipsum dolor sit amet"], + [None], + ], + schema_pii, + ) + custom_nlp_engine_config = { + "nlp_engine_name": "spacy", + "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}], + } + actual = test_df.select(does_not_contain_pii("col1", entities=["PERSON"], nlp_engine_config=custom_nlp_engine_config)) + + checked_schema = "col1_contains_pii: string" + expected = spark.createDataFrame( + [ + [ + """Column 'col1' contains PII: [{"entity_type": "PERSON", "score": 1.0, "text": "Jane Smith"},{"entity_type": "PERSON", "score": 1.0, "text": "John Doe"}]""" + ], + [None], + [None], + ], + checked_schema, + ) + transforms = [ + lambda df: df.select( + F.ilike("col1_contains_pii", F.lit("Column 'col1' contains PII: %")).alias("col1_contains_pii"), + ) + ] + assert_df_equality(actual, expected, transforms=transforms) + +test_does_not_contain_pii_with_custom_nlp_config_dict() diff --git a/tests/e2e/test_pii_detection_checks.py b/tests/e2e/test_pii_detection_checks.py new file mode 100644 index 000000000..bb02dce44 --- /dev/null +++ b/tests/e2e/test_pii_detection_checks.py @@ -0,0 +1,49 @@ +import logging + +from datetime import timedelta +from pathlib import Path +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.workspace import ImportFormat +from databricks.sdk.service.jobs import NotebookTask, Run, Task, TerminationTypeType + + +logger = logging.getLogger(__name__) +RETRY_INTERVAL_SECONDS = 30 + + +def test_run_pii_detection_notebook(make_notebook, make_job, library_ref): + ws = WorkspaceClient() + path = Path(__file__).parent / "notebooks" / "pii_detection_notebook.py" + with open(path, "rb") as f: + notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + + notebook_path = notebook.as_fuse().as_posix() + notebook_task = NotebookTask( + notebook_path=notebook_path, + base_parameters={"test_library_ref": library_ref}, + ) + job = make_job( + tasks=[Task(task_key="pii_detection_notebook", notebook_task=notebook_task)], + ) + + waiter = ws.jobs.run_now_and_wait(job.job_id) + run = ws.jobs.wait_get_run_job_terminated_or_skipped( + run_id=waiter.run_id, + timeout=timedelta(minutes=30), + callback=lambda r: validate_run_status(r, client=ws), + ) + logging.info(f"Job run {run.run_id} completed successfully for pii_detection_notebook") + + +def validate_run_status(run: Run, client: WorkspaceClient) -> None: + """ + Validates that a job task run completed successfully. + :param run: `Run` object returned from a `WorkspaceClient.jobs.submit(...)` command + :param client: `WorkspaceClient` object for getting task output + """ + task = run.tasks[0] + termination_details = run.status.termination_details + + assert ( + termination_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with message: {client.jobs.get_run_output(task.run_id).error}" diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index e6a86539a..2f3d809f4 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,5 +1,4 @@ import logging -import os from datetime import timedelta from pathlib import Path @@ -9,20 +8,14 @@ from databricks.sdk.service.pipelines import NotebookLibrary, PipelinesEnvironment, PipelineLibrary from databricks.sdk.service.jobs import NotebookTask, PipelineTask, Run, Task, TerminationTypeType -logging.getLogger("tests").setLevel("DEBUG") -logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") -logger = logging.getLogger(__name__) +logger = logging.getLogger(__name__) RETRY_INTERVAL_SECONDS = 30 -TEST_LIBRARY_REF = "git+https://github.com/databrickslabs/dqx" -if os.getenv("REF_NAME"): - TEST_LIBRARY_REF = f"{TEST_LIBRARY_REF}.git@refs/pull/{os.getenv('REF_NAME')}" -logger.info(f"Running demo tests from {TEST_LIBRARY_REF}") -def test_run_dqx_demo_library(make_notebook, make_schema, make_job): - path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py" +def test_run_dqx_demo_library(make_notebook, make_schema, make_job, library_ref): ws = WorkspaceClient() + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py" with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) directory = notebook.as_fuse().parent.as_posix() @@ -36,7 +29,7 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): "demo_database_name": catalog, "demo_schema_name": schema, "demo_file_directory": directory, - "test_library_ref": TEST_LIBRARY_REF, + "test_library_ref": library_ref, }, ) job = make_job(tasks=[Task(task_key="dqx_demo_library", notebook_task=notebook_task)]) @@ -45,14 +38,14 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_library") -def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, make_job): - path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py" +def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, make_job, library_ref): ws = WorkspaceClient() + path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py" with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) folder = notebook.as_fuse().parent / "quality_rules" @@ -63,7 +56,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, - base_parameters={"demo_database": catalog, "demo_schema": schema, "test_library_ref": TEST_LIBRARY_REF}, + base_parameters={"demo_database": catalog, "demo_schema": schema, "test_library_ref": library_ref}, ) job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) @@ -71,64 +64,62 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo") -def test_run_dqx_quick_start_demo_library(make_notebook, make_job): - path = Path(__file__).parent.parent.parent / "demos" / "dqx_quick_start_demo_library.py" +def test_run_dqx_quick_start_demo_library(make_notebook, make_job, library_ref): ws = WorkspaceClient() + path = Path(__file__).parent.parent.parent / "demos" / "dqx_quick_start_demo_library.py" with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() - notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters={"test_library_ref": TEST_LIBRARY_REF}) + notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters={"test_library_ref": library_ref}) job = make_job(tasks=[Task(task_key="dqx_quick_start_demo_library", notebook_task=notebook_task)]) waiter = ws.jobs.run_now_and_wait(job.job_id) run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_quick_start_demo_library") -def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): - path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_pii_detection.py" +def test_run_dqx_demo_pii_detection(make_notebook, make_job, library_ref): ws = WorkspaceClient() + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_pii_detection.py" with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() - notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters={"test_library_ref": TEST_LIBRARY_REF}) - cluster = make_cluster(single_node=True, spark_version="15.4.x-scala2.12") - job = make_job( - tasks=[ - Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task, existing_cluster_id=cluster.cluster_id) - ] + notebook_task = NotebookTask( + notebook_path=notebook_path, + base_parameters={"test_library_ref": library_ref}, ) + job = make_job(tasks=[Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task)]) waiter = ws.jobs.run_now_and_wait(job.job_id) run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_pii_detection") -def test_run_dqx_dlt_demo(make_notebook, make_pipeline, make_job): - path = Path(__file__).parent.parent.parent / "demos" / "dqx_dlt_demo.py" +def test_run_dqx_dlt_demo(make_notebook, make_pipeline, make_job, library_ref): ws = WorkspaceClient() + path = Path(__file__).parent.parent.parent / "demos" / "dqx_dlt_demo.py" with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() pipeline = make_pipeline( libraries=[PipelineLibrary(notebook=NotebookLibrary(notebook_path))], - environment=PipelinesEnvironment(dependencies=[TEST_LIBRARY_REF]), + environment=PipelinesEnvironment(dependencies=[library_ref]), ) pipeline_task = PipelineTask(pipeline_id=pipeline.pipeline_id) job = make_job(tasks=[Task(task_key="dqx_dlt_demo", pipeline_task=pipeline_task)]) @@ -137,7 +128,7 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline, make_job): run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_dlt_demo") @@ -175,26 +166,12 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_tool") -def validate_demo_run_status(run: Run, client: WorkspaceClient) -> None: - """ - Validates that a demo run completed successfully. - :param run: `Run` object returned from a `WorkspaceClient.jobs.submit(...)` command - :param client: `WorkspaceClient` object for getting task output - """ - task = run.tasks[0] - termination_details = run.status.termination_details - - assert ( - termination_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with message: {client.jobs.get_run_output(task.run_id).error}" - - -def test_run_dqx_streaming_demo_native(make_notebook, make_schema, make_job, tmp_path): +def test_run_dqx_streaming_demo_native(make_notebook, make_schema, make_job, tmp_path, library_ref): ws = WorkspaceClient() path = Path(__file__).parent.parent.parent / "demos" / "dqx_streaming_demo_native.py" @@ -212,7 +189,7 @@ def test_run_dqx_streaming_demo_native(make_notebook, make_schema, make_job, tmp "demo_schema_name": schema, "silver_checkpoint": f"{base_output_path}/silver_checkpoint", "quarantine_checkpoint": f"{base_output_path}/quarantine_checkpoint", - "test_library_ref": TEST_LIBRARY_REF, + "test_library_ref": library_ref, } notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters=base_parameters) job = make_job(tasks=[Task(task_key="dqx_streaming_demo", notebook_task=notebook_task)]) @@ -220,12 +197,12 @@ def test_run_dqx_streaming_demo_native(make_notebook, make_schema, make_job, tmp run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, client=ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_streaming_demo") -def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path): +def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path, library_ref): ws = WorkspaceClient() path = Path(__file__).parent.parent.parent / "demos" / "dqx_streaming_demo_diy.py" @@ -241,7 +218,7 @@ def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path): "silver_table": f"{base_output_path}/silver_table", "quarantine_checkpoint": f"{base_output_path}/quarantine_checkpoint", "quarantine_table": f"{base_output_path}/quarantine_table", - "test_library_ref": TEST_LIBRARY_REF, + "test_library_ref": library_ref, } notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters=base_parameters) job = make_job(tasks=[Task(task_key="dqx_streaming_demo", notebook_task=notebook_task)]) @@ -249,6 +226,20 @@ def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path): run = ws.jobs.wait_get_run_job_terminated_or_skipped( run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=lambda r: validate_demo_run_status(r, client=ws), + callback=lambda r: validate_run_status(r, client=ws), ) logging.info(f"Job run {run.run_id} completed successfully for dqx_streaming_demo") + + +def validate_run_status(run: Run, client: WorkspaceClient) -> None: + """ + Validates that a job task run completed successfully. + :param run: `Run` object returned from a `WorkspaceClient.jobs.submit(...)` command + :param client: `WorkspaceClient` object for getting task output + """ + task = run.tasks[0] + termination_details = run.status.termination_details + + assert ( + termination_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with message: {client.jobs.get_run_output(task.run_id).error}" diff --git a/tests/integration/test_row_checks.py b/tests/integration/test_row_checks.py index 789b1d26f..c3879d10d 100644 --- a/tests/integration/test_row_checks.py +++ b/tests/integration/test_row_checks.py @@ -1,5 +1,6 @@ from datetime import datetime from decimal import Decimal +import pytest import pyspark.sql.functions as F from chispa.dataframe_comparer import assert_df_equality # type: ignore from databricks.labs.dqx.check_funcs import ( @@ -25,6 +26,7 @@ is_ipv4_address_in_cidr, is_data_fresh, ) +from databricks.labs.dqx.pii import pii_detection_funcs SCHEMA = "a: string, b: int" @@ -1224,6 +1226,25 @@ def test_is_ipv4_address_in_cidr(spark): assert_df_equality(actual, expected, ignore_nullable=True) +def test_contains_pii_fails_session_validation(spark): + """ + We restrict running `does_not_contain_pii` using Databricks Connect due to limitations + on the size of UDF dependencies. Because tests are run from a Databricks Connect session, + we can only validate that the correct error is raised. + + Complete testing of `does_not_contain_pii` and its options has been added to e2e tests. + We run many scenarios in `test_pii_detection_checks` to validate `contains_pii` from + a Databricks workspace. + """ + schema_pii = "col1: string" + test_df = spark.createDataFrame([["Contact us at info@company.com"]], schema_pii) + + with pytest.raises( + ImportError, match="'does_not_contain_pii' is not supported when running checks with Databricks Connect" + ): + test_df.select(pii_detection_funcs.does_not_contain_pii("col1")) + + def test_is_data_fresh(spark, set_utc_timezone): input_schema = "a: string, b: timestamp, c: date, d: timestamp" test_df = spark.createDataFrame( diff --git a/tests/unit/test_row_checks.py b/tests/unit/test_row_checks.py index 29e46dfad..e1c558c80 100644 --- a/tests/unit/test_row_checks.py +++ b/tests/unit/test_row_checks.py @@ -9,6 +9,7 @@ is_aggr_not_greater_than, is_ipv4_address_in_cidr, ) +from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii LIMIT_VALUE_ERROR = "Limit is not provided" @@ -58,3 +59,15 @@ def test_col_is_ipv4_address_in_cidr_missing_cidr_block(): def test_col_is_ipv4_address_in_cidr_empty_cidr_block(): with pytest.raises(ValueError, match="'cidr_block' must be a non-empty string"): is_ipv4_address_in_cidr("a", cidr_block="") + + +def test_col_does_not_contain_pii_invalid_engine_config(): + nlp_engine_config = "'model': 'my_model'" + with pytest.raises(ValueError, match=f"Invalid type provided for 'nlp_engine_config': {type(nlp_engine_config)}"): + does_not_contain_pii("a", nlp_engine_config=nlp_engine_config) + + +@pytest.mark.parametrize("threshold", [-10.0, -0.1, 1.1, 10.0]) +def test_col_does_not_contain_pii_invalid_threshold(threshold: float): + with pytest.raises(ValueError, match=f"Provided threshold {threshold} must be between 0.0 and 1.0"): + does_not_contain_pii("a", threshold=threshold)