Skip to content
Merged
Show file tree
Hide file tree
Changes from 16 commits
Commits
Show all changes
31 commits
Select commit Hold shift + click to select a range
7a7e092
first draft
IvannKurchenko Jul 29, 2026
268ad0c
linter fixes
IvannKurchenko Jul 29, 2026
c68523f
profiling metrics docs
IvannKurchenko Jul 29, 2026
1e37af9
tests
IvannKurchenko Jul 29, 2026
7b22dfc
Clean up blank lines in profiler unit tests
IvannKurchenko Jul 29, 2026
fc6cc28
integration test of metric via public API
IvannKurchenko Jul 29, 2026
8224e1c
Merge branch 'main' into feature/profiler_additional_metrics
IvannKurchenko Jul 29, 2026
7c4fc5f
Merge branch 'main' into feature/profiler_additional_metrics
mwojtyczka Jul 31, 2026
e86cec8
Merge branch 'main' into feature/profiler_additional_metrics
mwojtyczka Aug 3, 2026
e4ea9bd
code review comments addressed
IvannKurchenko Aug 6, 2026
9b10735
conflicts wigh main resolved
IvannKurchenko Aug 6, 2026
da2e916
code review comments addressed
IvannKurchenko Aug 14, 2026
5457f02
better type hints
IvannKurchenko Aug 14, 2026
49e0877
Merge branch 'main' into feature/profiler_additional_metrics
IvannKurchenko Aug 14, 2026
0f68549
Merge branch 'main' into feature/profiler_additional_metrics
IvannKurchenko Aug 29, 2026
9931ddc
Merge branch 'main' into feature/profiler_additional_metrics
mwojtyczka Aug 31, 2026
fe680d6
Merge branch 'main' into feature/profiler_additional_metrics
IvannKurchenko Sep 1, 2026
998a66d
code review comments addressed
IvannKurchenko Sep 1, 2026
14ebbb4
Merge branch 'feature/profiler_additional_metrics' of github.com:Ivan…
IvannKurchenko Sep 1, 2026
b804e05
Merge branch 'main' into feature/profiler_additional_metrics
mwojtyczka Sep 2, 2026
2507246
code review comments addressed
IvannKurchenko Sep 3, 2026
42a808c
Merge branch 'feature/profiler_additional_metrics' of github.com:Ivan…
IvannKurchenko Sep 3, 2026
ae7e302
Merge branch 'main' into feature/profiler_additional_metrics
IvannKurchenko Sep 8, 2026
fe60249
Merge branch 'main' into feature/profiler_additional_metrics
mwojtyczka Sep 9, 2026
12ede78
Merge branch 'main' into feature/profiler_additional_metrics
IvannKurchenko Sep 11, 2026
e732ce3
code review comments fixes
IvannKurchenko Sep 11, 2026
863ecf2
Merge branch 'feature/profiler_additional_metrics' of github.com:Ivan…
IvannKurchenko Sep 11, 2026
0b94233
docs(profiler): correct summary() percentile metric keys (25%/50%/75%)
mwojtyczka Sep 11, 2026
993583f
Merge branch 'main' into feature/profiler_additional_metrics
mwojtyczka Sep 11, 2026
3103199
docs(profiler): clarify custom metric/builder interfaces and cross-re…
mwojtyczka Sep 11, 2026
a88902e
test(profiler): cover register_profile_builder end-to-end and metric …
mwojtyczka Sep 11, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions docs/dqx/docs/guide/data_profiling.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -715,3 +715,59 @@ 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 `register_profile_column_metric` to add your own per-column metrics.
Metrics are computed once per column before any profile builder runs, so they are available to every builder at no extra cost.
The *profile_column_metric_type* passed to the decorator becomes the key under which the value is available inside 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")
Comment thread
mwojtyczka marked this conversation as resolved.
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>

Combine this with `register_profile_builder` to generate rules based on the metric:

<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>
34 changes: 34 additions & 0 deletions docs/dqx/docs/reference/profiler.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,40 @@ 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 `register_profile_column_metric` to add your own per-column metrics. Metrics are computed once per column before any profile builder runs and are available to all 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 for that column type. The *profile_column_metric_type* passed to the decorator 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

The following metrics are available to all profile builders:

