Skip to content

Refactor profiler column metrics into an extensible registry - #1384

Open
IvannKurchenko wants to merge 14 commits into
databrickslabs:mainfrom
IvannKurchenko:feature/profiler_additional_metrics
Open

Refactor profiler column metrics into an extensible registry#1384
IvannKurchenko wants to merge 14 commits into
databrickslabs:mainfrom
IvannKurchenko:feature/profiler_additional_metrics

Conversation

@IvannKurchenko

@IvannKurchenko IvannKurchenko commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Changes

Sets up an extension point for the profiler by moving inline column-metric aggregation in DQProfiler._profile into a registry-based system. Users can now register their own column metric functions via the @register_profile_column_metric() decorator, mirroring the existing register_rule / register_profile_builder extension patterns.

The three existing metrics (count_non_null, count_distinct, empty_count) are moved into this registry with no behavioural change. _build_column_metrics builds the aggregation from whatever is registered.

Rationale for landing the refactoring without new metrics

The profiling pipeline is: column metrics → profile builder → check. Adding a new metric only adds value once a builder consumes it and generates a check. Two paths were considered for #1067:

  1. Add percentile metrics (p10/p90) with a matching profile builder that emits is_aggr_not_less_than / is_aggr_not_greater_than checks. But these checks are mostly useful for measurement data (revenue, sales amount, latency, temperature) and not meaningful for keys or categorical columns. Applying them indiscriminately would generate false positives. Selective, purpose-aware application is tracked in [FEATURE]: Profile classification support #1343.
  2. Land the refactoring only, expose the extension point, and let users register the metrics they need. This keeps the profiler flexible without shipping metrics that don't yet have a purpose-fit builder.

This PR takes option 2. New built-in metrics can be added later, together with the specific builder that consumes them, once the classification work in #1343 makes selective application safe.

What changed

  • New PROFILE_COLUMN_METRIC_REGISTRY and register_profile_column_metric decorator in profiler/profiler_column_metrics.py
  • Existing metrics (count_non_null, count_distinct, empty_count) moved into the registry
  • DQProfiler._profile refactored: inline aggregation extracted into _build_column_metrics, which iterates the registry
  • is_text helper moved from profile_builder.py to profiler/common.py (now used across modules)

Linked issues

Relates to #1067, related to #1343

Tests

  • added unit tests
  • manually tested
  • added integration tests
  • added end-to-end tests
  • added performance tests

Unit tests cover: registry (register, overwrite, function-name key); built-in metric functions across column types; is_text helper; _build_column_metrics (alias correctness, count_null derivation, summary merge, empty DataFrame, None-returning metrics). Existing integration tests already cover the refactored aggregation path end-to-end.

Documentation and Demos

  • added/updated docs
  • added/updated agent skills
  • added/updated demos

New sections in data_profiling.mdx guide and profiler.mdx reference showing how to register custom metrics; dqx-profile-and-generate/SKILL.md updated with the extension point.

🤖 Generated with Claude Code

@IvannKurchenko
IvannKurchenko marked this pull request as ready for review July 29, 2026 19:35
@IvannKurchenko
IvannKurchenko requested a review from a team as a code owner July 29, 2026 19:35
@IvannKurchenko
IvannKurchenko requested review from pratikk-databricks and removed request for a team July 29, 2026 19:35

@mwojtyczka mwojtyczka left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the profiler column-metrics registry refactor. The extension-point design is reasonable, but the wiring has a blocking crash plus a few correctness regressions in _build_column_metrics — details inline.

Comment thread src/databricks/labs/dqx/profiler/profiler.py
Comment thread src/databricks/labs/dqx/profiler/profiler.py Outdated
Comment thread src/databricks/labs/dqx/profiler/profiler.py Outdated
Comment thread src/databricks/labs/dqx/profiler/profiler.py Outdated
Comment thread src/databricks/labs/dqx/profiler/profiler_column_metrics.py
Comment thread tests/integration/test_profiler.py
@mwojtyczka mwojtyczka added the under-review This PR is currently being reviewed by one of DQX maintainers. label Jul 31, 2026

@mwojtyczka mwojtyczka left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up review of the head commit. The six earlier threads are all addressed — resolved them. A few new points on the extension point and conventions:

"""

def wrapper(metric_func: Callable) -> Callable:
PROFILE_COLUMN_METRIC_REGISTRY[metric_func.__name__] = metric_func

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The registry is keyed by metric_func.__name__, but the sibling extension points this mirrors — register_rule(rule_type) and register_profile_builder(profile_type) — take an explicit name string. Two custom metric functions with the same __name__ (e.g. both named count_non_null in different modules, or wrapped/functools.partial funcs) will silently overwrite each other in PROFILE_COLUMN_METRIC_REGISTRY with no error.

Also, since no argument is captured, the parameterless decorator-factory layer (the required ()) adds nothing over a plain def register_profile_column_metric(metric_func):. Consider either taking an explicit name: str (consistent with the siblings, and lets callers disambiguate) or dropping the factory layer.


# 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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is_text() was moved here, but the identical TextType / TEXT_TYPES definitions still live in profile_builder.py (lines 35-36), which profiler.py continues to import from. The same configuration is now defined in two modules and can drift. Since profile_builder.py already imports is_text from this module, it could import TEXT_TYPES / TextType from here too and drop its own copies (DRY — per the AGENTS.md guideline).

from databricks.labs.dqx.profiler.common import is_text


DQProfileColumnMetricFunc = Callable[[T.StructField, str], Column | None]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DQProfileColumnMetricFunc = Callable[[T.StructField, str], Column | None] is defined but not used to type the decorator: register_profile_column_metric returns a bare Callable and wrapper takes/returns a bare Callable. Using the alias in these signatures lets mypy/basedpyright verify that registered functions match the metric signature. (Note register_profile_builder has the same looseness, so this is optional consistency.)

@register_profile_column_metric()
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docstring uses backticks around 0 (and None on line 28). AGENTS.md forbids backticks around object names in docstrings — they break API-doc rendering; use italics (0, None) instead.

@mwojtyczka mwojtyczka added the needs-changes Changes required after review label Aug 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

needs-changes Changes required after review under-review This PR is currently being reviewed by one of DQX maintainers.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants