Demo for Spark Structured Streaming

demos/dqx_streaming_demo.py

@@ -0,0 +1,214 @@
+# Databricks notebook source
+# MAGIC
+# MAGIC %md
+# MAGIC ### Environment Setup for DQX with Structured Streaming
+# MAGIC
+# MAGIC #### 1\. Library Installation
+# MAGIC
+# MAGIC **Option A — Notebook-scoped installation**
+# MAGIC
+# MAGIC - Use notebook-scoped pip installs, which take effect only for the current interactive notebook session.
+# MAGIC
+# MAGIC ```py
+# MAGIC # Install DQX library in the current notebook session
+# MAGIC # Run this cell before any imports from databricks-labs-dqx
+# MAGIC
+# MAGIC %pip install databricks-labs-dqx
+# MAGIC ```
+# MAGIC
+# MAGIC - If you are using a cluster that was running prior to this install, you may need to restart the Python process for the changes to take full effect:
+# MAGIC
+# MAGIC ```py
+# MAGIC # Optional: Force restart of the Python process if required by your environment
+# MAGIC dbutils.library.restartPython()
+# MAGIC ```
+# MAGIC
+# MAGIC **Option B — Cluster-scoped installation (recommended for production streaming jobs)**
+# MAGIC
+# MAGIC - Install the DQX library to your cluster via the Databricks Libraries UI or by specifying it as a PyPI library dependency during cluster creation.  
+# MAGIC - This ensures the package is available for all jobs or notebooks attached to the cluster, including those running streaming workloads.  
+# MAGIC - The syntax for specifying the package is:  
+# MAGIC   - Library Source: PyPI  
+# MAGIC   - Package: `databricks-labs-dqx`
+# MAGIC
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Install DQX as Library <br>
+# MAGIC For this demo, we will install DQX as library
+# MAGIC
+
+# 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")}'
+else:
+    %pip install databricks-labs-dqx
+
+%restart_python
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Bronze Stream
+# MAGIC
+# MAGIC This cell reads the NYC Taxi 2019 dataset from a Delta table as a streaming DataFrame, which will be used as the input for downstream data quality checks and transformations.
+
+# COMMAND ----------
+
+bronze_stream = (
+    spark.readStream
+    .format("delta")
+    .load("/databricks-datasets/delta-sharing/samples/nyctaxi_2019")
+)
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Define Data Quality Checks
+# MAGIC
+# MAGIC This cell defines a set of data quality checks using YAML syntax. The checks are designed to validate key columns in the NYC Taxi dataset, such as `vendor_id`, `pickup_datetime`, `dropoff_datetime`, `passenger_count`, and `trip_distance`. Each check specifies a function, its arguments, a name, and a criticality level (error or warn). These checks will be used by the DQEngine to enforce data quality rules on the streaming data.
+
+# COMMAND ----------
+
+# Define Data Quality checks
+import yaml
+
+
+checks = yaml.safe_load("""
+- check:
+    function: is_not_null
+    arguments:
+      column: vendor_id
+  name: vendor_id_is_null
+  criticality: error
+- check:
+    function: is_not_null_and_not_empty
+    arguments:
+      column: vendor_id
+      trim_strings: true
+  name: vendor_id_is_null_or_empty
+  criticality: error
+
+- check:
+    function: is_not_null
+    arguments:
+      column: pickup_datetime
+  name: pickup_datetime_is_null
+  criticality: error
+- check:
+    function: is_not_in_future
+    arguments:
+      column: pickup_datetime
+  name: pickup_datetime_isnt_in_range
+  criticality: warn
+
+- check:
+    function: is_not_in_future
+    arguments:
+      column: pickup_datetime
+  name: pickup_datetime_not_in_future
+  criticality: warn
+- check:
+    function: is_not_in_future
+    arguments:
+      column: dropoff_datetime
+  name: dropoff_datetime_not_in_future
+  criticality: warn
+- check:
+    function: is_not_null
+    arguments:
+      column: passenger_count
+  name: passenger_count_is_null
+  criticality: error
+- check:
+    function: is_in_range
+    arguments:
+      column: passenger_count
+      min_limit: 0
+      max_limit: 6
+  name: passenger_incorrect_count
+  criticality: warn
+- check:
+    function: is_not_null
+    arguments:
+      column: trip_distance
+  name: trip_distance_is_null
+  criticality: error
+""")
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Inputs
+# MAGIC - Checkpoint and table locations are set via widgets
+
+# COMMAND ----------
+
+dbutils.widgets.text("silver_checkpoint", "/tmp/dq/structured/bronze_to_silver_checkpoint")
+dbutils.widgets.text("quarantine_checkpoint", "/tmp/dq/structured/bronze_to_quarantine_checkpoint")
+dbutils.widgets.text("silver_table", "/tmp/dq/structured/silver_table")
+dbutils.widgets.text("quarantine_table", "/tmp/dq/structured/quarantine_table")
+
+# COMMAND ----------
+
+silver_checkpoint = dbutils.widgets.get("silver_checkpoint")
+quarantine_checkpoint = dbutils.widgets.get("quarantine_checkpoint")
+silver_table = dbutils.widgets.get("silver_table")
+quarantine_table = dbutils.widgets.get("quarantine_table")
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Apply data quality checks and split the bronze stream into silver and quarantine DataFrames
+
+# COMMAND ----------
+
+from databricks.labs.dqx.engine import DQEngine
+from databricks.sdk import WorkspaceClient
+from pyspark.sql import DataFrame
+
+dq_engine = DQEngine(WorkspaceClient())
+
+def apply_data_quality_and_split(df: DataFrame):
+    return dq_engine.apply_checks_by_metadata_and_split(df, checks)
+
+silver_df, quarantine_df = apply_data_quality_and_split(bronze_stream)
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Write streaming DataFrames to Delta tables with checkpointing
+
+# COMMAND ----------
+
+PATH_PREFIXES = ("/", "dbfs:/", "s3://", "abfss://", "gs://")
+
+def write_stream(df, checkpoint_location, target):
+    writer = (
+        df.writeStream
+            .format("delta")
+            .outputMode("append")
+            .option("checkpointLocation", checkpoint_location)
+            .trigger(availableNow=True)
+    )
+    if target.startswith(PATH_PREFIXES):
+        return writer.start(target)
+    else:
+        return writer.toTable(target)
+
+silver_query = write_stream(silver_df, silver_checkpoint, silver_table)
+quarantine_query = write_stream(quarantine_df, quarantine_checkpoint, quarantine_table)
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Wait for the streams to finish or keep them running for interactive sessions if required
+
+# COMMAND ----------
+
+silver_query.awaitTermination()
+quarantine_query.awaitTermination()

