Adding detailed telemetry

demos/dqx_demo_pii_detection.py

@@ -3,7 +3,7 @@
 # MAGIC # Using DQX for PII Detection
 # 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 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.
+# MAGIC DQX provides in-transit 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 ----------
 

docs/dqx/docs/guide/index.mdx

@@ -21,7 +21,7 @@ For more details, see the [Installation Guide](/docs/installation/).
 
 | Quality checking type | Integration with processing pipelines | Description                                                                                                                                                                                                    |
 |---------------------- | ------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| **In-flight**         | Code-level only                       | DQX allows data quality to be validated on the fly while the data is being processed, before it is written to storage. This requires DQX to be used as a library and integrated directly into user pipelines.  |
+| **In-transit**        | Code-level only                       | DQX allows data quality to be validated on the fly while the data is being processed, before it is written to storage. This requires DQX to be used as a library and integrated directly into user pipelines.  |
 | **At-rest**           | Code-level or No-code (Workflows)     | DQX enables data quality checking on existing data stored in tables. For no-code integration, DQX must first be installed in the workspace as a tool to deploy workflows.                                      |
 
 **Integration options:**

docs/dqx/docs/reference/engine.mdx

@@ -41,7 +41,7 @@ spark = DatabricksSession.builder.getOrCreate()
 dq_engine = DQEngine(ws, spark)
 ```
 
-For local execution without a Databricks workspace, please refer to the [local testing section](/docs/reference/testing/#local-testing-with-dqengine).
+For local execution without a Databricks workspace, please refer to the [local testing section](/docs/reference/testing/#local-execution-and-testing-with-dqengine).
 
 ## DQX engine methods
 
@@ -60,11 +60,11 @@ The following table outlines the available methods of the `DQEngine` and their f
 | `validate_checks`                              | Validates the provided quality checks to ensure they conform to the expected structure and types.                                                                                                                                 | `checks`: List of checks to validate; `custom_check_functions`: (optional) dictionary of custom check functions that can be used; `validate_custom_check_functions`: (optional) if set to True, validates custom check functions (defaults to True).                                                                                                                                                                                                                                                                                                                                                                                              | Yes                                       |
 | `get_invalid`                                  | Retrieves records from the DataFrame that violate data quality checks (records with warnings and errors).                                                                                                                         | `df`: Input DataFrame.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            | Yes                                       |
 | `get_valid`                                    | Retrieves records from the DataFrame that pass all data quality checks.                                                                                                                                                           | `df`: Input DataFrame.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            | Yes                                       |
-| `load_checks`                                  | Loads quality rules (checks) from storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed sources where the location is inferred automatically from run config  | `config`: Configuration for loading checks from a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON) using absolute paths; `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: installation-managed storage backend, using the `checks_location` field from the run configuration. See more details below.                                                     | Yes (only with `FileChecksStorageConfig`) |
-| `save_checks`                                  | Saves quality rules (checks) to storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed targets where the location is inferred automatically from run config    | `checks`: List of checks defined as dictionary; `config`: Configuration for saving checks in a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON); `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: storage defined in the installation context, using the `checks_location` field from the run configuration. See more details below.                      | Yes (only with `FileChecksStorageConfig`) |
+| `load_checks`                                  | Loads quality rules (checks) from storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed sources where the location is inferred automatically from run config. | `config`: Configuration for loading checks from a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON) using absolute paths; `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: installation-managed storage backend, using the `checks_location` field from the run configuration. See more details below.                                                     | Yes (only with `FileChecksStorageConfig`) |
+| `save_checks`                                  | Saves quality rules (checks) to storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed targets where the location is inferred automatically from run config.   | `checks`: List of checks defined as dictionary; `config`: Configuration for saving checks in a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON); `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: storage defined in the installation context, using the `checks_location` field from the run configuration. See more details below.                      | Yes (only with `FileChecksStorageConfig`) |
 | `save_results_in_table`                        | Save quality checking results in delta table(s).                                                                                                                                                                                  | `output_df`: (optional) Dataframe containing the output data; `quarantine_df`: (optional) Dataframe containing the output data; `output_config`: `OutputConfig` object with the table name, output mode, and options for the output data; `quarantine_config`: `OutputConfig` object with the table name, output mode, and options for the quarantine data - if provided, data will be split; `run_config_name`: Name of the run config to use; `assume_user`: If True, assume user installation.                                                                                                                                                 | No                                        |
 
-The 'Supports local execution' in the above table indicates which methods can be used for local testing without a Databricks workspace (see the usage in [local testing section](/docs/reference/testing/#local-testing-with-dqengine)).
+The 'Supports local execution' in the above table indicates which methods can be used for local testing without a Databricks workspace (see the usage in [local testing section](/docs/reference/testing/#local-execution-and-testing-with-dqengine)).
 
 `InputConfig` support the following parameters:
 * `location`: The location of the input data source (e.g. table name or file path).

docs/dqx/docs/reference/quality_checks.mdx

@@ -2650,8 +2650,8 @@ This is because the metadata format relies on string representations that can't
 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]
-   ```
+pip install databricks-labs-dqx[pii]
+```
 
 <Admonition type="info" title="PII Detection Guarantees">
 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.

docs/dqx/docs/reference/testing.mdx

@@ -114,14 +114,21 @@ To run the integration tests on serverless compute, add the `DATABRICKS_SERVERLE
 ```
 When `DATABRICKS_SERVERLESS_COMPUTE_ID` is set, the `DATABRICKS_CLUSTER_ID` is ignored, and tests run on serverless compute.
 
-### Local testing with DQEngine
+### Local execution and testing with DQEngine
 
-If workspace-level access is unavailable in your testing environment, you can perform local testing by installing the latest `pyspark` package and mocking the workspace client.
-Below is an example test.
+If workspace-level access is unavailable in your environment, you can execute DQX locally without workspace connection by installing the latest `pyspark` package and mocking the workspace client:
+```
+# user latest dqx version
+pip install databricks-labs-dqx pyspark==3.5.0
+
+# or use local build
+hatch build
+pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark==3.5.0
+```
 
 <Admonition type="tip" title="Usage tips">
 **This approach should be treated as experimental!** It does not offer the same level of testing as the standard approach, and it is only applicable to selected methods (see [here](/docs/reference/engine/#dqx-engine-methods)).
-We strongly recommend following the standard testing procedure outlined above, which includes proper initialization of the workspace client.
+We strongly recommend following the standard execution procedure outlined above, which includes proper initialization of the workspace client.
 </Admonition>
 
 <Tabs>
@@ -130,33 +137,25 @@ We strongly recommend following the standard testing procedure outlined above, w
     from unittest.mock import MagicMock
     from databricks.sdk import WorkspaceClient
     from pyspark.sql import SparkSession
-    from chispa.dataframe_comparer import assert_df_equality
-    from databricks.labs.dqx.check_funcs import is_not_null_and_not_empty
+    from databricks.labs.dqx.check_funcs import is_not_null
     from databricks.labs.dqx.engine import DQEngine
     from databricks.labs.dqx.rule import DQRowRule
-    from databricks.labs.dqx.schema import dq_result_schema
-
 
-    def test_dq():
-      spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
-      ws = MagicMock(spec=WorkspaceClient, **{"current_user.me.return_value": None}) # mock the workspace client
+    spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
+    ws = MagicMock(spec=WorkspaceClient) # mock the workspace client
 
-      schema = "a: int, b: int, c: int"
-      expected_schema = schema + f", _errors: {dq_result_schema.simpleString()}, _warnings: {dq_result_schema.simpleString()}"
-      test_df = spark.createDataFrame([[1, None, 3]], schema)
+    schema = "a: int, b: int, c: int"
+    test_df = spark.createDataFrame([[1, None, 3]], schema)
 
-      checks = [
-        DQRowRule(name="col_a_is_null_or_empty", criticality="warn", check_func=is_not_null_and_not_empty, column="a"),
-        DQRowRule(name="col_b_is_null_or_empty", criticality="error", check_func=is_not_null_and_not_empty, column="b"),
-      ]
+    checks = [
+      DQRowRule(name="col_a_is_null", criticality="warn", check_func=is_not_null, column="a"),
+      DQRowRule(name="col_b_is_null", criticality="error", check_func=is_not_null, column="b"),
+    ]
 
-      dq_engine = DQEngine(ws)
-      df = dq_engine.apply_checks(test_df, checks)
+    dq_engine = DQEngine(spark=spark, workspace_client=ws)
+    df = dq_engine.apply_checks(test_df, checks)
 
-      expected_df = spark.createDataFrame(
-        [[1, None, 3, {"b_is_null_or_empty": "Column b is null or empty"}, None]], expected_schema
-      )
-      assert_df_equality(df, expected_df)
+    df.show(truncate=False)
     ```
   </TabItem>
 </Tabs>

