Skip to content

Commit 15ae1ec

Browse files
mwojtyczkamwojtyczka
authored
Merge branch 'main' into summary_metrics
2 parents 49303d6 + 8c0bc01 commit 15ae1ec

17 files changed

Lines changed: 185 additions & 57 deletions

File tree

‎demos/dqx_demo_pii_detection.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
# MAGIC # Using DQX for PII Detection
44
# 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.
55
# MAGIC
6-
# 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.
6+
# 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.
77

88
# COMMAND ----------
99

‎docs/dqx/docs/guide/index.mdx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ For more details, see the [Installation Guide](/docs/installation/).
2121

2222
| Quality checking type | Integration with processing pipelines | Description |
2323
|---------------------- | ------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
24-
| **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. |
24+
| **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. |
2525
| **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. |
2626

2727
**Integration options:**

‎docs/dqx/docs/reference/engine.mdx‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ spark = DatabricksSession.builder.getOrCreate()
4141
dq_engine = DQEngine(ws, spark)
4242
```
4343

44-
For local execution without a Databricks workspace, please refer to the [local testing section](/docs/reference/testing/#local-testing-with-dqengine).
44+
For local execution without a Databricks workspace, please refer to the [local testing section](/docs/reference/testing/#local-execution-and-testing-with-dqengine).
4545

4646
## DQX engine methods
4747

@@ -60,11 +60,11 @@ The following table outlines the available methods of the `DQEngine` and their f
6060
| `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 |
6161
| `get_invalid` | Retrieves records from the DataFrame that violate data quality checks (records with warnings and errors). | `df`: Input DataFrame. | Yes |
6262
| `get_valid` | Retrieves records from the DataFrame that pass all data quality checks. | `df`: Input DataFrame. | Yes |
63-
| `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`) |
64-
| `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`) |
63+
| `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`) |
64+
| `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`) |
6565
| `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 |
6666

67-
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)).
67+
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)).
6868

6969
`InputConfig` support the following parameters:
7070
* `location`: The location of the input data source (e.g. table name or file path).

‎docs/dqx/docs/reference/quality_checks.mdx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2650,8 +2650,8 @@ This is because the metadata format relies on string representations that can't
26502650
DQX provides both built-in PII detection and support for custom PII detection implementations. PII detection checks can be installed and run as extras:
26512651

26522652
```bash
2653-
pip install databricks-labs-dqx[pii]
2654-
```
2653+
pip install databricks-labs-dqx[pii]
2654+
```
26552655

26562656
<Admonition type="info" title="PII Detection Guarantees">
26572657
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‎

Lines changed: 27 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -114,14 +114,25 @@ To run the integration tests on serverless compute, add the `DATABRICKS_SERVERLE
114114
```
115115
When `DATABRICKS_SERVERLESS_COMPUTE_ID` is set, the `DATABRICKS_CLUSTER_ID` is ignored, and tests run on serverless compute.
116116

117-
### Local testing with DQEngine
117+
### Local execution and testing with DQEngine
118118

119-
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.
120-
Below is an example test.
119+
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:
120+
```
121+
# user latest dqx version
122+
pip install databricks-labs-dqx pyspark==3.5.0
123+
124+
# or use local build
125+
hatch build
126+
pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark==3.5.6
127+
```
128+
129+
<Admonition type="warning" title="Python requirements">
130+
You need a python version that is compatible with the pyspark version, e.g. Python 3.9, 3.10, 3.11, and 3.12 are compatible with PySpark 3.5, but Python 3.13 is not supported.
131+
</Admonition>
121132

122133
<Admonition type="tip" title="Usage tips">
123134
**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)).
124-
We strongly recommend following the standard testing procedure outlined above, which includes proper initialization of the workspace client.
135+
We strongly recommend following the standard execution procedure outlined above, which includes proper initialization of the workspace client.
125136
</Admonition>
126137

127138
<Tabs>
@@ -130,35 +141,25 @@ We strongly recommend following the standard testing procedure outlined above, w
130141
from unittest.mock import MagicMock
131142
from databricks.sdk import WorkspaceClient
132143
from pyspark.sql import SparkSession
133-
from chispa.dataframe_comparer import assert_df_equality
134-
from databricks.labs.dqx.check_funcs import is_not_null_and_not_empty
144+
from databricks.labs.dqx.check_funcs import is_not_null
135145
from databricks.labs.dqx.engine import DQEngine
136146
from databricks.labs.dqx.rule import DQRowRule
137-
from databricks.labs.dqx.schema import dq_result_schema
138147

