Add documentation for using DQEngine in foreachBatch

docs/dqx/docs/guide/quality_checks_apply.mdx

@@ -418,6 +418,55 @@ You can define checks using DQX classes or as metadata (list of dictionaries) or
   </TabItem>
 </Tabs>
 
+### Applying checks using foreachBatch  
+
+Spark's `foreachBatch` allows users to define an arbitrary processing function run for each streaming batch. DQX requires authentication with a Databricks SDK `WorkspaceClient` which is not supported inside `foreachBatch`. To use DQX in a `foreachBatch` function, create a `DQEngine` instance with a mock `WorkspaceClient`.  
+
+<Admonition type="warning" title="Engine functionality with mocked WorkspaceClient">  
+Some DQX engine methods require a valid Databricks `WorkspaceClient`. Methods must support local execution to be run in a `foreachBatch` function with a mocked `WorkspaceClient`. Check the **Supports local execution** column [here](/docs/reference/engine/#dqx-engine-methods) to verify that DQX engine methods can be executed with a mocked client.  
+</Admonition>  
+
+```python  
+from unittest.mock import MagicMock
+from pyspark.sql import DataFrame
+from databricks.labs.dqx import check_funcs
+from databricks.labs.dqx.engine import DQEngine
+from databricks.labs.dqx.rule import DQRowRule
+from databricks.sdk import WorkspaceClient
+
+# Define DQX checks
+checks = [
+  DQRowRule(
+    name="col_timestamp_is_null",
+    criticality="error",
+    check_func=check_funcs.is_not_null,
+    column="timestamp"
+  ),
+  DQRowRule(
+    name="col_value_is_null",
+    criticality="error",
+    check_func=check_funcs.is_not_null,
+    column="value"
+  )
+]
+
+# Define a `foreachBatch` function to check and write the data
+def validate_and_write_micro_batch_data(batch_df: DataFrame, _: int) -> None:
+  # Create a mock `WorkspaceClient`
+  mock_ws = MagicMock(spec=WorkspaceClient)
+
+  # Create the `DQEngine`
+  dq_engine = DQEngine(mock_ws)
+
+  # Apply the DQX checks and write the data
+  valid_and_invalid_df = dq_engine.apply_checks(batch_df, checks)
+  valid_and_invalid_df.write.format("delta").saveAsTable("catalog.schema.output")  
+
+# Read the data and apply the `foreachBatch` function
+input_data = spark.readStream.format("rate-micro-batch").option("rowsPerBatch", 100).load()
+input_data.writeStream.foreachBatch(validate_and_write_micro_batch_data).start()
+```
+
 ### Saving quality results
 
 The quality check results (DataFrames) can be saved by the user to arbitrary tables using Spark APIs.

tests/e2e/notebooks/apply_checks_foreachbatch_notebook.py

@@ -0,0 +1,48 @@
+# Databricks notebook source
+
+dbutils.widgets.text("test_library_ref", "", "Test Library Ref")
+%pip install 'databricks-labs-dqx @ {dbutils.widgets.get("test_library_ref")}' chispa==0.10.1
+
+# COMMAND ----------
+
+dbutils.library.restartPython()
+
+# COMMAND ----------
+# DBTITLE 1,test_apply_checks_foreachbatch
+
+from unittest.mock import MagicMock
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import lit
+from chispa import assert_df_equality
+
+from databricks.labs.dqx import check_funcs
+from databricks.labs.dqx.engine import DQEngine
+from databricks.labs.dqx.rule import DQRowRule
+from databricks.labs.dqx.schema import dq_result_schema
+from databricks.sdk import WorkspaceClient
+
+def test_apply_checks_foreachbatch():
+    input_df = spark.readStream.format("rate-micro-batch").option("rowsPerBatch", 100).load()
+    checks = [
+        DQRowRule(  # 'rate-micro-batch' source has a non-null 'value' column of type 'long'
+            name="value_is_null_or_empty",
+            criticality="warn",
+            check_func=check_funcs.is_not_null_and_not_empty,
+            column="value",
+        ),
+    ]
+
+    def batch_handler_function(batch_df: DataFrame, _: int) -> None:
+        mock_ws = MagicMock(spec=WorkspaceClient)
+        engine = DQEngine(mock_ws)
+        expected = batch_df.select(
+            "*",
+            lit(None).cast(dq_result_schema.simpleString()).alias("_errors"),
+            lit(None).cast(dq_result_schema.simpleString()).alias("_warnings"),
+        )
+        actual = engine.apply_checks(batch_df, checks)
+        assert_df_equality(actual, expected)
+
+    input_df.writeStream.trigger(availableNow=True).foreachBatch(batch_handler_function).start().awaitTermination()
+
+test_apply_checks_foreachbatch()

tests/e2e/test_apply_checks_foreachbatch.py

@@ -0,0 +1,18 @@
+import logging
+from pathlib import Path
+
+from tests.e2e.conftest import run_notebook_job
+
+
+logger = logging.getLogger(__name__)
+
+
+def test_run_apply_checks_foreachbatch_notebook(make_notebook, make_job, library_ref):
+    notebook_path = Path(__file__).parent / "notebooks" / "apply_checks_foreachbatch_notebook.py"
+    run_notebook_job(
+        notebook_path=notebook_path,
+        make_notebook=make_notebook,
+        make_job=make_job,
+        library_reference=library_ref,
+        task_key="observable_metrics_notebook",
+    )

tests/integration/test_apply_checks.py

@@ -3,12 +3,14 @@
 from datetime import datetime
 from pathlib import Path
 from collections.abc import Callable
+
 import yaml
 import pyspark.sql.functions as F
 import pytest
 from pyspark.sql import Column, DataFrame, SparkSession
 from chispa.dataframe_comparer import assert_df_equality  # type: ignore
 
+import databricks.labs.dqx.geo.check_funcs as geo_check_funcs
 from databricks.labs.dqx.errors import MissingParameterError, InvalidCheckError, InvalidParameterError
 from databricks.labs.dqx.check_funcs import sql_query
 from databricks.labs.dqx.config import OutputConfig, FileChecksStorageConfig, ExtraParams, RunConfig
@@ -22,10 +24,9 @@
 )
 from databricks.labs.dqx.schema import dq_result_schema
 from databricks.labs.dqx import check_funcs
-import databricks.labs.dqx.geo.check_funcs as geo_check_funcs
-from tests.integration.conftest import REPORTING_COLUMNS, RUN_TIME, EXTRA_PARAMS, RUN_ID
 
 from tests.conftest import TEST_CATALOG
+from tests.integration.conftest import REPORTING_COLUMNS, RUN_TIME, EXTRA_PARAMS, RUN_ID
 
 
 SCHEMA = "a: int, b: int, c: int"