| Metric key | Applicable types | Description |
|---|---|---|
| `count_non_null` | All | Non-null value count |
| `count_null` | All | Null value count |
Comment thread
mwojtyczka marked this conversation as resolved.
Outdated
| `count_distinct` | All | Distinct non-null value count |
| `empty_count` | Text only | Empty-string count; `0` for non-text |

Spark's `DataFrame.summary()` additionally contributes `count`, `mean`, `stddev`, `min`, `25`, `50`, `75`, and `max` for numeric columns.

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).
Expand Down
19 changes: 19 additions & 0 deletions src/databricks/labs/dqx/profiler/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
mwojtyczka marked this conversation as resolved.


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):
"""
Expand Down
20 changes: 2 additions & 18 deletions src/databricks/labs/dqx/profiler/profile_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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).
Expand Down Expand Up @@ -74,7 +71,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)
Expand Down Expand Up @@ -169,19 +166,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:
Expand Down
82 changes: 47 additions & 35 deletions src/databricks/labs/dqx/profiler/profiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -30,6 +31,7 @@
PROFILE_OPTION_SAMPLE_SEED,
PROFILE_OPTION_TRIM_STRINGS,
)
from databricks.labs.dqx.profiler.profiler_column_metrics import PROFILE_COLUMN_METRIC_REGISTRY
from databricks.labs.dqx.utils import list_tables
from databricks.labs.dqx.telemetry import telemetry_logger

Expand Down Expand Up @@ -102,9 +104,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)
Expand Down Expand Up @@ -439,44 +439,56 @@ 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 = self._prepare_column_df(df, field, opts)
metrics = self._build_column_metrics(column_df, column_label, 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)

def _prepare_column_df(self, 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

def _build_column_metrics(
Comment thread
mwojtyczka marked this conversation as resolved.
Comment thread
mwojtyczka marked this conversation as resolved.
self,
column_df: DataFrame,
column_label: str,
field: T.StructField,
summary_stats: dict[str, Any],
Comment thread
mwojtyczka marked this conversation as resolved.
Outdated
total_count: int,
) -> dict[str, Any]:
field_metric_aggregations = []
for metric_name, metric_function in PROFILE_COLUMN_METRIC_REGISTRY.items():
metric_col = metric_function(field, column_label)
if metric_col is not None:
field_metric_aggregations.append(metric_col.alias(metric_name))

field_aggregation_stats: dict[str, Any] = {}
if field_metric_aggregations:
field_aggregation_row = column_df.agg(*field_metric_aggregations).first()
if field_aggregation_row:
field_aggregation_stats = field_aggregation_row.asDict()
Comment thread
mwojtyczka marked this conversation as resolved.
Outdated

field_summary_stats = summary_stats.get(field.name, {})
metrics: dict[str, Any] = {**field_summary_stats, **field_aggregation_stats}
metrics["count"] = total_count
Comment thread
mwojtyczka marked this conversation as resolved.
metrics["count_null"] = total_count - metrics.get("count_non_null", 0)
Comment thread
mwojtyczka marked this conversation as resolved.
Outdated
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],
Expand All @@ -494,7 +506,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)
Expand Down
67 changes: 67 additions & 0 deletions src/databricks/labs/dqx/profiler/profiler_column_metrics.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
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.profiler.common import is_text


DQProfileColumnMetricFunc = Callable[[T.StructField, str], Column | None]
Comment thread
mwojtyczka marked this conversation as resolved.
PROFILE_COLUMN_METRIC_REGISTRY: dict[str, DQProfileColumnMetricFunc] = {}
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.

Args:
profile_column_metric_type: Key under which the metric is registered and exposed to profile builders.
"""

def wrapper(metric_func: DQProfileColumnMetricFunc) -> DQProfileColumnMetricFunc:
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


@register_profile_column_metric("empty_count")
def empty_count(field: T.StructField, column_label: str) -> Column | None:
"""
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)
Comment thread
mwojtyczka marked this conversation as resolved.


@register_profile_column_metric("count_distinct")
def count_distinct(_field: T.StructField, column_label: str) -> Column | None:
"""
Profiling column metric for count distinct. Applicable for all columns.
"""
return F.countDistinct(column_label)


@register_profile_column_metric("count_non_null")
def count_non_null(_field: T.StructField, column_label: str) -> Column | None:
"""
Profiling column metric for count not null values. Applicable for all columns.
"""
return F.count(column_label)
Loading