pyproject.toml

@@ -172,7 +172,7 @@ extend-exclude = 'demos/|tests/e2e/notebooks'
 cache-dir = ".venv/ruff-cache"
 target-version = "py310"
 line-length = 120
-exclude = ["demos/*", "tests/e2e/notebooks"]
+exclude = ["demos/*", "tests/e2e/notebooks", "ui/*"]
 
 [tool.ruff.lint.isort]
 known-first-party = ["databricks.labs.dqx"]

src/databricks/labs/dqx/base.py

@@ -2,7 +2,7 @@
 from collections.abc import Callable
 from functools import cached_property
 from typing import final
-from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import DataFrame
 
 from databricks.labs.dqx.checks_validator import ChecksValidationStatus
 from databricks.labs.dqx.rule import DQRule
@@ -12,8 +12,7 @@
 
 class DQEngineBase(abc.ABC):
     def __init__(self, workspace_client: WorkspaceClient):
-        self._workspace_client = workspace_client
-        self._spark = SparkSession.builder.getOrCreate()
+        self._workspace_client = self._verify_workspace_client(workspace_client)
 
     @cached_property
     def ws(self) -> WorkspaceClient:
@@ -22,16 +21,7 @@ def ws(self) -> WorkspaceClient:
         Ensures workspace connectivity and sets the product info used for
         telemetry so that requests are attributed to *dqx*.
         """
-        return self._verify_workspace_client(self._workspace_client)
-
-    @cached_property
-    def spark(self) -> SparkSession:
-        """Return the *SparkSession* associated with this engine.
-
-        The session is created during initialization using
-        *SparkSession.builder.getOrCreate()*.
-        """
-        return self._spark
+        return self._workspace_client
 
     @staticmethod
     @final
@@ -45,7 +35,8 @@ def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
             setattr(ws.config, '_product_info', ('dqx', __version__))
 
         # make sure Databricks workspace is accessible
-        ws.current_user.me()
+        # use api that works on all workspaces and clusters including group assigned clusters
+        ws.clusters.select_spark_version()
         return ws