diff --git a/docs/dqx/docs/guide/data_profiling.mdx b/docs/dqx/docs/guide/data_profiling.mdx index 2f1683bd1..28f6faa52 100644 --- a/docs/dqx/docs/guide/data_profiling.mdx +++ b/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. + +## 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. + + + + ```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) + ``` + + + + +`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). + + +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: + + + + ```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}, + ) + ``` + + + + +`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). + diff --git a/docs/dqx/docs/reference/profiler.mdx b/docs/dqx/docs/reference/profiler.mdx index ddc30694c..5fe8e7a99 100644 --- a/docs/dqx/docs/reference/profiler.mdx +++ b/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. + For comprehensive examples, advanced options, and best practices, see the [Data Profiling Guide](/docs/guide/data_profiling). diff --git a/src/databricks/labs/dqx/profiler/common.py b/src/databricks/labs/dqx/profiler/common.py index 07f72442c..09960320a 100644 --- a/src/databricks/labs/dqx/profiler/common.py +++ b/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): """ diff --git a/src/databricks/labs/dqx/profiler/profile_builder.py b/src/databricks/labs/dqx/profiler/profile_builder.py index 556254e06..471385ae6 100644 --- a/src/databricks/labs/dqx/profiler/profile_builder.py +++ b/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: diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 0d7b24b53..9a88667ee 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/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) diff --git a/src/databricks/labs/dqx/profiler/profiler_column_metrics.py b/src/databricks/labs/dqx/profiler/profiler_column_metrics.py new file mode 100644 index 000000000..a3f920ad3 --- /dev/null +++ b/src/databricks/labs/dqx/profiler/profiler_column_metrics.py @@ -0,0 +1,125 @@ +import logging +from collections.abc import Callable + +from pyspark.sql import Column +from pyspark.sql import functions as F +from pyspark.sql import types as T + +from databricks.labs.dqx.errors import InvalidParameterError +from databricks.labs.dqx.profiler.common import is_text + + +DQProfileColumnMetricFunc = Callable[[T.StructField, str], Column | None] +PROFILE_COLUMN_METRIC_REGISTRY: dict[str, DQProfileColumnMetricFunc] = {} +# Reserved keys that always-on profiler internals own. They are computed inline in +# `_build_column_metrics` (count_non_null) or derived from it (count_null, count), and downstream +# profile builders read them by string key. Allowing registry entries to shadow these keys would +# let a user metric overwrite the inline aggregation via Row.asDict() (last-value-wins on duplicate +# aliases), corrupting count_null/count and every builder that depends on them. +RESERVED_PROFILE_COLUMN_METRIC_KEYS: frozenset[str] = frozenset({"count", "count_non_null", "count_null"}) +logger = logging.getLogger(__name__) + + +def register_profile_column_metric( + profile_column_metric_type: str, +) -> Callable[[DQProfileColumnMetricFunc], DQProfileColumnMetricFunc]: + """ + Registers data quality profile column metric function. The function that may create a column metric depending on + the column type of the input column or other internal logic. Result column is used in an aggregation function + resulting in a single value for a given column and data frame. The aggregation value will be used further to at the profiling + stage to supply common column level metrics to construct corresponding builders. + + Expected signature of the function is as follows: + (field,column_label) -> Column | None + where: + - field: struct field of the profiling column + - column_label: name of the column that is present in the dataframe to be aggregated + The function may return *None* if aggregation is not applicable. + + The following keys are reserved for always-on profiler internals and cannot be used as a + registration key: *count*, *count_non_null*, *count_null*. Registering under any of these keys + raises *InvalidParameterError*. + + Args: + profile_column_metric_type: Key under which the metric is registered and exposed to profile builders. + + Raises: + InvalidParameterError: If *profile_column_metric_type* is one of the reserved keys. + """ + + def wrapper(metric_func: DQProfileColumnMetricFunc) -> DQProfileColumnMetricFunc: + if profile_column_metric_type in RESERVED_PROFILE_COLUMN_METRIC_KEYS: + raise InvalidParameterError( + f"'{profile_column_metric_type}' is a reserved profile column metric key and cannot be registered. " + f"Reserved keys: {sorted(RESERVED_PROFILE_COLUMN_METRIC_KEYS)}." + ) + if profile_column_metric_type in PROFILE_COLUMN_METRIC_REGISTRY: + logger.warning(f"Overwriting profile column metric registered as '{profile_column_metric_type}'") + PROFILE_COLUMN_METRIC_REGISTRY[profile_column_metric_type] = metric_func + return metric_func + + return wrapper + + +def deregister_profile_column_metric(profile_column_metric_type: str) -> None: + """ + Removes a previously registered profile column metric from *PROFILE_COLUMN_METRIC_REGISTRY*. + No-op if no metric is registered under the given key. + + The reserved keys *count*, *count_non_null*, *count_null* are always-on and cannot be + deregistered; passing any of them raises *InvalidParameterError* so callers do not silently + assume a reserved metric has been disabled. + + Args: + profile_column_metric_type: Key under which the metric was registered. + + Raises: + InvalidParameterError: If *profile_column_metric_type* is one of the reserved keys. + """ + if profile_column_metric_type in RESERVED_PROFILE_COLUMN_METRIC_KEYS: + raise InvalidParameterError( + f"'{profile_column_metric_type}' is a reserved profile column metric key and cannot be deregistered. " + f"Reserved keys: {sorted(RESERVED_PROFILE_COLUMN_METRIC_KEYS)}." + ) + PROFILE_COLUMN_METRIC_REGISTRY.pop(profile_column_metric_type, None) + + +def build_registered_metric_aggregations(field: T.StructField, column_label: str) -> list[Column]: + """ + Return aliased aggregation columns from *PROFILE_COLUMN_METRIC_REGISTRY*, skipping reserved keys + and any metric function that returns *None* for the given field. + + Reserved keys (*count*, *count_non_null*, *count_null*) are computed inline by the profiler and + must never be shadowed by a registry entry — see *RESERVED_PROFILE_COLUMN_METRIC_KEYS*. The + reserved-key filter here is defence-in-depth for entries injected directly into the registry + (bypassing the *register_profile_column_metric* guard). + + Args: + field: Struct field of the profiling column. + column_label: Name of the column present in the dataframe to be aggregated. + """ + aggregations: list[Column] = [] + for metric_name, metric_function in PROFILE_COLUMN_METRIC_REGISTRY.items(): + if metric_name in RESERVED_PROFILE_COLUMN_METRIC_KEYS: + continue + metric_col = metric_function(field, column_label) + if metric_col is not None: + aggregations.append(metric_col.alias(metric_name)) + return aggregations + + +@register_profile_column_metric("empty_count") +def empty_count(field: T.StructField, column_label: str) -> Column: + """ + Profiling column metric for empty count. Applicable for text columns only, otherwise returns literal *0* for + backward compatibility. + """ + return F.count_if(F.col(column_label) == "") if is_text(field.dataType) else F.lit(0) + + +@register_profile_column_metric("count_distinct") +def count_distinct(_field: T.StructField, column_label: str) -> Column: + """ + Profiling column metric for count distinct. Applicable for all columns. + """ + return F.countDistinct(column_label) diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index cdd7fc799..efcb89ff5 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -4,15 +4,45 @@ import pytest import pyspark.sql.types as T +from pyspark.sql import functions as F from databricks.sdk.errors import NotFound from databricks.labs.dqx.config import InputConfig, LLMModelConfig -from databricks.labs.dqx.errors import InvalidConfigError +from databricks.labs.dqx.errors import InvalidConfigError, InvalidParameterError from databricks.labs.dqx.profiler.profiler import DQProfiler, DQProfile +from databricks.labs.dqx.profiler.profiler_column_metrics import ( + PROFILE_COLUMN_METRIC_REGISTRY, + RESERVED_PROFILE_COLUMN_METRIC_KEYS, + register_profile_column_metric, +) +from databricks.labs.dqx.profiler.profile_builder import ( + PROFILE_BUILDER_REGISTRY, + register_profile_builder, +) from tests.constants import TEST_CATALOG +@pytest.fixture +def snapshot_profile_column_metric_registry(): + original_registry = dict(PROFILE_COLUMN_METRIC_REGISTRY) + try: + yield + finally: + PROFILE_COLUMN_METRIC_REGISTRY.clear() + PROFILE_COLUMN_METRIC_REGISTRY.update(original_registry) + + +@pytest.fixture +def snapshot_profile_builder_registry(): + original_registry = dict(PROFILE_BUILDER_REGISTRY) + try: + yield + finally: + PROFILE_BUILDER_REGISTRY.clear() + PROFILE_BUILDER_REGISTRY.update(original_registry) + + def test_profiler(spark, ws): inp_schema = T.StructType( [ @@ -190,6 +220,162 @@ def test_profiler_timestamp_precision_and_rounding( assert min_max_profiles[0].parameters == expected_parameters +def test_profiler_column_metrics_flow_into_generated_profiles(spark, ws): + # Exercises the _build_column_metrics → _build_profiles_for_column flow through the public profile() API. + # The generated is_not_null + min_max profiles depend on count_non_null, count_null, min and max being + # correctly aggregated and merged for the column. + schema = T.StructType([T.StructField("amount", T.IntegerType())]) + input_df = spark.createDataFrame([[10], [20], [30], [40], [50]], schema=schema) + + profiler = DQProfiler(ws) + _, profiles = profiler.profile( + input_df, + options={"sample_fraction": None, "llm_primary_key_detection": False, "remove_outliers": False}, + ) + + assert DQProfile(name="is_not_null", column="amount", description=None, parameters=None) in profiles + min_max = next(p for p in profiles if p.name == "min_max" and p.column == "amount") + assert min_max.parameters == {"min": 10, "max": 50} + + +def test_profiler_high_null_ratio_column_skips_is_not_null(spark, ws): + # Exercises the count_null derivation in _build_column_metrics (total_count - count_non_null): with + # null_ratio above max_null_ratio the null_or_empty builder must skip is_not_null generation for the column. + schema = T.StructType([T.StructField("sparse", T.IntegerType())]) + input_df = spark.createDataFrame([[None], [None], [None], [None], [1]], schema=schema) + + profiler = DQProfiler(ws) + _, profiles = profiler.profile( + input_df, + options={"sample_fraction": None, "llm_primary_key_detection": False, "max_null_ratio": 0.1}, + ) + + assert not [p for p in profiles if p.name == "is_not_null" and p.column == "sparse"] + + +def test_profiler_uses_registered_custom_column_metric(spark, ws, snapshot_profile_column_metric_registry): + # Verifies the register_profile_column_metric extension point end-to-end: + # a user-registered metric is executed against each column during profiling and its + # aggregated value is exposed under the registered key in the returned summary_stats. + metric_key = "p50" + + @register_profile_column_metric(metric_key) + def _p50(_field, column_label): + return F.percentile_approx(F.col(column_label), 0.5) + + schema = T.StructType([T.StructField("amount", T.IntegerType())]) + input_df = spark.createDataFrame([[10], [20], [30], [40], [50]], schema=schema) + + profiler = DQProfiler(ws) + summary_stats, _ = profiler.profile( + input_df, + options={"sample_fraction": None, "llm_primary_key_detection": False, "remove_outliers": False}, + ) + + assert summary_stats["amount"][metric_key] == 30 + + +def test_profiler_drops_registered_custom_column_metric_that_evaluates_to_null( + spark, ws, snapshot_profile_column_metric_registry +): + # A user-registered metric whose aggregation evaluates to SQL NULL must not surface in + # summary_stats — otherwise downstream consumers (profile builders, count_null derivation) + # would need to defensively handle None for every metric key. + metric_key = "always_null" + + @register_profile_column_metric(metric_key) + def _always_null(_field, _column_label): + return F.max(F.lit(None).cast(T.LongType())) + + schema = T.StructType([T.StructField("amount", T.IntegerType())]) + input_df = spark.createDataFrame([[10], [20], [30], [40], [50]], schema=schema) + + profiler = DQProfiler(ws) + summary_stats, _ = profiler.profile( + input_df, + options={"sample_fraction": None, "llm_primary_key_detection": False, "remove_outliers": False}, + ) + + assert metric_key not in summary_stats["amount"] + + +def test_profiler_rejects_reserved_metric_key_and_keeps_count_derivation_intact( + spark, ws, snapshot_profile_column_metric_registry +): + # Guards against overwriting builtin metric: a user metric registered under the reserved + # key "count_non_null" (or "count" / "count_null") must be refused at registration time so it + # cannot collide with the inline aggregation and corrupt count_null (which is derived as + # total_count - count_non_null). + with pytest.raises(InvalidParameterError): + register_profile_column_metric("count_non_null")(lambda _field, _column_label: F.count("*")) + + for reserved_key in RESERVED_PROFILE_COLUMN_METRIC_KEYS: + assert PROFILE_COLUMN_METRIC_REGISTRY.get(reserved_key) is None + + schema = T.StructType([T.StructField("amount", T.IntegerType())]) + input_df = spark.createDataFrame([[10], [20], [None], [40], [None]], schema=schema) + + profiler = DQProfiler(ws) + summary_stats, _ = profiler.profile( + input_df, + options={"sample_fraction": None, "llm_primary_key_detection": False, "remove_outliers": False}, + ) + + amount_stats = summary_stats["amount"] + assert amount_stats["count"] == 5 + assert amount_stats["count_non_null"] == 3 + assert amount_stats["count_null"] == 2 + assert amount_stats["count_non_null"] + amount_stats["count_null"] == amount_stats["count"] + + +def test_profiler_registered_custom_builder_consumes_custom_metric( + spark, ws, snapshot_profile_column_metric_registry, snapshot_profile_builder_registry +): + # Verifies the full extension path end-to-end (register_profile_column_metric + + # register_profile_builder together): a custom column metric is computed during profiling and + # exposed under its key in profiler_metrics, and a custom builder reads that value to emit a + # DQProfile through the public profile() API. This is the workflow documented in the profiling + # guide, and the composition of the two registries is not covered by any other test. + metric_key = "p10" + + @register_profile_column_metric(metric_key) + def _p10(_field, column_label): + return F.percentile_approx(F.col(column_label), 0.1) + + @register_profile_builder("p10_lower_bound") + def _p10_lower_bound(_df, column_name, _column_type, profiler_metrics, _profiler_options): + p10 = profiler_metrics.get(metric_key) + 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}, + ) + + schema = T.StructType([T.StructField("amount", T.IntegerType())]) + input_df = spark.createDataFrame([[10], [20], [30], [40], [50]], schema=schema) + + profiler = DQProfiler(ws) + summary_stats, profiles = profiler.profile( + input_df, + options={"sample_fraction": None, "llm_primary_key_detection": False, "remove_outliers": False}, + ) + + # The custom metric reached summary_stats, and the built-in count_distinct metric is present too. + expected_p10 = summary_stats["amount"][metric_key] + assert summary_stats["amount"]["count_distinct"] == 5 + + # The custom builder consumed that metric and emitted a profile with the metric's value. + custom_profiles = [ + p for p in profiles if p.description and p.description.startswith("Lower bound set to 10th percentile") + ] + assert len(custom_profiles) == 1 + assert custom_profiles[0].column == "amount" + assert custom_profiles[0].parameters == {"min": expected_p10} + + def test_profiler_rounding_midnight_behavior(spark, ws, set_utc_timezone): inp_schema = T.StructType( [ @@ -2356,6 +2542,25 @@ def test_profiler_count_distinct_computed(spark, ws): assert stats["value"]["count_distinct"] == 3 +def test_profiler_empty_count_computed(spark, ws): + schema = T.StructType( + [ + T.StructField("label", T.StringType()), + T.StructField("value", T.IntegerType()), + ] + ) + input_df = spark.createDataFrame( + [["a", 1], ["", 2], ["", 3], ["b", 4]], + schema=schema, + ) + + profiler = DQProfiler(ws) + stats, _ = profiler.profile(input_df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + + assert stats["label"]["empty_count"] == 2 + assert stats["value"]["empty_count"] == 0 + + def test_profiler_generates_has_no_outliers_for_clean_numeric_data(spark, ws): """End-to-end: has_no_outliers profile is emitted when the outlier fraction is below the threshold. diff --git a/tests/unit/test_profile_builder.py b/tests/unit/test_profile_builder.py index 22f2aacd0..d2a020a9c 100644 --- a/tests/unit/test_profile_builder.py +++ b/tests/unit/test_profile_builder.py @@ -9,6 +9,7 @@ from databricks.labs.dqx.profiler.profile import DQProfile from databricks.labs.dqx.profiler.profile_builder import ( PROFILE_BUILDER_REGISTRY, + deregister_profile_builder, make_has_no_outliers_profile, make_is_in_profile, make_min_max_profile, @@ -18,6 +19,20 @@ ) +@pytest.fixture +def restore_profile_builder_registry(): + """ + Snapshot PROFILE_BUILDER_REGISTRY before the test and restore it after, + so tests can freely add/overwrite entries without leaking into other tests. + """ + original_registry = dict(PROFILE_BUILDER_REGISTRY) + try: + yield + finally: + PROFILE_BUILDER_REGISTRY.clear() + PROFILE_BUILDER_REGISTRY.update(original_registry) + + @pytest.fixture def mock_df(): df = create_autospec(DataFrame) @@ -81,6 +96,25 @@ def _my_builder(*_): PROFILE_BUILDER_REGISTRY.pop("_test_return", None) +def test_deregister_profile_builder_removes_registered_builder(restore_profile_builder_registry): + @register_profile_builder("_test_deregister") + def _custom_builder(*_): + return None + + assert "_test_deregister" in PROFILE_BUILDER_REGISTRY + + deregister_profile_builder("_test_deregister") + + assert "_test_deregister" not in PROFILE_BUILDER_REGISTRY + + +def test_deregister_profile_builder_missing_key_is_noop(restore_profile_builder_registry): + # Deregistering an unregistered key must not raise, so callers can use it unconditionally in cleanup. + deregister_profile_builder("_never_registered_key") + + assert "_never_registered_key" not in PROFILE_BUILDER_REGISTRY + + # --------------------------------------------------------------------------- # make_null_or_empty_profile — text types # --------------------------------------------------------------------------- diff --git a/tests/unit/test_profiler_column_metrics.py b/tests/unit/test_profiler_column_metrics.py new file mode 100644 index 000000000..0fb0ff477 --- /dev/null +++ b/tests/unit/test_profiler_column_metrics.py @@ -0,0 +1,123 @@ +import pytest +from pyspark.sql import functions as F +from pyspark.sql import types as T + +from databricks.labs.dqx.errors import InvalidParameterError +from databricks.labs.dqx.profiler.profiler_column_metrics import ( + PROFILE_COLUMN_METRIC_REGISTRY, + RESERVED_PROFILE_COLUMN_METRIC_KEYS, + build_registered_metric_aggregations, + deregister_profile_column_metric, + register_profile_column_metric, +) + + +@pytest.fixture +def restore_profile_column_metric_registry(): + """ + Snapshot PROFILE_COLUMN_METRIC_REGISTRY before the test and restore it after, + so tests can freely add/overwrite entries without leaking into other tests. + """ + original_registry = dict(PROFILE_COLUMN_METRIC_REGISTRY) + try: + yield + finally: + PROFILE_COLUMN_METRIC_REGISTRY.clear() + PROFILE_COLUMN_METRIC_REGISTRY.update(original_registry) + + +def test_register_profile_column_metric_registers_under_explicit_type(restore_profile_column_metric_registry): + @register_profile_column_metric("custom_metric_key") + def _test_metric(_field, _column_label): + return None + + assert "custom_metric_key" in PROFILE_COLUMN_METRIC_REGISTRY + assert PROFILE_COLUMN_METRIC_REGISTRY["custom_metric_key"] is _test_metric + + +def test_register_profile_column_metric_overwrites_and_warns(caplog, restore_profile_column_metric_registry): + # Re-registering a non-reserved key replaces the previous function (last-value-wins) and logs a + # warning so an accidental shadowing is visible rather than silent. + @register_profile_column_metric("custom_metric_key") + def _first(_field, _column_label): + return None + + with caplog.at_level("WARNING"): + + @register_profile_column_metric("custom_metric_key") + def _second(_field, _column_label): + return None + + assert PROFILE_COLUMN_METRIC_REGISTRY["custom_metric_key"] is _second + assert "custom_metric_key" in caplog.text + + +def test_deregister_profile_column_metric_removes_registered_metric(restore_profile_column_metric_registry): + @register_profile_column_metric("custom_metric_key") + def _test_metric(_field, _column_label): + return None + + assert "custom_metric_key" in PROFILE_COLUMN_METRIC_REGISTRY + deregister_profile_column_metric("custom_metric_key") + assert "custom_metric_key" not in PROFILE_COLUMN_METRIC_REGISTRY + + +def test_deregister_profile_column_metric_missing_key_is_noop(restore_profile_column_metric_registry): + # Deregistering an unregistered key must not raise, so callers can use it unconditionally in cleanup. + deregister_profile_column_metric("never_registered_key") + assert "never_registered_key" not in PROFILE_COLUMN_METRIC_REGISTRY + + +@pytest.mark.parametrize("reserved_key", sorted(RESERVED_PROFILE_COLUMN_METRIC_KEYS)) +def test_register_profile_column_metric_rejects_reserved_key(reserved_key, restore_profile_column_metric_registry): + snapshot = dict(PROFILE_COLUMN_METRIC_REGISTRY) + + with pytest.raises(InvalidParameterError): + + @register_profile_column_metric(reserved_key) + def _shadow_metric(_field, _column_label): + return None + + # Registry must be left unmodified. + assert PROFILE_COLUMN_METRIC_REGISTRY == snapshot + + +@pytest.mark.parametrize("reserved_key", sorted(RESERVED_PROFILE_COLUMN_METRIC_KEYS)) +def test_deregister_profile_column_metric_rejects_reserved_key(reserved_key, restore_profile_column_metric_registry): + snapshot = dict(PROFILE_COLUMN_METRIC_REGISTRY) + + with pytest.raises(InvalidParameterError): + deregister_profile_column_metric(reserved_key) + + assert PROFILE_COLUMN_METRIC_REGISTRY == snapshot + + +@pytest.mark.parametrize("reserved_key", sorted(RESERVED_PROFILE_COLUMN_METRIC_KEYS)) +def test_build_registered_metric_aggregations_skips_reserved_key_collision( + reserved_key, restore_profile_column_metric_registry +): + # Even when a colliding entry is injected directly into the registry (bypassing the + # register_profile_column_metric guard), the reserved key must be excluded from the + # aggregation list so it cannot shadow the inline count_non_null alias via Row.asDict(). + PROFILE_COLUMN_METRIC_REGISTRY.clear() + + def _shadow(_field, _column_label): + return F.lit(None).cast(T.LongType()) + + PROFILE_COLUMN_METRIC_REGISTRY[reserved_key] = _shadow + + aggregations = build_registered_metric_aggregations(T.StructField("amount", T.IntegerType()), "amount") + assert not aggregations + + +def test_build_registered_metric_aggregations_skips_metrics_returning_none(restore_profile_column_metric_registry): + # Metric functions may return None to opt out for a given field type; those entries must + # not appear in the aggregation list. + PROFILE_COLUMN_METRIC_REGISTRY.clear() + + @register_profile_column_metric("always_none") + def _always_none(_field, _column_label): + return None + + aggregations = build_registered_metric_aggregations(T.StructField("amount", T.IntegerType()), "amount") + assert not aggregations diff --git a/tests/unit/test_profiler_common.py b/tests/unit/test_profiler_common.py index c187236f9..e03d1c602 100644 --- a/tests/unit/test_profiler_common.py +++ b/tests/unit/test_profiler_common.py @@ -1,7 +1,23 @@ import datetime from decimal import Decimal -from databricks.labs.dqx.profiler.common import val_maybe_to_str, val_to_str +import pytest +import pyspark.sql.types as T + +from databricks.labs.dqx.profiler.common import is_text, val_maybe_to_str, val_to_str + + +@pytest.mark.parametrize("column_type", [T.StringType(), T.CharType(10), T.VarcharType(50)]) +def test_is_text_returns_true_for_text_types(column_type): + assert is_text(column_type) is True + + +@pytest.mark.parametrize( + "column_type", + [T.IntegerType(), T.LongType(), T.DoubleType(), T.FloatType(), T.DateType(), T.TimestampType(), T.BooleanType()], +) +def test_is_text_returns_false_for_non_text_types(column_type): + assert is_text(column_type) is False def test_val_to_str():