148+
spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
149+
ws = MagicMock(spec=WorkspaceClient) # mock the workspace client
139150

140-
def test_dq():
141-
spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
142-
ws = MagicMock(spec=WorkspaceClient, **{"current_user.me.return_value": None}) # mock the workspace client
151+
schema = "a: int, b: int, c: int"
152+
test_df = spark.createDataFrame([[1, None, 3]], schema)
143153

144-
schema = "a: int, b: int, c: int"
145-
expected_schema = schema + f", _errors: {dq_result_schema.simpleString()}, _warnings: {dq_result_schema.simpleString()}"
146-
test_df = spark.createDataFrame([[1, None, 3]], schema)
154+
checks = [
155+
DQRowRule(name="col_a_is_null", criticality="warn", check_func=is_not_null, column="a"),
156+
DQRowRule(name="col_b_is_null", criticality="error", check_func=is_not_null, column="b"),
157+
]
147158

148-
checks = [
149-
DQRowRule(name="col_a_is_null_or_empty", criticality="warn", check_func=is_not_null_and_not_empty, column="a"),
150-
DQRowRule(name="col_b_is_null_or_empty", criticality="error", check_func=is_not_null_and_not_empty, column="b"),
151-
]
159+
dq_engine = DQEngine(spark=spark, workspace_client=ws)
160+
df = dq_engine.apply_checks(test_df, checks)
152161

153-
dq_engine = DQEngine(ws)
154-
df = dq_engine.apply_checks(test_df, checks)
155-
156-
expected_df = spark.createDataFrame(
157-
[[1, None, 3, {"b_is_null_or_empty": "Column b is null or empty"}, None]], expected_schema
158-
)
159-
assert_df_equality(df, expected_df)
162+
df.show(truncate=False)
160163
```
161164
</TabItem>
162165
</Tabs>
163-
164-

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,7 @@ extend-exclude = 'demos/|tests/e2e/notebooks'
172172
cache-dir = ".venv/ruff-cache"
173173
target-version = "py310"
174174
line-length = 120
175-
exclude = ["demos/*", "tests/e2e/notebooks"]
175+
exclude = ["demos/*", "tests/e2e/notebooks", "ui/*"]
176176

177177
[tool.ruff.lint.isort]
178178
known-first-party = ["databricks.labs.dqx"]

‎src/databricks/labs/dqx/base.py‎

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
from functools import cached_property
44
from typing import final
55
from pyspark.sql import DataFrame, Observation, SparkSession
6-
76
from databricks.labs.dqx.checks_validator import ChecksValidationStatus
87
from databricks.labs.dqx.rule import DQRule
98
from databricks.sdk import WorkspaceClient
@@ -12,8 +11,7 @@
1211

1312
class DQEngineBase(abc.ABC):
1413
def __init__(self, workspace_client: WorkspaceClient):
15-
self._workspace_client = workspace_client
16-
self._spark = SparkSession.builder.getOrCreate()
14+
self._workspace_client = self._verify_workspace_client(workspace_client)
1715

1816
@cached_property
1917
def ws(self) -> WorkspaceClient:
@@ -22,16 +20,7 @@ def ws(self) -> WorkspaceClient:
2220
Ensures workspace connectivity and sets the product info used for
2321
telemetry so that requests are attributed to *dqx*.
2422
"""
25-
return self._verify_workspace_client(self._workspace_client)
26-
27-
@cached_property
28-
def spark(self) -> SparkSession:
29-
"""Return the *SparkSession* associated with this engine.
30-
31-
The session is created during initialization using
32-
*SparkSession.builder.getOrCreate()*.
33-
"""
34-
return self._spark
23+
return self._workspace_client
3524

3625
@staticmethod
3726
@final
@@ -45,7 +34,8 @@ def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
4534
setattr(ws.config, '_product_info', ('dqx', __version__))
4635

4736
# make sure Databricks workspace is accessible
48-
ws.current_user.me()
37+
# use api that works on all workspaces and clusters including group assigned clusters
38+
ws.clusters.select_spark_version()
4939
return ws
5040

5141

0 commit comments

Comments
 (0)