Extend summary metrics with per-check-name breakdowns

src/databricks/labs/dqx/engine.py

@@ -165,6 +165,10 @@ def apply_checks(
             ref_dfs,
             rule_set_fingerprint=rule_set_fingerprint,
         )
+
+        if self.observer and self.observer.track_extended_metrics:
+            self.observer.set_check_names([check.name for check in checks])
+
         observed_result = self._observe_metrics(result_df)
 
         if isinstance(observed_result, tuple):

src/databricks/labs/dqx/metrics_observer.py

@@ -1,4 +1,5 @@
-from dataclasses import dataclass
+import re
+from dataclasses import dataclass, field
 from datetime import datetime
 from functools import cached_property
 from typing import Any
@@ -7,6 +8,27 @@
 from pyspark.sql import DataFrame, Observation, SparkSession
 import pyspark.sql.functions as F
 
+from databricks.labs.dqx.errors import DQXError
+
+
+def _sanitize_metric_alias(name: str) -> str:
+    """Sanitize a check name to produce a valid Spark SQL column alias.
+
+    Replaces any non-alphanumeric character (except underscore) with an underscore
+    and collapses consecutive underscores.
+
+    Args:
+        name: The raw check name.
+
+    Returns:
+        A sanitized string safe for use as a Spark SQL alias.
+    """
+    sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", name)
+    sanitized = re.sub(r"_+", "_", sanitized).strip("_")
+    if not sanitized:
+        raise DQXError(f"Sanitizing check '{name}' produces an empty alias")
+    return sanitized
+
 
 OBSERVATION_TABLE_SCHEMA = (
     "run_id string, run_name string, input_location string, output_location string, quarantine_location string, "
@@ -66,10 +88,12 @@ class DQMetricsObserver:
 
     name: str = "dqx"
     custom_metrics: list[str] | None = None
+    track_extended_metrics: bool = False
     id_overwrite: str | None = None
 
     _error_column_name: str = "_errors"
     _warning_column_name: str = "_warnings"
+    _check_names: list[str] = field(default_factory=list)
 
     @cached_property
     def id(self) -> str:
@@ -81,24 +105,64 @@ def id(self) -> str:
         """
         return self.id_overwrite or str(uuid4())
 
-    @cached_property
+    @property
     def metrics(self) -> list[str]:
         """
         Gets the observer metrics as Spark SQL expressions.
 
         Returns:
-            A list of Spark SQL expressions defining the observer metrics (both default and custom).
+            A list of Spark SQL expressions defining the observer metrics (both default, per-check, and custom).
         """
         default_metrics = [
             "count(1) as input_row_count",
             f"count(case when {self._error_column_name} is not null then 1 end) as error_row_count",
             f"count(case when {self._warning_column_name} is not null then 1 end) as warning_row_count",
             f"count(case when {self._error_column_name} is null and {self._warning_column_name} is null then 1 end) as valid_row_count",
         ]
+        if self.track_extended_metrics:
+            default_metrics.extend(self._build_per_check_metrics())
         if self.custom_metrics:
             default_metrics.extend(self.custom_metrics)
         return default_metrics
 
+    def set_check_names(self, check_names: list[str]) -> None:
+        """
+        Sets the check names used to generate per-check summary metrics.
+
+        Args:
+            check_names: List of check names from the applied quality rules.
+        """
+        self._check_names = check_names
+
+    def _build_per_check_metrics(self) -> list[str]:
+        """
+        Builds per-check metric SQL expressions for each registered check name.
+
+        For each check name, generates a count of rows where that check appears in the errors or warnings array.
+
+        Returns:
+            A list of Spark SQL expressions for per-check metrics.
+        """
+        per_check_metrics: list[str] = []
+        seen_aliases: set[str] = set()
+        for check_name in self._check_names:
+            safe_alias = _sanitize_metric_alias(check_name)
+            if safe_alias in seen_aliases:
+                raise DQXError(
+                    f"Check name '{check_name}' produces alias '{safe_alias}' which collides with another check"
+                )
+            seen_aliases.add(safe_alias)
+            escaped_name = check_name.replace("'", "''")
+            per_check_metrics.append(
+                f"count(case when exists({self._error_column_name}, x -> x.name = '{escaped_name}') then 1 end) "
+                f"as {safe_alias}_error_count"
+            )
+            per_check_metrics.append(
+                f"count(case when exists({self._warning_column_name}, x -> x.name = '{escaped_name}') then 1 end) "
+                f"as {safe_alias}_warning_count"
+            )
+        return per_check_metrics
+
     @property
     def observation(self) -> Observation:
         """

tests/integration/test_summary_metrics.py

@@ -2408,3 +2408,115 @@ def test_streaming_observer_metrics_output_and_quarantine_with_empty_checks(
     assert (
         spark.table(quarantine_config.location).count() == 0
     ), f"Quarantine table {quarantine_config.location} has {spark.table(quarantine_config.location).count()} rows"
+
+
+@pytest.mark.parametrize("apply_checks_method", [DQEngine.apply_checks, DQEngine.apply_checks_by_metadata])
+def test_observer_extended_metrics(ws, spark, apply_checks_method):
+    """Test that per-check metrics are included when track_extended_metrics is True."""
+    observer = DQMetricsObserver(name="test_observer", track_extended_metrics=True)
+    dq_engine = DQEngine(workspace_client=ws, spark=spark, observer=observer, extra_params=EXTRA_PARAMS)
+
+    test_df = spark.createDataFrame(
+        [
+            [1, "Alice", 30, 50000],
+            [2, "Bob", 25, 45000],
+            [None, "Charlie", 35, 60000],
+            [4, None, 28, 55000],
+        ],
+        TEST_SCHEMA,
+    )
+
+    if apply_checks_method == DQEngine.apply_checks:
+        checks = deserialize_checks(TEST_CHECKS)
+        checked_df, observation = dq_engine.apply_checks(test_df, checks)
+    elif apply_checks_method == DQEngine.apply_checks_by_metadata:
+        checked_df, observation = dq_engine.apply_checks_by_metadata(test_df, TEST_CHECKS)
+    else:
+        raise ValueError("Invalid 'apply_checks_method' used for testing observable metrics.")
+
+    checked_df.count()  # Trigger an action to get the metrics
+    actual_metrics = observation.get
+
+    # Default metrics
+    assert actual_metrics["input_row_count"] == 4
+    assert actual_metrics["error_row_count"] == 1
+    assert actual_metrics["warning_row_count"] == 1
+    assert actual_metrics["valid_row_count"] == 2
+
+    # Per-check extended metrics
+    assert actual_metrics["id_is_not_null_error_count"] == 1
+    assert actual_metrics["id_is_not_null_warning_count"] == 0
+    assert actual_metrics["name_is_not_null_and_not_empty_error_count"] == 0
+    assert actual_metrics["name_is_not_null_and_not_empty_warning_count"] == 1
+
+
+@pytest.mark.parametrize("apply_checks_method", [DQEngine.apply_checks, DQEngine.apply_checks_by_metadata])
+def test_observer_extended_metrics_checks_change_between_runs(ws, spark, apply_checks_method):
+    """Test that extended metrics reflect the correct checks when the rule set changes between runs."""
+    full_checks_metadata = TEST_CHECKS
+    reduced_checks_metadata = [TEST_CHECKS[0]]
+
+    test_df = spark.createDataFrame(
+        [
+            [1, "Alice", 30, 50000],
+            [2, "Bob", 25, 45000],
+            [None, "Charlie", 35, 60000],
+            [4, None, 28, 55000],
+        ],
+        TEST_SCHEMA,
+    )
+
+    expected_full = {
+        "input_row_count": 4,
+        "error_row_count": 1,
+        "warning_row_count": 1,
+        "valid_row_count": 2,
+        "id_is_not_null_error_count": 1,
+        "id_is_not_null_warning_count": 0,
+        "name_is_not_null_and_not_empty_error_count": 0,
+        "name_is_not_null_and_not_empty_warning_count": 1,
+    }
+    expected_reduced = {
+        "input_row_count": 4,
+        "error_row_count": 1,
+        "warning_row_count": 0,
+        "valid_row_count": 3,
+        "id_is_not_null_error_count": 1,
+        "id_is_not_null_warning_count": 0,
+    }
+
+    if apply_checks_method == DQEngine.apply_checks:
+        full_checks = deserialize_checks(full_checks_metadata)
+        reduced_checks = deserialize_checks(reduced_checks_metadata)
+
+        # First run: both checks
+        observer1 = DQMetricsObserver(name="test_observer", track_extended_metrics=True)
+        dq_engine1 = DQEngine(workspace_client=ws, spark=spark, observer=observer1, extra_params=EXTRA_PARAMS)
+        checked_df1, observation1 = dq_engine1.apply_checks(test_df, full_checks)
+        checked_df1.count()
+        assert observation1.get == expected_full
+
+        # Second run: reduced checks
+        observer2 = DQMetricsObserver(name="test_observer", track_extended_metrics=True)
+        dq_engine2 = DQEngine(workspace_client=ws, spark=spark, observer=observer2, extra_params=EXTRA_PARAMS)
+        checked_df2, observation2 = dq_engine2.apply_checks(test_df, reduced_checks)
+        checked_df2.count()
+        assert observation2.get == expected_reduced
+
+    elif apply_checks_method == DQEngine.apply_checks_by_metadata:
+        # First run: both checks
+        observer1 = DQMetricsObserver(name="test_observer", track_extended_metrics=True)
+        dq_engine1 = DQEngine(workspace_client=ws, spark=spark, observer=observer1, extra_params=EXTRA_PARAMS)
+        checked_df1, observation1 = dq_engine1.apply_checks_by_metadata(test_df, full_checks_metadata)
+        checked_df1.count()
+        assert observation1.get == expected_full
+
+        # Second run: reduced checks
+        observer2 = DQMetricsObserver(name="test_observer", track_extended_metrics=True)
+        dq_engine2 = DQEngine(workspace_client=ws, spark=spark, observer=observer2, extra_params=EXTRA_PARAMS)
+        checked_df2, observation2 = dq_engine2.apply_checks_by_metadata(test_df, reduced_checks_metadata)
+        checked_df2.count()
+        assert observation2.get == expected_reduced
+
+    else:
+        raise ValueError("Invalid 'apply_checks_method' used for testing observable metrics.")

tests/unit/test_observer.py

@@ -1,8 +1,10 @@
 """Unit tests for DQMetricsObserver class."""
 
+import pytest
 from pyspark.sql import Observation
 from pyspark.sql.connect.observation import Observation as SparkConnectObservation
-from databricks.labs.dqx.metrics_observer import DQMetricsObserver
+from databricks.labs.dqx.errors import DQXError
+from databricks.labs.dqx.metrics_observer import DQMetricsObserver, _sanitize_metric_alias
 from databricks.labs.dqx.reporting_columns import DefaultColumnNames
 
 
@@ -100,3 +102,105 @@ def test_dq_observer_observation_property():
     observation = observer.observation
     assert observation is not None
     assert isinstance(observation, Observation | SparkConnectObservation)
+
+
+def test_dq_observer_per_check_metrics_without_check_names():
+    """Test that no per-check metrics are generated when no check names are set."""
+    observer = DQMetricsObserver(track_extended_metrics=True)
+    expected_default_metrics = [
+        "count(1) as input_row_count",
+        "count(case when _errors is not null then 1 end) as error_row_count",
+        "count(case when _warnings is not null then 1 end) as warning_row_count",
+        "count(case when _errors is null and _warnings is null then 1 end) as valid_row_count",
+    ]
+    assert observer.metrics == expected_default_metrics
+
+
+def test_dq_observer_per_check_metrics_with_check_names():
+    """Test that per-check metrics are generated after setting check names."""
+    observer = DQMetricsObserver(track_extended_metrics=True)
+    observer.set_check_names(["id_is_not_null", "name_is_not_empty"])
+
+    metrics = observer.metrics
+    assert "count(1) as input_row_count" in metrics
+
+    assert (
+        "count(case when exists(_errors, x -> x.name = 'id_is_not_null') then 1 end) " "as id_is_not_null_error_count"
+    ) in metrics
+    assert (
+        "count(case when exists(_warnings, x -> x.name = 'id_is_not_null') then 1 end) "
+        "as id_is_not_null_warning_count"
+    ) in metrics
+    assert (
+        "count(case when exists(_errors, x -> x.name = 'name_is_not_empty') then 1 end) "
+        "as name_is_not_empty_error_count"
+    ) in metrics
+    assert (
+        "count(case when exists(_warnings, x -> x.name = 'name_is_not_empty') then 1 end) "
+        "as name_is_not_empty_warning_count"
+    ) in metrics
+
+
+def test_dq_observer_per_check_metrics_ordering():
+    """Test that per-check metrics appear between default and custom metrics."""
+    custom_metrics = ["avg(age) as avg_age"]
+    observer = DQMetricsObserver(custom_metrics=custom_metrics, track_extended_metrics=True)
+    observer.set_check_names(["my_check"])
+
+    metrics = observer.metrics
+    # Default metrics come first (4 items), then per-check (2 per check), then custom
+    assert metrics[0] == "count(1) as input_row_count"
+    assert "my_check_error_count" in metrics[4]
+    assert "my_check_warning_count" in metrics[5]
+    assert metrics[6] == "avg(age) as avg_age"
+
+
+def test_dq_observer_per_check_metrics_with_custom_column_names():
+    """Test that per-check metrics use custom error/warning column names."""
+    observer = DQMetricsObserver(track_extended_metrics=True)
+    observer.set_column_names(error_column_name="dq_errors", warning_column_name="dq_warnings")
+    observer.set_check_names(["my_check"])
+
+    metrics = observer.metrics
+    assert (
+        "count(case when exists(dq_errors, x -> x.name = 'my_check') then 1 end) " "as my_check_error_count"
+    ) in metrics
+    assert (
+        "count(case when exists(dq_warnings, x -> x.name = 'my_check') then 1 end) " "as my_check_warning_count"
+    ) in metrics
+
+
+def test_dq_observer_per_check_metrics_sanitizes_special_characters():
+    """Test that check names with special characters are sanitized in metric aliases."""
+    observer = DQMetricsObserver(track_extended_metrics=True)
+    observer.set_check_names(["check-with-dashes", "check.with.dots"])
+
+    metrics = observer.metrics
+    # Aliases should be sanitized: dashes and dots become underscores
+    assert any("check_with_dashes_error_count" in m for m in metrics)
+    assert any("check_with_dots_error_count" in m for m in metrics)
+
+
+def test_dq_observer_set_check_names_replaces_previous():
+    """Test that calling set_check_names replaces previously set names."""
+    observer = DQMetricsObserver(track_extended_metrics=True)
+    observer.set_check_names(["first_check"])
+    assert any("first_check_error_count" in m for m in observer.metrics)
+
+    observer.set_check_names(["second_check"])
+    assert not any("first_check_error_count" in m for m in observer.metrics)
+    assert any("second_check_error_count" in m for m in observer.metrics)
+
+
+def test_sanitize_metric_alias_empty_raises_error():
+    """Test that _sanitize_metric_alias raises DQXError for names that produce empty aliases."""
+    with pytest.raises(DQXError, match="Sanitizing check '!!!' produces an empty alias"):
+        _sanitize_metric_alias("!!!")
+
+
+def test_dq_observer_duplicate_alias_collision_raises_error():
+    """Test that check names sanitizing to the same alias raise DQXError."""
+    observer = DQMetricsObserver(track_extended_metrics=True)
+    observer.set_check_names(["check-1", "check.1"])
+    with pytest.raises(DQXError, match="produces alias 'check_1' which collides"):
+        observer.metrics