Fork test: PR #1384 - Refactor profiler column metrics into an extensible registry

docs/dqx/docs/guide/data_profiling.mdx

@@ -715,3 +715,89 @@ Combining either limit with `sample_by_column` may filter out values and lead to
 Summary statistics from limited samples may not reflect the characteristics of the overall dataset. Balance the sampling rate and limits
 with your desired profile accuracy. Manually review and tune rules generated from profiles on sample data to ensure correctness.
 </Admonition>
+
+## Extending the profiler with custom column metrics
+
+Use the `register_profile_column_metric` decorator to add your own per-column metrics.
+Each metric is computed once per column (before any profile builder runs) and shared across every profile builder, so builders never recompute it — but each additional registered metric adds an aggregation to **every** profiling run (see the warning below).
+The string passed to the decorator's *profile_column_metric_type* argument (e.g. `"percentile_10"` in the example below) becomes the key under which the value is exposed to profile builders.
+
+A metric function receives the column's *field* (`StructField`) and *column_label* (its name in the DataFrame), and returns a PySpark aggregation `Column`, or `None` to skip the metric for that column type.
+
+See [Custom Column Metrics](/docs/reference/profiler#custom-column-metrics) in the profiler reference for the built-in metric keys (both always-on internals and registry-managed built-ins) that are also available to profile builders.
+
+<Tabs>
+  <TabItem value="Python" label="Python" default>
+    ```python
+    from pyspark.sql import Column
+    from pyspark.sql import functions as F
+    from pyspark.sql import types as T
+    from databricks.labs.dqx.profiler.profiler_column_metrics import register_profile_column_metric
+
+    @register_profile_column_metric("percentile_10")
+    def percentile_10(field: T.StructField, column_label: str) -> Column | None:
+        # Return None to skip for column types where the metric does not apply
+        if not isinstance(field.dataType, T.NumericType):
+            return None
+        return F.percentile_approx(column_label, 0.1)
+    ```
+  </TabItem>
+</Tabs>
+
+<Admonition type="warning" title="Registration has global side effects">
+`register_profile_column_metric` mutates a module-level registry, so a registered metric persists for the lifetime of the Python process and is aggregated on **every** subsequent `DQProfiler.profile(...)` call for every column it applies to. Leaving unused metrics registered adds work to every profiling run and can slow it down. Remove a metric you no longer need with `deregister_profile_column_metric`:
+
+```python
+from databricks.labs.dqx.profiler.profiler_column_metrics import deregister_profile_column_metric
+
+deregister_profile_column_metric("percentile_10")
+```
+
+The call is a no-op if the key is not registered, so it is safe to use unconditionally in cleanup paths (e.g. notebook teardown, test fixtures).
+</Admonition>
+
+Combine this with the `register_profile_builder` decorator to generate data quality rule suggestions based on the metric.
+A builder function receives the *df* (`DataFrame`), *column_name* (`str`), *column_type* (`DataType`), *profiler_metrics* (`dict[str, Any]` — the column-level statistics the profiler computed, keyed by metric type, including any custom metrics you registered), and *profiler_options* (`dict[str, Any]` — profiler configuration).
+It returns a `DQProfile` (a data quality rule suggestion), or `None` to emit nothing for that column:
+
+<Tabs>
+  <TabItem value="Python" label="Python" default>
+    ```python
+    from databricks.labs.dqx.profiler.profile import DQProfile
+    from databricks.labs.dqx.profiler.profile_builder import register_profile_builder
+    from pyspark.sql import DataFrame
+    from pyspark.sql import types as T
+    from typing import Any
+
+    @register_profile_builder("p10_lower_bound")
+    def make_p10_lower_bound_profile(
+        df: DataFrame,
+        column_name: str,
+        column_type: T.DataType,
+        profiler_metrics: dict[str, Any],
+        profiler_options: dict[str, Any],
+    ) -> DQProfile | None:
+        p10 = profiler_metrics.get("percentile_10")
+        if p10 is None:
+            return None
+        return DQProfile(
+            name="min_max",
+            column=column_name,
+            description=f"Lower bound set to 10th percentile ({p10})",
+            parameters={"min": p10},
+        )
+    ```
+  </TabItem>
+</Tabs>
+
+<Admonition type="warning" title="Registration has global side effects">
+`register_profile_builder` shares the same global-registry lifecycle as `register_profile_column_metric` — a registered builder persists for the lifetime of the Python process and runs on **every** subsequent `DQProfiler.profile(...)` call. Remove a builder you no longer need with `deregister_profile_builder`:
+
+```python
+from databricks.labs.dqx.profiler.profile_builder import deregister_profile_builder
+
+deregister_profile_builder("p10_lower_bound")
+```
+
+See [Registration has global side effects](#extending-the-profiler-with-custom-column-metrics) above for the full caveat (unused registrations slow down every profiling run; cleanup is safe to call unconditionally).
+</Admonition>

docs/dqx/docs/reference/profiler.mdx

@@ -124,6 +124,52 @@ The `DQDltGenerator` class creates Delta Live Tables expectation statements from
 | -------------------- | ------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------ |
 | `generate_dlt_rules` | Generates Delta Live Table rules in the specified language.        | `rules`: List of DQProfile objects; `action`: Optional violation action ("drop", "fail", or None); `language`: Target language ("SQL", "Python", or "Python_Dict").     | Yes                      |
 
+## Custom Column Metrics
+
+Use the `register_profile_column_metric` decorator to add your own per-column metrics. Each metric is computed once per column (before any profile builder runs) and shared across every profile builder — but each additional registered metric adds an aggregation to every profiling run.
+
+A metric function receives the column's *field* (`StructField`) and *column_label* (its name in the DataFrame), and returns a PySpark aggregation `Column` or `None` to skip for that column type. The string passed to the decorator's *profile_column_metric_type* argument (e.g. `"percentile_10"`) becomes the key in the metrics dictionary that all profile builders receive.
+
+```python
+from pyspark.sql import Column
+from pyspark.sql import functions as F
+from pyspark.sql import types as T
+from databricks.labs.dqx.profiler.profiler_column_metrics import register_profile_column_metric
+
+@register_profile_column_metric("percentile_10")
+def percentile_10(field: T.StructField, column_label: str) -> Column | None:
+    if not isinstance(field.dataType, T.NumericType):
+        return None
+    return F.percentile_approx(column_label, 0.1)
+```
+
+### Built-in column metrics
+
+Two groups of metrics are available to all profile builders: **always-on internals** that the profiler owns, and **registry-managed** metrics that can be added, replaced, or removed via `register_profile_column_metric` / `deregister_profile_column_metric`.
+
+**Always-on internals (not deregisterable):**
+
+| Metric key | Applicable types | Description |
+|---|---|---|
+| `count` | All | Total row count (before null filtering) |
+| `count_non_null` | All | Non-null value count |
+| `count_null` | All | Null value count (derived from `count - count_non_null`) |
+
+These keys are reserved: calling `register_profile_column_metric("count")` / `"count_non_null"` / `"count_null"` — or `deregister_profile_column_metric` with any of them — raises `InvalidParameterError`. Profile builders can rely on these keys always being present with integer values.
+
+**Registry-managed built-ins:**
+
+| Metric key | Applicable types | Description |
+|---|---|---|
+| `count_distinct` | All | Distinct non-null value count |
+| `empty_count` | Text only | Empty-string count; `0` for non-text |
+
+Both are registered via `@register_profile_column_metric(...)` at import time and can be deregistered like any custom metric.
+
+Spark's `DataFrame.summary()` additionally contributes `mean`, `stddev`, `min`, `25%`, `50%`, `75%`, and `max` for numeric columns (the percentile keys carry the literal `%` suffix that `summary()` emits, so a builder reads them as `profiler_metrics.get("50%")`).
+
+See the [Data Profiling Guide](/docs/guide/data_profiling#extending-the-profiler-with-custom-column-metrics) for a complete example including a custom profile builder that consumes a custom metric.
+
 <Admonition type="info" title="Complete Profiling Guide">
     For comprehensive examples, advanced options, and best practices, see the [Data Profiling
     Guide](/docs/guide/data_profiling).

src/databricks/labs/dqx/profiler/common.py

@@ -3,6 +3,25 @@
 from decimal import Decimal
 from typing import Any
 
+from pyspark.sql import types as T
+
+# Type alias for annotations; use TEXT_TYPES for isinstance() checks.
+TextType = T.CharType | T.StringType | T.VarcharType
+TEXT_TYPES: tuple[type[TextType], ...] = (T.CharType, T.StringType, T.VarcharType)
+
+
+def is_text(column_type: T.DataType) -> bool:
+    """
+    Validates that the input column type is a Spark text type.
+
+    Args:
+        column_type: Input column type
+
+    Returns:
+        True if the column is a Spark text type, otherwise False
+    """
+    return isinstance(column_type, TEXT_TYPES)
+
 
 def val_to_str(value: Any, include_sql_quotes: bool = True):
     """

src/databricks/labs/dqx/profiler/profile_builder.py

@@ -10,6 +10,7 @@
 
 from databricks.labs.dqx.check_funcs import get_limit_expr
 from databricks.labs.dqx.errors import InvalidParameterError
+from databricks.labs.dqx.profiler.common import TEXT_TYPES, is_text
 from databricks.labs.dqx.profiler.profile import DQProfile, DQProfileBuilder
 from databricks.labs.dqx.profiling_utils import calculate_median_absolute_deviation_bounds
 from databricks.labs.dqx.profiler.profile_options import (
@@ -30,10 +31,6 @@
     DEFAULT_PROFILE_OPTIONS,
 )
 
-# Type alias for annotations; use TEXT_TYPES for isinstance() checks.
-TextType = T.CharType | T.StringType | T.VarcharType
-TEXT_TYPES: tuple[type[TextType], ...] = (T.CharType, T.StringType, T.VarcharType)
-
 # Matched pair for serializing timestamp min/max through the Spark fallback: Spark renders with six
 # fractional-second digits and Python parses them back. Kept together as constants so the two patterns
 # can never drift apart (a mismatch would raise ValueError at parse time).
@@ -53,6 +50,17 @@ def wrapper(builder_func: Callable) -> Callable:
     return wrapper
 
 
+def deregister_profile_builder(profile_type: str) -> None:
+    """
+    Removes a previously registered profile builder from *PROFILE_BUILDER_REGISTRY*.
+    No-op if no builder is registered under the given key.
+
+    Args:
+        profile_type: Key under which the builder was registered.
+    """
+    PROFILE_BUILDER_REGISTRY.pop(profile_type, None)
+
+
 @register_profile_builder("null_or_empty")
 def make_null_or_empty_profile(
     _: DataFrame,
@@ -74,7 +82,7 @@ def make_null_or_empty_profile(
     Returns:
         A DQProfile if the correct conditions are met, otherwise None
     """
-    if _is_text(column_type):
+    if is_text(column_type):
         return _make_null_or_empty_profile(column_name, profiler_metrics, profiler_options)
 
     return _make_null_profile(column_name, profiler_metrics, profiler_options)
@@ -169,19 +177,6 @@ def make_min_max_profile(
     )
 
 
-def _is_text(column_type: T.DataType) -> bool:
-    """
-    Validates that the input column type is a Spark text type.
-
-    Args:
-        column_type: Input column type
-
-    Returns:
-        True if the column is a Spark text type, otherwise False
-    """
-    return isinstance(column_type, TEXT_TYPES)
-
-
 def _make_null_or_empty_profile(
     column_name: str, profiler_metrics: dict[str, Any], profiler_options: dict[str, Any]
 ) -> DQProfile | None:

src/databricks/labs/dqx/profiler/profiler.py

@@ -17,8 +17,9 @@
 from databricks.labs.dqx.config import InputConfig, LLMModelConfig
 from databricks.labs.dqx.errors import MissingParameterError, InvalidConfigError
 from databricks.labs.dqx.io import read_input_data, STORAGE_PATH_PATTERN
+from databricks.labs.dqx.profiler.common import TEXT_TYPES, is_text
 from databricks.labs.dqx.profiler.profile import DQProfile
-from databricks.labs.dqx.profiler.profile_builder import PROFILE_BUILDER_REGISTRY, TEXT_TYPES, validate_profile_options
+from databricks.labs.dqx.profiler.profile_builder import PROFILE_BUILDER_REGISTRY, validate_profile_options
 from databricks.labs.dqx.profiler.profile_options import (
     DEFAULT_PROFILE_OPTIONS,
     PROFILE_OPTION_FILTER,
@@ -30,6 +31,9 @@
     PROFILE_OPTION_SAMPLE_SEED,
     PROFILE_OPTION_TRIM_STRINGS,
 )
+from databricks.labs.dqx.profiler.profiler_column_metrics import (
+    build_registered_metric_aggregations,
+)
 from databricks.labs.dqx.utils import list_tables
 from databricks.labs.dqx.telemetry import telemetry_logger
 
@@ -102,9 +106,7 @@ def profile(
         df_columns = [f for f in df.schema.fields if f.name in columns]
         df = df.select(*[f.name for f in df_columns])
 
-        if options is None:
-            options = {}
-
+        options = options or {}
         options = {**DEFAULT_PROFILE_OPTIONS, **options}  # merge default options with user-provided options
         validate_profile_options(options)  # fail fast on misconfiguration before any profiling work
         df = self._sample(df, options)
@@ -439,44 +441,69 @@ def _profile(
             summary_stats: Summary statistics dictionary to update with profiler results.
             total_count: Total number of rows in the input DataFrame.
         """
-        trim_strings = opts.get(PROFILE_OPTION_TRIM_STRINGS, True)
-
         for field in self.get_columns_or_fields(df_cols):
-            field_name = field.name
-            field_type = field.dataType
-            if field_name not in summary_stats:
-                summary_stats[field_name] = {}
-            metrics = summary_stats[field_name]
-
-            column_df = df.select(field_name).dropna()
-            column_label = column_df.columns[0]
-            is_text = isinstance(field_type, TEXT_TYPES)
-            if is_text and trim_strings:
-                column_df = column_df.select(F.trim(F.col(column_label)).alias(column_label))
-
-            aggr_stats = column_df.agg(
-                F.count(column_label).alias("cnt"),
-                F.countDistinct(column_label).alias("cnt_distinct"),
-            ).first()
-            count_non_null = aggr_stats[0] if aggr_stats else 0
-            metrics["count"] = total_count
-            metrics["count_null"] = total_count - count_non_null
-            metrics["count_non_null"] = count_non_null
-            metrics["count_distinct"] = aggr_stats[1] if aggr_stats else 0
-            if is_text:
-                metrics["empty_count"] = column_df.filter(F.col(column_label) == "").count()
-            else:
-                metrics["empty_count"] = 0
+            column_df, column_label = DQProfiler._prepare_column_df(df, field, opts)
+            field_summary_stats = summary_stats.get(field.name, {})
+            metrics = DQProfiler._build_column_metrics(column_df, column_label, field, field_summary_stats, total_count)
+            summary_stats[field.name] = metrics
 
-            self._build_profiles_for_column(column_df, field_name, field_type, metrics, opts, dq_rules)
+            self._build_profiles_for_column(column_df, field, metrics, opts, dq_rules)
 
         self._add_llm_primary_key_for_dataframe(df, dq_rules, summary_stats, opts)
 
+    @staticmethod
+    def _prepare_column_df(df: DataFrame, field: T.StructField, opts: dict[str, Any]) -> tuple[DataFrame, str]:
+        trim_strings = opts.get(PROFILE_OPTION_TRIM_STRINGS, True)
+        field_name = field.name
+        field_type = field.dataType
+
+        column_df = df.select(field_name).dropna()
+        column_label = column_df.columns[0]
+        if is_text(field_type) and trim_strings:
+            column_df = column_df.select(F.trim(F.col(column_label)).alias(column_label))
+        return column_df, column_label
+
+    @staticmethod
+    def _build_column_metrics(
+        column_df: DataFrame,
+        column_label: str,
+        field: T.StructField,
+        field_summary_stats: dict[str, Any],
+        total_count: int,
+    ) -> dict[str, Any]:
+        # count_non_null / count_null / count are always-on internals aggregated (or derived) here,
+        # not through PROFILE_COLUMN_METRIC_REGISTRY. build_registered_metric_aggregations skips
+        # reserved keys as defence-in-depth so a colliding entry injected directly into the registry
+        # can never shadow the inline alias via Row.asDict().
+        count_non_null_alias = "count_non_null"
+        field_metric_aggregations = [F.count(column_label).alias(count_non_null_alias)]
+        field_metric_aggregations.extend(build_registered_metric_aggregations(field, column_label))
+
+        field_aggregation_stats: dict[str, Any] = {}
+        field_aggregation_row = column_df.agg(*field_metric_aggregations).first()
+        if field_aggregation_row:
+            # Drop keys whose aggregation evaluated to SQL NULL so downstream consumers
+            # (profile builders, count_null derivation) can rely on int/comparable values.
+            field_aggregation_stats = {
+                metric_name: metric_value
+                for metric_name, metric_value in field_aggregation_row.asDict().items()
+                if metric_value is not None
+            }
+
+        # Merge order guarantees reserved keys are authoritative: summary_stats first, then
+        # registry aggregations, then the always-on internals (count_non_null / count / count_null)
+        # last, so nothing upstream can shadow them.
+        count_non_null = field_aggregation_stats.get(count_non_null_alias, 0)
+        metrics: dict[str, Any] = {**field_summary_stats, **field_aggregation_stats}
+        metrics["count_non_null"] = count_non_null
+        metrics["count"] = total_count
+        metrics["count_null"] = total_count - count_non_null
+        return metrics
+
     def _build_profiles_for_column(
         self,
         column_df: DataFrame,
-        field_name: str,
-        field_type: T.DataType,
+        field: T.StructField,
         metrics: dict[str, Any],
         opts: dict[str, Any],
         dq_rules: list[DQProfile],
@@ -494,7 +521,7 @@ def _build_profiles_for_column(
         without triggering a second Spark action.
         """
         for profile_type in PROFILE_BUILDER_REGISTRY.values():
-            profile = profile_type.builder(column_df, field_name, field_type, metrics, opts)
+            profile = profile_type.builder(column_df, field.name, field.dataType, metrics, opts)
             if not profile:
                 continue
             dq_rules.append(profile)