docs/dqx/docs/demos.mdx

@@ -11,6 +11,7 @@ Import the following notebooks in the Databricks workspace to try DQX out:
 * [DQX Quick Start Demo Notebook (library)](https://github.com/databrickslabs/dqx/blob/v0.7.1/demos/dqx_quick_start_demo_library.py) - quickstart on how to use DQX as a library.
 * [DQX Demo Notebook (library)](https://github.com/databrickslabs/dqx/blob/v0.7.1/demos/dqx_demo_library.py) - demonstrates how to use DQX as a library.
 * [DQX Lakeflow Pipelines Demo Notebook](https://github.com/databrickslabs/dqx/blob/v0.7.1/demos/dqx_dlt_demo.py) - demonstrates how to use DQX with Lakeflow Pipelines (formerly Delta Live Tables (DLT)).
+* [DQX Spark Structured Streaming Demo Notebook](https://github.com/databrickslabs/dqx/blob/v0.7.1/demos/dqx_streaming_demo.py) - demonstrates how to use DQX with Spark Structured Streaming.
 * [DQX Asset Bundles Demo](https://github.com/databrickslabs/dqx/blob/v0.7.1/demos/dqx_demo_asset_bundle/README.md) - demonstrates how to use DQX as a library with Databricks Asset Bundles.
 * [DQX Demo with dbt](https://github.com/databrickslabs/dqx/blob/v0.7.1/demos/dqx_demo_dbt/README.md) - demonstrates how to use DQX as a library with dbt projects.
 

tests/e2e/test_run_demos.py

@@ -3,6 +3,7 @@
 
 from datetime import timedelta
 from pathlib import Path
+from uuid import uuid4
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.workspace import ImportFormat
 from databricks.sdk.service.pipelines import NotebookLibrary, PipelinesEnvironment, PipelineLibrary
@@ -191,3 +192,32 @@ def validate_demo_run_status(run: Run, client: WorkspaceClient) -> None:
     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(make_notebook, make_job, tmp_path):
+
+    ws = WorkspaceClient()
+    path = Path(__file__).parent.parent.parent / "demos" / "dqx_streaming_demo.py"
+    with open(path, "rb") as f:
+        notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
+    notebook_path = notebook.as_fuse().as_posix()
+
+    # Use the temporary directory for outputs
+    run_id = str(uuid4())
+    base_output_path = tmp_path / run_id
+    base_parameters = {
+        "silver_checkpoint": f"{base_output_path}/silver_checkpoint",
+        "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,
+    }
+    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)])
+    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),
+    )
+    logging.info(f"Job run {run.run_id} completed successfully for dqx_streaming_demo")