From 4722a4418e9cf5b03fb8458a18c5dbd2922ddc15 Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 29 Aug 2026 13:35:55 +0200 Subject: [PATCH 1/9] draft implementation --- docs/dqx/docs/reference/profiler.mdx | 158 +++++- src/databricks/labs/dqx/profiler/profile.py | 76 ++- .../labs/dqx/profiler/profile_builder.py | 194 +++---- src/databricks/labs/dqx/profiler/profiler.py | 148 +++++- src/databricks/labs/dqx/profiler/semantic.py | 369 +++++++++++++ tests/integration/profiler/__init__.py | 0 .../profiler/test_semantic_profiling.py | 205 ++++++++ tests/integration/test_profile_builder.py | 37 +- tests/unit/profiler/__init__.py | 0 tests/unit/profiler/test_semantic.py | 319 ++++++++++++ tests/unit/test_profile_builder.py | 487 +++++++++++------- 11 files changed, 1678 insertions(+), 315 deletions(-) create mode 100644 src/databricks/labs/dqx/profiler/semantic.py create mode 100644 tests/integration/profiler/__init__.py create mode 100644 tests/integration/profiler/test_semantic_profiling.py create mode 100644 tests/unit/profiler/__init__.py create mode 100644 tests/unit/profiler/test_semantic.py diff --git a/docs/dqx/docs/reference/profiler.mdx b/docs/dqx/docs/reference/profiler.mdx index ddc30694c..03aed5593 100644 --- a/docs/dqx/docs/reference/profiler.mdx +++ b/docs/dqx/docs/reference/profiler.mdx @@ -5,6 +5,7 @@ sidebar_position: 503 import Admonition from '@theme/Admonition'; import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; +import { FeatureLifecycleStage, AvailableSinceVersion, FeatureTags } from '@site/src/components/FeatureTags'; # DQX Profiler @@ -89,6 +90,7 @@ class DQProfile: description: str | None = None # Optional description of how the rule was generated parameters: dict[str, Any] | None = None # Optional parameters for the rule filter: str | None = None # Optional filter to be applied to the data source + semantic_type: str | None = None # Detected semantic type name (only set when semantic profiling is enabled) ``` ## DQProfile Types @@ -127,4 +129,158 @@ The `DQDltGenerator` class creates Delta Live Tables expectation statements from For comprehensive examples, advanced options, and best practices, see the [Data Profiling Guide](/docs/guide/data_profiling). - + + +## Semantic-aware profiling + + + + + + +Semantic-aware profiling introduces a lightweight classification stage between metric +collection and check generation so that a column receives *one* consistent family of +checks. For example, a *vehicle_type* column classified as an *enum* now receives an +*is_in* candidate only — not both *is_in* and *min_max* — while a continuous +*cargo_weight* column is classified as a *measurement* and receives *min_max* only. + +The feature is fully opt-in: it runs iff a `semantic_registry` is supplied to the +`DQProfiler` constructor. When no registry is supplied, the profiler output is +byte-identical to the pre-feature behaviour and every `DQProfile.semantic_type` +field is `None`. + +### Composing a registry + +`SemanticRegistry` is an immutable, ordered, name-unique collection of detectors. +Its four composition patterns: + +```python +from databricks.labs.dqx.profiler.profiler import DQProfiler +from databricks.labs.dqx.profiler.semantic import ( + DQSemanticType, + DQSemanticTypeDetector, + SemanticRegistry, + default_semantic_detectors, +) + +# 1. Default chain (enum → key → measurement → text) +default_registry = SemanticRegistry.default() + +# 2. Prepend a custom detector so it wins first +def _detect_uuid(ctx): + # ... your logic ... + return DQSemanticType(name="uuid") if ... else None + +uuid_detector = DQSemanticTypeDetector(name="uuid", detect=_detect_uuid) +prepended = SemanticRegistry.default().prepend(uuid_detector) + +# 3. Replace the chain entirely +custom = SemanticRegistry.default().replace([uuid_detector, *default_semantic_detectors()]) + +# 4. Empty registry — semantic types are always None +empty = SemanticRegistry() + +profiler = DQProfiler(ws, semantic_registry=default_registry) +``` + +Both `prepend(...)` and `replace(...)` return a **new** `SemanticRegistry` and +route through the constructor so name-uniqueness is enforced on the derived +instance. Attempting to prepend a detector whose name clashes with an existing +one raises `pydantic.ValidationError`. + +### Built-in detectors + +| Detector | Applicability | Produced `properties` | +|---------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------| +| `enum` | String or integer columns where `cardinality < max_in_count` AND `cardinality / count_non_null ≤ ENUM_MAX_CARDINALITY_RATIO` (default 0.95). Collects distinct values once and hands them downstream. | `values`: tuple of distinct values | +| `key` | Any column with `distinctness ≥ 0.99` AND a type-specific second signal — numeric density ≥ `KEY_MIN_DENSITY_RATIO` (default 0.99) or string length_stability ≥ `KEY_MIN_LENGTH_STABILITY_RATIO` (default 0.95). | `signal`: `"density"` or `"length_stability"` | +| `measurement` | Any numeric column not consumed by an earlier detector. | `distribution`: best-effort guess (`normal` / `uniform` / …) | +| `text` | Any string column not consumed by an earlier detector. | (none) | + +`key` requires **two** positive signals — distinctness alone was over-classifying +continuous numeric measurements and free-form short strings. For a UUID-like +`order_id` the second signal is length stability (every value is 36 chars); for +an auto-incrementing `user_id` it is density (distinct values densely cover the +observed value range). Users can tune the thresholds via a custom detector. + +### Planned specialised detectors (example snippets) + +Specialised format detectors are documented here as example patterns; they are +not yet included in the default chain. Supply them via `prepend(...)` if you +need one today: + +| Planned name | Idea | +|--------------|--------------------------------------------------------------------------------------------------| +| `uuid` | Regex match against the canonical RFC 9562 form on string columns | +| `email` | Basic email regex match on string columns | +| `h3_hash` | H3 index format on string / integer columns | +| `geo_coord` | Numeric columns whose distribution and range look like WGS-84 latitude/longitude | +| `date_int` | Integer columns storing packed calendar dates (e.g. `20260101`) | + +Minimal example — a UUID detector via regex: + +```python +import re +from databricks.labs.dqx.profiler.semantic import DQSemanticType, DQSemanticTypeDetector + +_UUID_RE = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$") + +def _detect_uuid(ctx): + if not isinstance(ctx.column_type, (StringType,)): + return None + # Cheap probe: look at a small sample instead of the whole column. + sample = ctx.df.limit(20).collect() + if not sample: + return None + values = [row[0] for row in sample if row[0] is not None] + if not values or not all(_UUID_RE.match(v) for v in values): + return None + return DQSemanticType(name="uuid") + +uuid_detector = DQSemanticTypeDetector(name="uuid", detect=_detect_uuid) +``` + +### Table metadata + +When `profile_table(...)` is used, DQX best-effort fetches table and per-column +metadata from Unity Catalog and exposes it to detectors via `ctx.metadata`. +The following fields are populated when available: + +* *table_name* +* *table_comment* +* *column_comment* + +Tags are not surfaced in this release — they will follow in a subsequent +release. Metadata fetches are best-effort and degrade gracefully: any SDK +failure (missing table, permission denied, non-UC location) is logged as a +warning and treated as an empty mapping, so metadata errors never block +profiling. On the raw-DataFrame path (`.profile(df, ...)`) there is no table +origin, so `ctx.metadata` is always `{}`. + +### Traceability via `DQProfile.semantic_type` + +When a detector fires, the resulting `DQProfile` records the semantic type name +in its `semantic_type` field so users can trace *why* a check was generated. +When the profiler is used without a `semantic_registry`, this field is `None` +on every profile. + +### Registering custom profile builders + +Existing user-authored builders registered via `@register_profile_builder(...)` +continue to work unchanged. The decorator now also accepts an optional +`type="context"` argument to opt into the preferred contextual callback shape, +which receives a `DQProfileContext` (giving access to *ctx.semantic_type*, +*ctx.metadata*, *ctx.metrics*, and *ctx.options*): + +```python +from databricks.labs.dqx.profiler.profile import DQProfile +from databricks.labs.dqx.profiler.profile_builder import register_profile_builder +from databricks.labs.dqx.profiler.semantic import DQProfileContext + +@register_profile_builder("my_custom", type="context") +def _my_custom_builder(ctx: DQProfileContext) -> DQProfile | None: + if ctx.semantic_type and ctx.semantic_type.name != "measurement": + return None + return DQProfile(name="my_custom", column=ctx.column_name) +``` + diff --git a/src/databricks/labs/dqx/profiler/profile.py b/src/databricks/labs/dqx/profiler/profile.py index e4b0bf660..0c9d7a264 100644 --- a/src/databricks/labs/dqx/profiler/profile.py +++ b/src/databricks/labs/dqx/profiler/profile.py @@ -2,42 +2,78 @@ from dataclasses import dataclass from typing import Any +from pydantic import BaseModel, ConfigDict, model_validator from pyspark.sql import DataFrame from pyspark.sql.types import DataType +from databricks.labs.dqx.profiler.semantic import DQProfileContext + @dataclass(frozen=True) class DQProfile: - """Data quality profile class representing a data quality rule candidate.""" + """Data quality profile class representing a data quality rule candidate. + + Attributes: + name: Profile name (e.g. *is_not_null_or_empty*, *min_max*). + column: Column the profile applies to. + description: Optional human-readable description. + parameters: Optional parameters that specialise the rule. + filter: Optional filter expression scoping the rule. + semantic_type: Optional detected semantic type name (e.g. *enum*, *key*, + *measurement*, *text*) — records *why* the profile was generated + when semantic profiling is enabled. Defaults to *None*, so + pre-existing YAML/JSON round-trips are unaffected. + """ name: str column: str description: str | None = None parameters: dict[str, Any] | None = None filter: str | None = None + semantic_type: str | None = None -@dataclass(frozen=True) -class DQProfileBuilder: - """Data quality profile builder class: a named builder that may produce a DQProfile for a column. +# Legacy 5-argument callback shape used by pre-semantic profile builders. Kept for backward +# compatibility with user-authored builders registered via @register_profile_builder without +# type="context". Prefer *ContextualProfileBuilder* for new code. +ProfileBuilder = Callable[ + [DataFrame, str, DataType, dict[str, Any], dict[str, Any]], + DQProfile | None, +] + +# Preferred single-argument callback shape. Receives *DQProfileContext* so it can inspect +# *ctx.semantic_type* and gate its output on the detected type — giving the builder access +# to the detected semantic type and other profiling metadata. +ContextualProfileBuilder = Callable[[DQProfileContext], DQProfile | None] + + +class DQProfileBuilder(BaseModel): + """Named builder that may produce a *DQProfile* for a column. + + Exactly one of *builder* or *contextual_builder* must be provided. Attributes: - name: Profile type identifier (e.g. "null_or_empty", "is_in", "min_max"). Used to - look up the builder in the registry and in generated rule metadata. - builder: Callable that inspects column data and options and returns a DQProfile when - the column matches the profile criteria, otherwise None. Signature: - - (df, column_name, column_type, profiler_metrics, profiler_options) -> DQProfile | None - - - df: DataFrame for this column (non-null rows only; strings trimmed - when profiler_options["trim_strings"] is True). Used for distinct/min/max etc. - - column_name: Name of the column being profiled. - - column_type: Spark DataType of the column (e.g. StringType(), LongType()). - - profiler_metrics: Column-level statistics from the profiler (e.g. count, - count_null, empty_count, count_non_null). Same key set as summary_stats[column_name]. - - profiler_options: Profiler options for this run (e.g. max_null_ratio, - max_empty_ratio, max_in_count, trim_strings, filter). + name: Profile type identifier (e.g. *min_max*). Used to look up the + builder in the registry and in generated rule metadata. + builder: Legacy 5-argument callback. Left in place for backward + compatibility so existing user-authored builders keep working + when wrapped in *DQProfileBuilder* or registered via + *@register_profile_builder* without *type="context"*. Does not + receive semantic-type information — prefer *contextual_builder* + for new code. + contextual_builder: Preferred single-argument callback receiving a + *DQProfileContext*. Sees the detected semantic type via + *ctx.semantic_type* and can gate its output accordingly. """ + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + name: str - builder: Callable[[DataFrame, str, DataType, dict[str, Any], dict[str, Any]], DQProfile | None] + builder: ProfileBuilder | None = None + contextual_builder: ContextualProfileBuilder | None = None + + @model_validator(mode="after") + def _exactly_one_callback(self) -> "DQProfileBuilder": + if (self.builder is None) == (self.contextual_builder is None): + raise ValueError("DQProfileBuilder requires exactly one of `builder` or `contextual_builder`.") + return self diff --git a/src/databricks/labs/dqx/profiler/profile_builder.py b/src/databricks/labs/dqx/profiler/profile_builder.py index 556254e06..d8c7309df 100644 --- a/src/databricks/labs/dqx/profiler/profile_builder.py +++ b/src/databricks/labs/dqx/profiler/profile_builder.py @@ -3,7 +3,7 @@ import logging from collections.abc import Callable import math -from typing import Any +from typing import Any, Literal, Mapping from pyspark.sql import DataFrame from pyspark.sql import types as T, functions as F @@ -11,6 +11,7 @@ from databricks.labs.dqx.check_funcs import get_limit_expr from databricks.labs.dqx.errors import InvalidParameterError from databricks.labs.dqx.profiler.profile import DQProfile, DQProfileBuilder +from databricks.labs.dqx.profiler.semantic import DQProfileContext from databricks.labs.dqx.profiling_utils import calculate_median_absolute_deviation_bounds from databricks.labs.dqx.profiler.profile_options import ( PROFILE_OPTION_DISTINCT_RATIO, @@ -45,74 +46,101 @@ logger = logging.getLogger(__name__) -def register_profile_builder(profile_type: str) -> Callable: +def register_profile_builder( + profile_type: str, + *, + type: Literal["legacy", "context"] | None = None, +) -> Callable: + """Register a profile builder in *PROFILE_BUILDER_REGISTRY*. + + Args: + profile_type: Registry key (e.g. *min_max*). + type: Callback shape. + * *None* (default) or *"legacy"* — builder is a 5-argument callback + matching the *ProfileBuilder* type alias. Registered as + *DQProfileBuilder(name=..., builder=fn)*. Backward-compatible path. + * *"context"* — builder is a single-argument callback matching + *ContextualProfileBuilder*. Registered as + *DQProfileBuilder(name=..., contextual_builder=fn)*. + """ + def wrapper(builder_func: Callable) -> Callable: - PROFILE_BUILDER_REGISTRY[profile_type] = DQProfileBuilder(name=profile_type, builder=builder_func) + if type == "context": + PROFILE_BUILDER_REGISTRY[profile_type] = DQProfileBuilder( + name=profile_type, contextual_builder=builder_func + ) + else: + PROFILE_BUILDER_REGISTRY[profile_type] = DQProfileBuilder(name=profile_type, builder=builder_func) return builder_func return wrapper -@register_profile_builder("null_or_empty") -def make_null_or_empty_profile( - _: DataFrame, - column_name: str, - column_type: T.DataType, - profiler_metrics: dict[str, Any], - profiler_options: dict[str, Any], -) -> DQProfile | None: +@register_profile_builder("null_or_empty", type="context") +def make_null_or_empty_profile(ctx: DQProfileContext) -> DQProfile | None: """ - Creates an 'is_not_null_or_empty', 'is_not_null', or 'is_not_empty' profile by checking the input column type, - profiled metrics, and profiler options. + Creates an *is_not_null_or_empty*, *is_not_null*, or *is_not_empty* profile by checking + the input column type, profiled metrics, and profiler options. Args: - column_name: Input column name - column_type: Input column type - profiler_metrics: Column-level statistics computed by the DQProfiler - profiler_options: Configuration options for the DQProfiler + ctx: Profile context (column, type, metrics, options). Returns: - A DQProfile if the correct conditions are met, otherwise None + A DQProfile if the correct conditions are met, otherwise None. """ - if _is_text(column_type): - return _make_null_or_empty_profile(column_name, profiler_metrics, profiler_options) + if _is_text(ctx.column_type): + return _make_null_or_empty_profile(ctx.column_name, ctx.metrics, ctx.options) - return _make_null_profile(column_name, profiler_metrics, profiler_options) + return _make_null_profile(ctx.column_name, ctx.metrics, ctx.options) -@register_profile_builder("is_in") -def make_is_in_profile( - df: DataFrame, - column_name: str, - column_type: T.DataType, - profiler_metrics: dict[str, Any], - profiler_options: dict[str, Any], -) -> DQProfile | None: +@register_profile_builder("is_in", type="context") +def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: """ - Creates an 'is_in' profile by checking the input column type, profiled metrics, and profiler options. + Creates an *is_in* profile. + + When a semantic registry is configured and the column was classified as *enum*, + the distinct values collected by the enum detector are reused (no extra Spark + action). Columns classified as other semantic types (*key*, *measurement*, + *text*, user-defined) are skipped. When no semantic type is present (default + path), applicability follows today's byte-identical rules. Args: - df: Single-column DataFrame - column_name: Input column name - column_type: Input column type - profiler_metrics: Column-level statistics computed by the DQProfiler - profiler_options: Configuration options for the DQProfiler + ctx: Profile context (column, type, metrics, options, semantic_type). Returns: - A DQProfile if the correct conditions are met, otherwise None + A DQProfile if the correct conditions are met, otherwise None. """ - if not _supports_distinct(column_type): + if not _supports_distinct(ctx.column_type): return None - total_count = profiler_metrics.get("count", 0) + total_count = ctx.metrics.get("count", 0) if total_count == 0: return None - max_in_count = profiler_options.get(PROFILE_OPTION_MAX_IN_COUNT, 0) - max_distinct_ratio = profiler_options.get(PROFILE_OPTION_DISTINCT_RATIO, 0.0) + semantic_type = ctx.semantic_type + if semantic_type is not None and semantic_type.name != "enum": + # A non-enum semantic type was assigned; do not emit an is_in candidate that would + # contradict the semantic classification. + return None + + if semantic_type is not None and semantic_type.name == "enum": + values = semantic_type.properties.get("values", ()) + distinct_values = list(values) if isinstance(values, tuple) else list(values) + if not distinct_values: + return None + return DQProfile( + name="is_in", + column=ctx.column_name, + parameters={"in": distinct_values}, + filter=ctx.options.get(PROFILE_OPTION_FILTER, None), + ) + + max_in_count = ctx.options.get(PROFILE_OPTION_MAX_IN_COUNT, 0) + max_distinct_ratio = ctx.options.get(PROFILE_OPTION_DISTINCT_RATIO, 0.0) - col = df.columns[0] - distinct_values = [row[0] for row in df.select(col).distinct().collect()] + col = ctx.df.columns[0] + distinct_values = [row[0] for row in ctx.df.select(col).distinct().collect()] distinct_count = len(distinct_values) if distinct_count == 0: # The df passed here has nulls already dropped by the caller. If distinct_count is 0, @@ -124,48 +152,44 @@ def make_is_in_profile( if distinct_count < max_in_count and distinct_ratio < max_distinct_ratio: return DQProfile( name="is_in", - column=column_name, + column=ctx.column_name, parameters={"in": distinct_values}, - filter=profiler_options.get(PROFILE_OPTION_FILTER, None), + filter=ctx.options.get(PROFILE_OPTION_FILTER, None), ) return None -@register_profile_builder("min_max") -def make_min_max_profile( - df: DataFrame, - column_name: str, - column_type: T.DataType, - profiler_metrics: dict[str, Any], - profiler_options: dict[str, Any], -) -> DQProfile | None: +@register_profile_builder("min_max", type="context") +def make_min_max_profile(ctx: DQProfileContext) -> DQProfile | None: """ - Creates a 'min_max' profile by checking the input column type, profiled metrics, and profiler options. + Creates a *min_max* profile. + + Gated on semantic type when semantic profiling is enabled: skipped unless + *ctx.semantic_type* is *None* or its name is *"measurement"*. Args: - df: Single-column DataFrame - column_name: Input column name (used for DQProfile output) - column_type: Input column type - profiler_metrics: Column-level statistics computed by the DQProfiler (includes summary stats) - profiler_options: Configuration options for the DQProfiler + ctx: Profile context (column, type, metrics, options, semantic_type). Returns: - A DQProfile if the correct conditions are met, otherwise None + A DQProfile if the correct conditions are met, otherwise None. """ - if profiler_metrics.get("count_non_null", 0) == 0: + if ctx.metrics.get("count_non_null", 0) == 0: + return None + + if not _supports_min_max(ctx.column_type): return None - if not _supports_min_max(column_type): + if ctx.semantic_type is not None and ctx.semantic_type.name != "measurement": return None - if _remove_outliers(column_name, profiler_options): + if _remove_outliers(ctx.column_name, ctx.options): return _make_min_max_profile_with_outlier_removal( - df, column_name, column_type, profiler_metrics, profiler_options + ctx.df, ctx.column_name, ctx.column_type, dict(ctx.metrics), dict(ctx.options) ) return _make_min_max_profile_without_outlier_removal( - df, column_name, column_type, profiler_metrics, profiler_options + ctx.df, ctx.column_name, ctx.column_type, dict(ctx.metrics), dict(ctx.options) ) @@ -183,7 +207,7 @@ def _is_text(column_type: T.DataType) -> bool: def _make_null_or_empty_profile( - column_name: str, profiler_metrics: dict[str, Any], profiler_options: dict[str, Any] + column_name: str, profiler_metrics: Mapping[str, Any], profiler_options: Mapping[str, Any] ) -> DQProfile | None: """ Creates an 'is_not_null_or_empty', 'is_not_null', or 'is_not_empty' profile for text type columns. @@ -256,7 +280,7 @@ def _make_null_or_empty_profile( def _make_null_profile( - column_name: str, profiler_metrics: dict[str, Any], profiler_options: dict[str, Any] + column_name: str, profiler_metrics: Mapping[str, Any], profiler_options: Mapping[str, Any] ) -> DQProfile | None: """ Builds an 'is_not_null' profile for non-text columns. @@ -319,7 +343,7 @@ def _supports_min_max(column_type: T.DataType) -> bool: ) -def _remove_outliers(column_name: str, profiler_options: dict[str, Any]) -> bool: +def _remove_outliers(column_name: str, profiler_options: Mapping[str, Any]) -> bool: """ Checks if outliers should be removed when generating 'min_max' profiles. @@ -340,7 +364,7 @@ def _remove_outliers(column_name: str, profiler_options: dict[str, Any]) -> bool return column_name in outlier_columns -def _is_has_no_outliers_enabled(column_name: str, profiler_options: dict[str, Any]) -> bool: +def _is_has_no_outliers_enabled(column_name: str, profiler_options: Mapping[str, Any]) -> bool: """ Checks if *has_no_outliers* profiling is enabled for given column. @@ -776,45 +800,43 @@ def _round_decimal(value: decimal.Decimal, rounding_direction: str) -> decimal.D return value -@register_profile_builder("has_no_outliers") -def make_has_no_outliers_profile( - df: DataFrame, - column_name: str, - column_type: T.DataType, - profiler_metrics: dict[str, Any], - profiler_options: dict[str, Any], -) -> DQProfile | None: +@register_profile_builder("has_no_outliers", type="context") +def make_has_no_outliers_profile(ctx: DQProfileContext) -> DQProfile | None: """ Creates a *has_no_outliers* profile using the same MAD method as the *has_no_outliers* check rule. A profile is returned when all the following conditions are met: - - The column type is child of `pyspark.sql.types.NumericType`. + - The column type is child of *pyspark.sql.types.NumericType*. - The DataFrame is non-empty. - The fraction of outliers (values outside *median* ± 3.5 × MAD) is at or below *outliers_ratio*. - Profile generation is enabled at configuration level for all columns or given column. + - When semantic profiling is enabled, the semantic type is unset or *measurement*. Args: - df: The DataFrame to create the profile for. - column_name: Input column name - column_type: Input column type - profiler_metrics: Column-level statistics computed by the DQProfiler - profiler_options: Configuration options for the DQProfiler + ctx: Profile context (column, type, metrics, options, semantic_type). Returns: A DQProfile if all conditions are met, otherwise None. """ + column_name = ctx.column_name + column_type = ctx.column_type + profiler_options = ctx.options + if not isinstance(column_type, T.NumericType): return None + if ctx.semantic_type is not None and ctx.semantic_type.name != "measurement": + return None + if not _is_has_no_outliers_enabled(column_name, profiler_options): return None - total_non_null_count = profiler_metrics.get("count_non_null", 0) + total_non_null_count = ctx.metrics.get("count_non_null", 0) if total_non_null_count == 0: logger.info(f"Column '{column_name}' has no non-null values. Skipping `has_no_outliers` profile generation") return None - bounds = calculate_median_absolute_deviation_bounds(df, column_name) + bounds = calculate_median_absolute_deviation_bounds(ctx.df, column_name) if bounds is None: logger.info( f"MAD bounds were not calculated for column '{column_name}'. Skipping `has_no_outliers` profile generation" @@ -838,7 +860,7 @@ def make_has_no_outliers_profile( below_lower_bound_expr = F.col(column_name) < get_limit_expr(lower_bound) above_upper_bound_expr = F.col(column_name) > get_limit_expr(upper_bound) outside_bounds_expr = below_lower_bound_expr | above_upper_bound_expr - outliers_count = df.filter(outside_bounds_expr).count() + outliers_count = ctx.df.filter(outside_bounds_expr).count() outliers_ratio = float(outliers_count) / total_non_null_count outliers_ratio_threshold = profiler_options.get( @@ -853,7 +875,7 @@ def make_has_no_outliers_profile( name="has_no_outliers", description=f"Column {safe_column_name} has {outliers_ratio * 100:.1f}% of outliers (allowed: {outliers_ratio_threshold * 100:.1f}%). Lower boundary - {lower_bound}, upper boundary - {upper_bound}.", column=column_name, - filter=profiler_options.get(PROFILE_OPTION_FILTER, None), + filter=ctx.options.get(PROFILE_OPTION_FILTER, None), ) return None diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 0d7b24b53..90f876aec 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -1,6 +1,8 @@ +import dataclasses import uuid import logging import os +from collections.abc import Mapping from concurrent import futures from decimal import Decimal, Context from difflib import SequenceMatcher @@ -19,6 +21,10 @@ from databricks.labs.dqx.io import read_input_data, STORAGE_PATH_PATTERN 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.semantic import ( + DQProfileContext, + SemanticRegistry, +) from databricks.labs.dqx.profiler.profile_options import ( DEFAULT_PROFILE_OPTIONS, PROFILE_OPTION_FILTER, @@ -51,6 +57,8 @@ def __init__( workspace_client: WorkspaceClient, spark: SparkSession | None = None, llm_model_config: LLMModelConfig | None = None, + *, + semantic_registry: SemanticRegistry | None = None, ): super().__init__(workspace_client=workspace_client) self.spark = SparkSession.builder.getOrCreate() if spark is None else spark @@ -59,6 +67,7 @@ def __init__( self.llm_engine = ( DQLLMPrimaryKeyEngine(model_config=llm_model_config, spark=self.spark) if LLM_ENABLED else None ) + self._semantic_registry = semantic_registry @staticmethod def get_columns_or_fields(columns: list[T.StructField]) -> list[T.StructField]: @@ -98,6 +107,24 @@ def profile( A tuple containing a dictionary of summary statistics and a list of data quality profiles. """ + return self._profile_dataframe(df, columns, options, table_metadata=None) + + def _profile_dataframe( + self, + df: DataFrame, + columns: list[str] | None, + options: dict[str, Any] | None, + *, + table_metadata: Mapping[str, Any] | None = None, + ) -> tuple[dict[str, Any], list[DQProfile]]: + """Shared private entry point for *.profile()* and *.profile_table()*. + + Accepts an optional *table_metadata* dict fetched from Unity Catalog (see + *_fetch_table_metadata*) that is threaded down to *_profile* so per-column + *DQProfileContext* instances get the corresponding *metadata* mapping. The + public *.profile()* signature stays unchanged; only the internal path carries + table metadata. + """ columns = columns or df.columns df_columns = [f for f in df.schema.fields if f.name in columns] df = df.select(*[f.name for f in df_columns]) @@ -115,7 +142,7 @@ def profile( if total_count == 0: return summary_stats, dq_rules - self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) + self._profile(df, df_columns, dq_rules, options, summary_stats, total_count, table_metadata=table_metadata) return summary_stats, dq_rules @@ -142,7 +169,8 @@ def profile_table( logger.info(f"Profiling {input_config.location} with options: {options}") df = read_input_data(spark=self.spark, input_config=input_config) - return self.profile(df=df, columns=columns, options=options) + table_metadata = self._fetch_table_metadata(input_config.location) + return self._profile_dataframe(df=df, columns=columns, options=options, table_metadata=table_metadata) @telemetry_logger("profiler", "profile_tables_for_patterns") def profile_tables_for_patterns( @@ -415,6 +443,40 @@ def _stratified_sample( logger.info(f"Stratified sampling on column '{sample_by_column}'") return df.sampleBy(sample_by_column, fractions=sample_fractions, seed=sample_seed) + def _fetch_table_metadata(self, location: str) -> dict[str, Any]: + """Best-effort fetch of Unity Catalog table + column metadata. + + Returns a mapping with keys *table_name*, *table_comment*, *columns* + (a dict keyed by column name, each carrying *column_comment*). Tags + are intentionally excluded — the SDK exposes them via separate + calls and current use cases only need the comment fields. Any SDK + failure (missing table, permission denied, non-UC location) is + logged as a warning and returns an empty dict, so metadata errors + never block profiling. + """ + try: + table = self.ws.tables.get(location) + # TODO (IK): Specify exception + except Exception as exc: # noqa: BLE001 — SDK raises many concrete types; degrade gracefully + safe_location = str(location).replace("\n", " ").replace("\r", " ") + logger.warning(f"Could not fetch table metadata for {safe_location}: {exc}") + return {} + + columns: dict[str, dict[str, Any]] = {} + for col in table.columns or []: + col_entry: dict[str, Any] = {} + if col.comment is not None: + col_entry["column_comment"] = col.comment + if col_entry: + columns[col.name] = col_entry + + metadata: dict[str, Any] = {"table_name": table.full_name or location} + if table.comment is not None: + metadata["table_comment"] = table.comment + if columns: + metadata["columns"] = columns + return metadata + def _profile( self, df: DataFrame, @@ -423,6 +485,8 @@ def _profile( opts: dict[str, Any], summary_stats: dict[str, Any], total_count: int, + *, + table_metadata: Mapping[str, Any] | None = None, ) -> None: """ Builds a list of DQProfiles by iterating through DQProfileBuilder builders. @@ -438,6 +502,8 @@ def _profile( opts: Dictionary of options for profiling. summary_stats: Summary statistics dictionary to update with profiler results. total_count: Total number of rows in the input DataFrame. + table_metadata: Table metadata fetched from Unity Catalog when available; + surfaced through *ctx.metadata* on each per-column *DQProfileContext*. """ trim_strings = opts.get(PROFILE_OPTION_TRIM_STRINGS, True) @@ -468,10 +534,45 @@ def _profile( else: metrics["empty_count"] = 0 - self._build_profiles_for_column(column_df, field_name, field_type, metrics, opts, dq_rules) + column_metadata = self._build_column_metadata(field_name, table_metadata) + + self._build_profiles_for_column( + column_df, + field_name, + field_type, + metrics, + opts, + dq_rules, + column_metadata=column_metadata, + ) self._add_llm_primary_key_for_dataframe(df, dq_rules, summary_stats, opts) + @staticmethod + def _build_column_metadata(field_name: str, table_metadata: Mapping[str, Any] | None) -> dict[str, Any]: + """Compose the per-column *metadata* mapping from table-level metadata. + + Fields whose source value is *None* are omitted so detectors can + *metadata.get(...)* cleanly. Returns an empty dict when no table + metadata was supplied. + """ + if not table_metadata: + return {} + metadata: dict[str, Any] = {} + table_name = table_metadata.get("table_name") + if table_name is not None: + metadata["table_name"] = table_name + table_comment = table_metadata.get("table_comment") + if table_comment is not None: + metadata["table_comment"] = table_comment + columns = table_metadata.get("columns") or {} + column_entry = columns.get(field_name) if isinstance(columns, Mapping) else None + if column_entry: + column_comment = column_entry.get("column_comment") + if column_comment is not None: + metadata["column_comment"] = column_comment + return metadata + def _build_profiles_for_column( self, column_df: DataFrame, @@ -480,6 +581,8 @@ def _build_profiles_for_column( metrics: dict[str, Any], opts: dict[str, Any], dq_rules: list[DQProfile], + *, + column_metadata: Mapping[str, Any] | None = None, ) -> None: """Run registered profile builders for a column and append profiles. @@ -493,10 +596,47 @@ def _build_profiles_for_column( *metrics* so that downstream consumers (e.g. LLM primary-key detection) can read them without triggering a second Spark action. """ + metadata_map = dict(column_metadata) if column_metadata else {} + detector_ctx = DQProfileContext( + df=column_df, + column_name=field_name, + column_type=field_type, + metrics=metrics, + options=opts, + metadata=metadata_map, + semantic_type=None, + ) + + semantic_type = None + if self._semantic_registry is not None: + for detector in self._semantic_registry.detectors: + match = detector.detect(detector_ctx) + if match is not None: + semantic_type = match + break + + # Reconstruct via the constructor (not model_copy) so validators run — same rationale as + # SemanticRegistry.prepend/replace. + builder_ctx = DQProfileContext( + df=column_df, + column_name=field_name, + column_type=field_type, + metrics=metrics, + options=opts, + metadata=metadata_map, + semantic_type=semantic_type, + ) + for profile_type in PROFILE_BUILDER_REGISTRY.values(): - profile = profile_type.builder(column_df, field_name, field_type, metrics, opts) + if profile_type.contextual_builder is not None: + profile = profile_type.contextual_builder(builder_ctx) + elif profile_type.builder is not None: + profile = profile_type.builder(column_df, field_name, field_type, dict(metrics), dict(opts)) + if not profile: continue + if semantic_type is not None and profile.semantic_type is None: + profile = dataclasses.replace(profile, semantic_type=semantic_type.name) dq_rules.append(profile) # Write resolved min/max back into metrics so callers (e.g. summary_stats consumers) # can access the final values without re-running Spark aggregates. diff --git a/src/databricks/labs/dqx/profiler/semantic.py b/src/databricks/labs/dqx/profiler/semantic.py new file mode 100644 index 000000000..ea02922c0 --- /dev/null +++ b/src/databricks/labs/dqx/profiler/semantic.py @@ -0,0 +1,369 @@ +"""Semantic-type classification for the DQProfiler. + +Introduces a lightweight classification stage between metric collection and check +generation so a column receives one consistent family of checks (e.g. an enum +column no longer receives both *is_in* and *min_max*). + +All three mutually-dependent public models — *DQSemanticType*, +*DQProfileContext*, and *DQSemanticTypeDetector* — live in this single module +to keep the forward-reference cycle +(*DQSemanticTypeDetector.detect* depends on *DQProfileContext*; +*DQProfileContext.semantic_type* depends on *DQSemanticType*) resolvable at +class-creation time without quoted forward refs or *model_rebuild()*. + +Public surface: + * models: *DQSemanticType*, *DQProfileContext*, *DQSemanticTypeDetector* + * registry: *SemanticRegistry* + * detectors: *DEFAULT_ENUM_DETECTOR*, *DEFAULT_KEY_DETECTOR*, + *DEFAULT_MEASUREMENT_DETECTOR*, *DEFAULT_TEXT_DETECTOR*, + *default_semantic_detectors()* + * thresholds: *ENUM_MAX_CARDINALITY_RATIO*, *KEY_MIN_DENSITY_RATIO*, + *KEY_MIN_LENGTH_STABILITY_RATIO* +""" + +import logging +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator +from pyspark.sql import DataFrame +from pyspark.sql import functions as F +from pyspark.sql import types as T +from pyspark.sql.types import DataType + +from databricks.labs.dqx.profiler.profile_options import PROFILE_OPTION_MAX_IN_COUNT + +logger = logging.getLogger(__name__) + + +Scalar = str | int | float | bool +PropertyValue = Scalar | tuple[Scalar, ...] + + +# TODO (IK): Make properties `BaseModel` too. +class DQSemanticType(BaseModel): + """Classification of a column's semantic meaning. + + Attributes: + name: Globally unique identifier (e.g. *key*, *enum*, *measurement*). + Used in generated check metadata and for detector lookup. + description: Optional human-readable description. + properties: Detector-specific properties (e.g. distribution family, + enum values). Sequence values must be *tuple* (the model's value + type is *Scalar | tuple[Scalar, ...]*), so individual entries + are immutable; the *frozen* model config blocks reassignment of + the *properties* field itself. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str + description: str | None = None + properties: Mapping[str, PropertyValue] = Field(default_factory=dict) + + +class DQProfileContext(BaseModel): + """Immutable context passed to semantic detectors and profile builders. + + Attributes: + df: DataFrame for this column with non-null rows only. Strings are + trimmed when *profiler_options["trim_strings"]* is True. + column_name: Name of the column being profiled. + column_type: Spark DataType of the column. + metrics: Column-level statistics (count, count_null, empty_count, + count_non_null, ...). Same shape as *summary_stats[column_name]*. + options: Profiler options for this run (max_null_ratio, max_empty_ratio, + max_in_count, trim_strings, filter, ...). + metadata: Table/column metadata sourced from Unity Catalog when the + profile is derived from a Delta table. Keys surfaced: + *table_name*, *table_comment*, *column_comment*. Tags are + intentionally excluded. The mapping is empty when profiling a + raw DataFrame with no table origin. + semantic_type: Detected semantic type for this column, or *None* if + no detector matched. Populated only for profile builders — always + *None* for semantic detectors. + """ + + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + + df: DataFrame + column_name: str + column_type: DataType + metrics: Mapping[str, Any] = Field(default_factory=dict) + options: Mapping[str, Any] = Field(default_factory=dict) + metadata: Mapping[str, Any] = Field(default_factory=dict) + semantic_type: DQSemanticType | None = None + + +class DQSemanticTypeDetector(BaseModel): + """Named detector that classifies a column's semantic type. + + Attributes: + name: Detector identifier (also used as the produced type's name when + the detector emits a single canonical type). + detect: Callable receiving a *DQProfileContext* and returning a + *DQSemanticType* when the column matches, otherwise *None*. + """ + + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + + name: str + detect: Callable[[DQProfileContext], DQSemanticType | None] + + +# Enum: cardinality / count_non_null must be at or below this ratio for a +# column to be classified as *enum* (in addition to +# cardinality < max_in_count). A higher ratio ceiling is more permissive +# — it lets slightly less-repeated columns still register as enum-like +# while still rejecting near-unique identifier columns. +ENUM_MAX_CARDINALITY_RATIO = 0.95 + +# Key (numeric): count_distinct / (max - min + 1) must be at or above this +# ratio for a numeric column to be classified as *key*. Tolerates small +# gaps while rejecting continuous distributions where max - min swamps +# count_distinct. +KEY_MIN_DENSITY_RATIO = 0.99 + +# Key (string): min(length) / max(length) over non-null values must be at +# or above this ratio for a string column to be classified as *key*. +# Real-world string keys (UUIDs, SKUs, hashes) have near-uniform length. +KEY_MIN_LENGTH_STABILITY_RATIO = 0.95 + + +# TODO (IK): Use standard PySpark numeric types for this. +_NUMERIC_TYPES: tuple[type[DataType], ...] = ( + T.IntegerType, + T.LongType, + T.ShortType, + T.DoubleType, + T.FloatType, + T.DecimalType, + T.ByteType, +) + +# TODO (IK): Reuse `TEXT_TYPES` from profile_builder.py +_TEXT_TYPES: tuple[type[DataType], ...] = (T.StringType, T.CharType, T.VarcharType) + + +def _is_numeric(column_type: DataType) -> bool: + return isinstance(column_type, _NUMERIC_TYPES) + + +def _is_text(column_type: DataType) -> bool: + return isinstance(column_type, _TEXT_TYPES) + + +def _detect_enum(ctx: DQProfileContext) -> DQSemanticType | None: + """Detect *enum*-like columns: low cardinality and highly repeated values. + + The detector is the sole owner of the enum value collection: when it + fires it runs a single distinct-collect on *ctx.df* and hands the result + to downstream builders through *ctx.semantic_type.properties["values"]*. + """ + column_type = ctx.column_type + if not (_is_text(column_type) or isinstance(column_type, (T.IntegerType, T.LongType, T.ShortType))): + return None + + count_non_null = ctx.metrics.get("count_non_null", 0) + if count_non_null == 0: + return None + + cardinality = ctx.metrics.get("count_distinct", 0) + if cardinality == 0: + return None + + max_in_count = ctx.options.get(PROFILE_OPTION_MAX_IN_COUNT, 0) + if cardinality >= max_in_count: + return None + + if (cardinality / count_non_null) > ENUM_MAX_CARDINALITY_RATIO: + return None + + col = ctx.df.columns[0] + distinct_rows = ctx.df.select(col).distinct().collect() + distinct_values = [row[0] for row in distinct_rows] + try: + sorted_distinct = sorted(distinct_values) + except TypeError: + sorted_distinct = distinct_values + + return DQSemanticType(name="enum", properties={"values": tuple(sorted_distinct)}) + + +def _detect_key(ctx: DQProfileContext) -> DQSemanticType | None: + """Detect *key*-like columns using two positive signals. + + Distinctness (>= 0.99) is the required primary signal for all types; a + second type-specific signal must also fire: + + * Numeric columns — density = count_distinct / (max - min + 1) must be + at or above *KEY_MIN_DENSITY_RATIO*. + * String columns — length_stability = min_length / max_length must be at + or above *KEY_MIN_LENGTH_STABILITY_RATIO*. + """ + count_non_null = ctx.metrics.get("count_non_null", 0) + if count_non_null == 0: + return None + + cardinality = ctx.metrics.get("count_distinct", 0) + if cardinality == 0: + return None + + distinctness = cardinality / count_non_null + if distinctness < 0.99: + return None + + column_type = ctx.column_type + + if isinstance(column_type, (T.IntegerType, T.LongType, T.ShortType)): + min_value = ctx.metrics.get("min") + max_value = ctx.metrics.get("max") + if min_value is None or max_value is None: + return None + span = max_value - min_value + 1 + if span <= 0: + return None + density = cardinality / span + if density < KEY_MIN_DENSITY_RATIO: + return None + return DQSemanticType(name="key", properties={"signal": "density"}) + + if _is_text(column_type): + col = ctx.df.columns[0] + agg = ctx.df.select( + F.min(F.length(F.col(col))).alias("min_len"), + F.max(F.length(F.col(col))).alias("max_len"), + ).first() + if agg is None: + return None + min_len = agg["min_len"] + max_len = agg["max_len"] + if min_len is None or max_len is None: + return None + if max_len == 0: + # Empty-string-only column — division by zero would follow. Not a key. + return None + length_stability = min_len / max_len + if length_stability < KEY_MIN_LENGTH_STABILITY_RATIO: + return None + return DQSemanticType(name="key", properties={"signal": "length_stability"}) + + return None + + +def _detect_measurement(ctx: DQProfileContext) -> DQSemanticType | None: + """Fallback numeric detector: emit *measurement* with a best-effort distribution guess.""" + if not _is_numeric(ctx.column_type): + return None + if ctx.metrics.get("count_non_null", 0) == 0: + return None + + mean = ctx.metrics.get("mean") + stddev = ctx.metrics.get("stddev") + min_value = ctx.metrics.get("min") + max_value = ctx.metrics.get("max") + + distribution = "unknown" + if mean is not None and stddev is not None and min_value is not None and max_value is not None: + try: + span = float(max_value) - float(min_value) + stddev_f = float(stddev) + mean_f = float(mean) + except (TypeError, ValueError): + span = 0.0 + stddev_f = 0.0 + mean_f = 0.0 + if span <= 0: + distribution = "constant" + elif stddev_f == 0: + distribution = "constant" + elif abs(mean_f) > 0 and stddev_f / max(abs(mean_f), 1e-12) > 1.0: + distribution = "exponential" + elif stddev_f / span < 0.15: + distribution = "uniform" + else: + distribution = "normal" + + return DQSemanticType(name="measurement", properties={"distribution": distribution}) + + +def _detect_text(ctx: DQProfileContext) -> DQSemanticType | None: + """Fallback string detector: any string column not otherwise consumed.""" + if not _is_text(ctx.column_type): + return None + if ctx.metrics.get("count_non_null", 0) == 0: + return None + return DQSemanticType(name="text") + + +DEFAULT_ENUM_DETECTOR = DQSemanticTypeDetector(name="enum", detect=_detect_enum) +DEFAULT_KEY_DETECTOR = DQSemanticTypeDetector(name="key", detect=_detect_key) +DEFAULT_MEASUREMENT_DETECTOR = DQSemanticTypeDetector(name="measurement", detect=_detect_measurement) +DEFAULT_TEXT_DETECTOR = DQSemanticTypeDetector(name="text", detect=_detect_text) + + +def default_semantic_detectors() -> tuple[DQSemanticTypeDetector, ...]: + """Return the built-in detector chain in first-match-wins order. + + Order: enum → key (statistical shape) → measurement (numeric fallback) + → text (string fallback). + """ + return ( + DEFAULT_ENUM_DETECTOR, + DEFAULT_KEY_DETECTOR, + DEFAULT_MEASUREMENT_DETECTOR, + DEFAULT_TEXT_DETECTOR, + ) + + +class SemanticRegistry(BaseModel): + """Immutable, ordered, name-unique collection of semantic-type detectors. + + Detectors are consulted in insertion order; the first non-None match wins. + The registry is immutable — the model is frozen (Pydantic rejects field + assignment) and the *detectors* field is a tuple, so external callers + cannot mutate the chain by aliasing. All construction paths (direct + instantiation, *default()*, *prepend()*, *replace()*) route through the + model validator, which enforces detector-name uniqueness. + + Attributes: + detectors: Ordered tuple of detectors in first-match-wins order. + Defaults to an empty tuple; use *SemanticRegistry.default()* to + obtain the built-in chain. + """ + + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + + detectors: tuple[DQSemanticTypeDetector, ...] = () + + @model_validator(mode="after") + def _validate_unique_names(self) -> "SemanticRegistry": + seen: set[str] = set() + for detector in self.detectors: + if detector.name in seen: + raise ValueError(f"Detector {detector.name!r} is registered more than once") + seen.add(detector.name) + return self + + @classmethod + def default(cls) -> "SemanticRegistry": + """Return a registry populated with *default_semantic_detectors()*.""" + return cls(detectors=default_semantic_detectors()) + + def prepend(self, detector: DQSemanticTypeDetector) -> "SemanticRegistry": + """Return a new registry with *detector* at position 0. + + The current instance is left unchanged. Re-invokes the constructor + so the *_validate_unique_names* model-validator runs on the derived + instance — using *model_copy(update=...)* would skip validation and + allow duplicate detector names. + """ + return type(self)(detectors=(detector, *self.detectors)) + + def replace(self, detectors: Sequence[DQSemanticTypeDetector]) -> "SemanticRegistry": + """Return a new registry whose chain is *detectors*. + + The current instance is left unchanged. Re-invokes the constructor + so the uniqueness validator runs on the derived instance. + """ + return type(self)(detectors=tuple(detectors)) diff --git a/tests/integration/profiler/__init__.py b/tests/integration/profiler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/profiler/test_semantic_profiling.py b/tests/integration/profiler/test_semantic_profiling.py new file mode 100644 index 000000000..38b11e580 --- /dev/null +++ b/tests/integration/profiler/test_semantic_profiling.py @@ -0,0 +1,205 @@ +import uuid + +import pyspark.sql.types as T + +from databricks.labs.dqx.config import InputConfig +from databricks.labs.dqx.profiler.profiler import DQProfiler +from databricks.labs.dqx.profiler.semantic import ( + DEFAULT_ENUM_DETECTOR, + DQSemanticType, + DQSemanticTypeDetector, + SemanticRegistry, + default_semantic_detectors, +) + +from tests.constants import TEST_CATALOG + + +DEMO_ROW_COUNT = 60 + + +def _make_demo_df(spark): + """Grounded design fixture: vehicle_type/cargo_weight/deal_value/user_id/order_id/user_name/work_description.""" + schema = T.StructType( + [ + T.StructField("vehicle_type", T.StringType()), + T.StructField("cargo_weight", T.DoubleType()), + T.StructField("deal_value", T.DoubleType()), + T.StructField("user_id", T.LongType()), + T.StructField("order_id", T.StringType()), + T.StructField("user_name", T.StringType()), + T.StructField("work_description", T.StringType()), + ] + ) + kinds = ["car", "truck", "van"] + rows = [] + for i in range(DEMO_ROW_COUNT): + rows.append( + ( + kinds[i % 3], + 100.0 + float(i) * 50.0, + 1000.0 + float(i) * 25.0, + i + 1, + # 36-char UUID-shaped keys — length-stability = 1.0 + str(uuid.UUID(int=i)), + # Free-form names 3..30 chars — length_stability ~= 0.1 + ("A" * (3 + (i % 27))), + "The quick brown fox jumps over the lazy dog " * (1 + i % 3), + ) + ) + return spark.createDataFrame(rows, schema=schema) + + +def _profile_by_column(profiles): + grouped: dict[str, list] = {} + for profile in profiles: + grouped.setdefault(profile.column, []).append(profile) + return grouped + + +def test_profile_without_semantic_registry_matches_pre_feature_output(spark, ws): + """No registry → profiler output is identical to the pre-feature behaviour. + + Every DQProfile.semantic_type must be None because detection did not run. + """ + df = _make_demo_df(spark) + profiler = DQProfiler(ws) + _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + assert profiles, "expected the pre-feature profiler to emit at least one profile" + for profile in profiles: + assert profile.semantic_type is None + + +def test_default_semantic_registry_classifies_grounded_columns(spark, ws): + df = _make_demo_df(spark) + profiler = DQProfiler(ws, semantic_registry=SemanticRegistry.default()) + _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + + by_column = _profile_by_column(profiles) + + vehicle_names = {p.name for p in by_column.get("vehicle_type", [])} + assert "is_in" in vehicle_names + assert "min_max" not in vehicle_names + assert all(p.semantic_type == "enum" for p in by_column.get("vehicle_type", []) if p.name == "is_in") + + for measurement_col in ("cargo_weight", "deal_value"): + names = {p.name for p in by_column.get(measurement_col, [])} + assert "min_max" in names + assert "is_in" not in names + for p in by_column.get(measurement_col, []): + if p.name == "min_max": + assert p.semantic_type == "measurement" + + user_id_names = {p.name for p in by_column.get("user_id", [])} + assert "min_max" not in user_id_names + assert "is_in" not in user_id_names + + order_id_names = {p.name for p in by_column.get("order_id", [])} + assert "is_in" not in order_id_names + assert "min_max" not in order_id_names + + user_name_profiles = by_column.get("user_name", []) + assert all(p.name != "is_in" for p in user_name_profiles) + # user_name has variable length → falls through to text (not key). text emits no additional + # rules; only null_or_empty candidates remain. + for p in user_name_profiles: + if p.semantic_type is not None: + assert p.semantic_type == "text" + + work_desc_profiles = by_column.get("work_description", []) + assert all(p.name != "is_in" for p in work_desc_profiles) + for p in work_desc_profiles: + if p.semantic_type is not None: + assert p.semantic_type == "text" + + +def test_registry_without_enum_falls_through_to_measurement(spark, ws): + df = _make_demo_df(spark) + chain_without_enum = default_semantic_detectors()[1:] # key → measurement → text (drops enum) + profiler = DQProfiler(ws, semantic_registry=SemanticRegistry(detectors=chain_without_enum)) + _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + + by_column = _profile_by_column(profiles) + vehicle_names = {p.name for p in by_column.get("vehicle_type", [])} + # vehicle_type is StringType, so it now falls through to text (no min_max, no is_in). + assert "is_in" not in vehicle_names + assert "min_max" not in vehicle_names + + +def test_prepend_custom_detector_takes_precedence(spark, ws): + """A prepended detector wins over the default enum detector.""" + + def _always_text(_ctx): + return DQSemanticType(name="text", description="forced") + + forced = DQSemanticTypeDetector(name="forced_text", detect=_always_text) + registry = SemanticRegistry.default().prepend(forced) + profiler = DQProfiler(ws, semantic_registry=registry) + df = _make_demo_df(spark) + _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + by_column = _profile_by_column(profiles) + vehicle_names = {p.name for p in by_column.get("vehicle_type", [])} + # forced_text emits "text" — same gate as DEFAULT_TEXT_DETECTOR — so is_in and min_max are skipped + assert "is_in" not in vehicle_names + assert "min_max" not in vehicle_names + + +def test_profile_table_populates_metadata_from_unity_catalog(spark, ws, make_schema, make_random): + """profile_table fetches UC metadata and threads it into ctx.metadata.""" + catalog_name = TEST_CATALOG + schema_name = make_schema(catalog_name=catalog_name).name + table_name = f"{catalog_name}.{schema_name}.t{make_random(10).lower()}" + + schema = T.StructType( + [ + T.StructField("vehicle_type", T.StringType(), metadata={"comment": "kind of vehicle"}), + T.StructField("cargo_weight", T.DoubleType()), + ] + ) + data = [("car", 1.0), ("truck", 2.0), ("van", 3.0)] * 20 + spark.createDataFrame(data, schema).write.format("delta").saveAsTable(table_name) + spark.sql(f"COMMENT ON TABLE {table_name} IS 'demo table with vehicle types'") + spark.sql(f"ALTER TABLE {table_name} ALTER COLUMN vehicle_type COMMENT 'kind of vehicle'") + + captured: list[dict] = [] + + def _spy_detect(ctx): + captured.append(dict(ctx.metadata)) + return None + + spy = DQSemanticTypeDetector(name="spy", detect=_spy_detect) + registry = SemanticRegistry(detectors=(spy, *default_semantic_detectors())) + profiler = DQProfiler(ws, semantic_registry=registry) + profiler.profile_table( + input_config=InputConfig(location=table_name), + options={"sample_fraction": None, "llm_primary_key_detection": False}, + ) + + vehicle_metadata = next((m for m in captured if m.get("column_comment") == "kind of vehicle"), None) + assert vehicle_metadata is not None + assert vehicle_metadata.get("table_name") == table_name + assert vehicle_metadata.get("table_comment") == "demo table with vehicle types" + + +def test_enum_detector_values_reused_by_is_in_builder(spark, ws): + """The is_in builder must reuse the enum detector's values — no double distinct-collect.""" + df = _make_demo_df(spark) + + invocations: list[str] = [] + original = DEFAULT_ENUM_DETECTOR.detect + + def _tracing_detect(ctx): + if ctx.column_name == "vehicle_type": + invocations.append(ctx.column_name) + return original(ctx) + + tracing = DQSemanticTypeDetector(name="enum", detect=_tracing_detect) + registry = SemanticRegistry.default().replace([tracing, *default_semantic_detectors()[1:]]) + profiler = DQProfiler(ws, semantic_registry=registry) + _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + + assert invocations == ["vehicle_type"] + is_in_profiles = [p for p in profiles if p.column == "vehicle_type" and p.name == "is_in"] + assert len(is_in_profiles) == 1 + is_in = is_in_profiles[0] + assert set(is_in.parameters["in"]) == {"car", "truck", "van"} diff --git a/tests/integration/test_profile_builder.py b/tests/integration/test_profile_builder.py index 3ba005a6e..18ff01deb 100644 --- a/tests/integration/test_profile_builder.py +++ b/tests/integration/test_profile_builder.py @@ -4,6 +4,17 @@ import pytest from databricks.labs.dqx.profiler.profile_builder import make_has_no_outliers_profile +from databricks.labs.dqx.profiler.semantic import DQProfileContext + + +def _ctx(df, column_type, metrics, options): + return DQProfileContext( + df=df, + column_name="col", + column_type=column_type, + metrics=metrics, + options=options, + ) @pytest.mark.parametrize( @@ -21,9 +32,7 @@ def test_make_has_no_outliers_profile_empty_data_frame(spark, col_type): """No profile when count_non_null is zero (early-exit path).""" df = spark.createDataFrame([], T.StructType([T.StructField("col", col_type)])) - profiler_metrics = {"count_non_null": 0} - profiler_options = {"outliers_ratio": 0.01} - profile = make_has_no_outliers_profile(df, "col", col_type, profiler_metrics, profiler_options) + profile = make_has_no_outliers_profile(_ctx(df, col_type, {"count_non_null": 0}, {"outliers_ratio": 0.01})) assert profile is None @@ -34,9 +43,7 @@ def test_make_has_no_outliers_profile_bounds_none(spark): calculate_median_absolute_deviation_bounds return None, exercising the bounds=None path. """ df = spark.createDataFrame([], T.StructType([T.StructField("col", T.IntegerType())])) - profiler_metrics = {"count_non_null": 5} - profiler_options = {"outliers_ratio": 0.01} - profile = make_has_no_outliers_profile(df, "col", T.IntegerType(), profiler_metrics, profiler_options) + profile = make_has_no_outliers_profile(_ctx(df, T.IntegerType(), {"count_non_null": 5}, {"outliers_ratio": 0.01})) assert profile is None @@ -74,9 +81,7 @@ def test_make_has_no_outliers_profile_outliers_below_threshold(spark, col_type, MAD bounds: median=6; MAD=3; lower=-4.5; upper=16.5; Hence only 1000 is an outlier. """ df = spark.createDataFrame(data, T.StructType([T.StructField("col", col_type)])) - profiler_metrics = {"count_non_null": len(data)} - profiler_options = {"outliers_ratio": 0.1} - profile = make_has_no_outliers_profile(df, "col", col_type, profiler_metrics, profiler_options) + profile = make_has_no_outliers_profile(_ctx(df, col_type, {"count_non_null": len(data)}, {"outliers_ratio": 0.1})) assert profile is not None assert profile.name == "has_no_outliers" assert profile.column == "col" @@ -115,9 +120,7 @@ def test_make_has_no_outliers_profile_outliers_above_threshold(spark, col_type, # [1..4] + three extreme values → 3 outliers out of 7 ≈ 43 %, threshold 10 % → None # MAD bounds: median=4, MAD=3 → lower=-6.5, upper=14.5 → 100, 200, 300 are outliers df = spark.createDataFrame(data, T.StructType([T.StructField("col", col_type)])) - profiler_metrics = {"count_non_null": len(data)} - profiler_options = {"outliers_ratio": 0.1} - profile = make_has_no_outliers_profile(df, "col", col_type, profiler_metrics, profiler_options) + profile = make_has_no_outliers_profile(_ctx(df, col_type, {"count_non_null": len(data)}, {"outliers_ratio": 0.1})) assert profile is None @@ -130,9 +133,9 @@ def test_make_has_no_outliers_profile_bounds_equal(spark): """ data = [(5,), (5,), (5,), (5,), (5,)] df = spark.createDataFrame(data, T.StructType([T.StructField("col", T.IntegerType())])) - profiler_metrics = {"count_non_null": len(data)} - profiler_options = {"outliers_ratio": 0.1} - profile = make_has_no_outliers_profile(df, "col", T.IntegerType(), profiler_metrics, profiler_options) + profile = make_has_no_outliers_profile( + _ctx(df, T.IntegerType(), {"count_non_null": len(data)}, {"outliers_ratio": 0.1}) + ) assert profile is None @@ -154,7 +157,5 @@ def test_make_has_no_outliers_profile_outliers_null_values(spark, col_type): """ data = [(None,), (None,)] df = spark.createDataFrame(data, T.StructType([T.StructField("col", col_type)])) - profiler_metrics = {"count_non_null": 0} - profiler_options = {"outliers_ratio": 0.1} - profile = make_has_no_outliers_profile(df, "col", col_type, profiler_metrics, profiler_options) + profile = make_has_no_outliers_profile(_ctx(df, col_type, {"count_non_null": 0}, {"outliers_ratio": 0.1})) assert profile is None diff --git a/tests/unit/profiler/__init__.py b/tests/unit/profiler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/profiler/test_semantic.py b/tests/unit/profiler/test_semantic.py new file mode 100644 index 000000000..7023a9558 --- /dev/null +++ b/tests/unit/profiler/test_semantic.py @@ -0,0 +1,319 @@ +from unittest.mock import create_autospec + +import pytest +from pydantic import ValidationError +import pyspark.sql.types as T +from pyspark.sql import DataFrame + +from databricks.labs.dqx.profiler.semantic import ( + DEFAULT_ENUM_DETECTOR, + DEFAULT_KEY_DETECTOR, + DEFAULT_MEASUREMENT_DETECTOR, + DEFAULT_TEXT_DETECTOR, + DQProfileContext, + DQSemanticType, + DQSemanticTypeDetector, + SemanticRegistry, + default_semantic_detectors, +) + + +# --------------------------------------------------------------------------- +# DQSemanticType immutability (Pydantic v2 guarantees) +# --------------------------------------------------------------------------- + + +def test_dq_semantic_type_tuple_property_round_trips(): + sem = DQSemanticType(name="enum", properties={"values": (1, 2, 3)}) + assert sem.properties["values"] == (1, 2, 3) + assert isinstance(sem.properties["values"], tuple) + + +def test_dq_semantic_type_frozen_blocks_field_reassignment(): + sem = DQSemanticType(name="enum") + with pytest.raises(ValidationError): + sem.name = "x" # type: ignore[misc] + + +def test_dq_semantic_type_rejects_unsupported_property_value(): + with pytest.raises(ValidationError): + DQSemanticType(name="enum", properties={"values": {"nested": "dict"}}) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# default_semantic_detectors chain shape +# --------------------------------------------------------------------------- + + +def test_default_semantic_detectors_order(): + chain = default_semantic_detectors() + assert isinstance(chain, tuple) + assert [d.name for d in chain] == ["enum", "key", "measurement", "text"] + + +# --------------------------------------------------------------------------- +# SemanticRegistry immutability + name uniqueness +# --------------------------------------------------------------------------- + + +def test_semantic_registry_default_matches_helper(): + registry = SemanticRegistry.default() + assert registry.detectors == default_semantic_detectors() + + +def test_semantic_registry_empty_constructor_has_empty_detectors(): + assert SemanticRegistry().detectors == () + + +def test_semantic_registry_prepend_returns_new_instance(): + original = SemanticRegistry.default() + snapshot = original.detectors + extra = DQSemanticTypeDetector(name="custom", detect=lambda _ctx: None) + new = original.prepend(extra) + assert new is not original + assert new.detectors[0] is extra + assert original.detectors == snapshot # original unchanged + + +def test_semantic_registry_replace_returns_new_instance(): + original = SemanticRegistry.default() + replacement = (DQSemanticTypeDetector(name="only", detect=lambda _ctx: None),) + new = original.replace(replacement) + assert new is not original + assert new.detectors == replacement + assert original.detectors == default_semantic_detectors() + + +def test_semantic_registry_assignment_raises(): + registry = SemanticRegistry.default() + with pytest.raises(ValidationError): + registry.detectors = () # type: ignore[misc] + + +def test_semantic_registry_prepend_duplicate_name_raises(): + dup = DQSemanticTypeDetector(name="enum", detect=lambda _ctx: None) + with pytest.raises(ValidationError): + SemanticRegistry.default().prepend(dup) + + +def test_semantic_registry_replace_duplicate_names_raises(): + dup1 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) + dup2 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) + with pytest.raises(ValidationError): + SemanticRegistry().replace([dup1, dup2]) + + +def test_semantic_registry_direct_constructor_duplicate_names_raises(): + dup1 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) + dup2 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) + with pytest.raises(ValidationError): + SemanticRegistry(detectors=(dup1, dup2)) + + +# --------------------------------------------------------------------------- +# Chain semantics via a hand-rolled 3-detector chain +# --------------------------------------------------------------------------- + + +def _fake_ctx(column_name="c", column_type=None, metrics=None, options=None): + df = create_autospec(DataFrame) + df.columns = [column_name] + return DQProfileContext( + df=df, + column_name=column_name, + column_type=column_type or T.IntegerType(), + metrics=metrics or {}, + options=options or {}, + ) + + +def test_chain_first_match_wins_and_stops(): + calls: list[str] = [] + + def make_detector(name: str, match: bool): + def _detect(ctx): + calls.append(name) + return DQSemanticType(name=name) if match else None + + return DQSemanticTypeDetector(name=name, detect=_detect) + + registry = SemanticRegistry( + detectors=( + make_detector("a", False), + make_detector("b", True), + make_detector("c", True), # should not be invoked + ) + ) + ctx = _fake_ctx() + picked = None + for detector in registry.detectors: + result = detector.detect(ctx) + if result is not None: + picked = result + break + assert picked is not None + assert picked.name == "b" + assert calls == ["a", "b"] + + +def test_empty_chain_leaves_ctx_semantic_type_none(): + registry = SemanticRegistry() + ctx = _fake_ctx() + match = None + for detector in registry.detectors: + match = detector.detect(ctx) + if match is not None: + break + assert match is None + + +# --------------------------------------------------------------------------- +# Individual detector guards (metric-only paths) +# --------------------------------------------------------------------------- + + +def test_measurement_detector_positive_numeric_column(): + ctx = _fake_ctx( + column_type=T.IntegerType(), + metrics={"count_non_null": 100, "min": 0, "max": 100, "mean": 50.0, "stddev": 15.0}, + ) + result = DEFAULT_MEASUREMENT_DETECTOR.detect(ctx) + assert result is not None + assert result.name == "measurement" + assert "distribution" in result.properties + + +def test_measurement_detector_rejects_string_column(): + ctx = _fake_ctx(column_type=T.StringType(), metrics={"count_non_null": 5}) + assert DEFAULT_MEASUREMENT_DETECTOR.detect(ctx) is None + + +def test_text_detector_positive_string_column(): + ctx = _fake_ctx(column_type=T.StringType(), metrics={"count_non_null": 5}) + result = DEFAULT_TEXT_DETECTOR.detect(ctx) + assert result is not None + assert result.name == "text" + + +def test_text_detector_rejects_numeric_column(): + ctx = _fake_ctx(column_type=T.IntegerType(), metrics={"count_non_null": 5}) + assert DEFAULT_TEXT_DETECTOR.detect(ctx) is None + + +def test_key_detector_numeric_dense_positive(): + """user_id-like: distinctness=1.0, density=1.0.""" + ctx = _fake_ctx( + column_type=T.LongType(), + metrics={"count_non_null": 100, "count_distinct": 100, "min": 1, "max": 100}, + ) + result = DEFAULT_KEY_DETECTOR.detect(ctx) + assert result is not None + assert result.name == "key" + assert result.properties.get("signal") == "density" + + +def test_key_detector_numeric_sparse_negative_falls_through(): + """cargo_weight-like: distinctness high but density << 0.99.""" + ctx = _fake_ctx( + column_type=T.DoubleType(), + metrics={"count_non_null": 100, "count_distinct": 100, "min": 0, "max": 1_000_000}, + ) + assert DEFAULT_KEY_DETECTOR.detect(ctx) is None + + +def test_key_detector_low_distinctness_returns_none(): + ctx = _fake_ctx( + column_type=T.IntegerType(), + metrics={"count_non_null": 100, "count_distinct": 50, "min": 1, "max": 50}, + ) + assert DEFAULT_KEY_DETECTOR.detect(ctx) is None + + +# --------------------------------------------------------------------------- +# Enum applicability guard +# --------------------------------------------------------------------------- + + +def _mock_df_with_distinct(column_name: str, distinct_values: list) -> DataFrame: + distinct_df = create_autospec(DataFrame) + distinct_df.collect.return_value = [[v] for v in distinct_values] + distinct_df.distinct.return_value = distinct_df + + df = create_autospec(DataFrame) + df.columns = [column_name] + df.select.return_value = distinct_df + return df + + +def test_enum_detector_rejects_all_distinct_integer_column(): + """A 20-row all-distinct integer column must not be classified as enum.""" + ctx = DQProfileContext( + df=_mock_df_with_distinct("c", list(range(20))), + column_name="c", + column_type=T.IntegerType(), + metrics={"count_non_null": 20, "count_distinct": 20}, + options={"max_in_count": 10}, + ) + assert DEFAULT_ENUM_DETECTOR.detect(ctx) is None + + +def test_enum_detector_positive_low_cardinality(): + ctx = DQProfileContext( + df=_mock_df_with_distinct("c", ["car", "truck", "van"]), + column_name="c", + column_type=T.StringType(), + metrics={"count_non_null": 300, "count_distinct": 3}, + options={"max_in_count": 10}, + ) + result = DEFAULT_ENUM_DETECTOR.detect(ctx) + assert result is not None + assert result.name == "enum" + assert isinstance(result.properties["values"], tuple) + assert set(result.properties["values"]) == {"car", "truck", "van"} + + +# --------------------------------------------------------------------------- +# String-key length stability guard (uses Spark aggregate via autospec) +# --------------------------------------------------------------------------- + + +def _string_key_ctx(column_name, count_non_null, count_distinct, min_len, max_len): + agg_row = {"min_len": min_len, "max_len": max_len} + first_row = create_autospec(dict) + first_row.__getitem__.side_effect = agg_row.__getitem__ + + projected = create_autospec(DataFrame) + projected.first.return_value = first_row + + df = create_autospec(DataFrame) + df.columns = [column_name] + df.select.return_value = projected + + return DQProfileContext( + df=df, + column_name=column_name, + column_type=T.StringType(), + metrics={"count_non_null": count_non_null, "count_distinct": count_distinct}, + options={}, + ) + + +def test_key_detector_string_uniform_length_positive(): + """UUID-like: every value 36 chars → stability = 1.0 passes.""" + ctx = _string_key_ctx("order_id", count_non_null=50, count_distinct=50, min_len=36, max_len=36) + result = DEFAULT_KEY_DETECTOR.detect(ctx) + assert result is not None + assert result.name == "key" + assert result.properties.get("signal") == "length_stability" + + +def test_key_detector_string_variable_length_negative(): + """Free-form names, min 4 / max 40 → stability=0.1 well below threshold.""" + ctx = _string_key_ctx("user_name", count_non_null=50, count_distinct=50, min_len=4, max_len=40) + assert DEFAULT_KEY_DETECTOR.detect(ctx) is None + + +def test_key_detector_string_empty_only_returns_none_without_error(): + """max_length == 0 must not trigger ZeroDivisionError; column is not classified as key.""" + ctx = _string_key_ctx("empty_col", count_non_null=5, count_distinct=5, min_len=0, max_len=0) + assert DEFAULT_KEY_DETECTOR.detect(ctx) is None diff --git a/tests/unit/test_profile_builder.py b/tests/unit/test_profile_builder.py index 22f2aacd0..2fd02232c 100644 --- a/tests/unit/test_profile_builder.py +++ b/tests/unit/test_profile_builder.py @@ -2,11 +2,12 @@ from unittest.mock import create_autospec import pytest +from pydantic import ValidationError import pyspark.sql.types as T from pyspark.sql import DataFrame from databricks.labs.dqx.errors import InvalidParameterError -from databricks.labs.dqx.profiler.profile import DQProfile +from databricks.labs.dqx.profiler.profile import DQProfile, DQProfileBuilder from databricks.labs.dqx.profiler.profile_builder import ( PROFILE_BUILDER_REGISTRY, make_has_no_outliers_profile, @@ -16,6 +17,7 @@ register_profile_builder, validate_profile_options, ) +from databricks.labs.dqx.profiler.semantic import DQProfileContext, DQSemanticType @pytest.fixture @@ -24,6 +26,17 @@ def mock_df(): return df +def _ctx(df, column_name, column_type, metrics, options, semantic_type=None): + return DQProfileContext( + df=df, + column_name=column_name, + column_type=column_type, + metrics=metrics, + options=options, + semantic_type=semantic_type, + ) + + # --------------------------------------------------------------------------- # Registry # --------------------------------------------------------------------------- @@ -45,8 +58,10 @@ def _custom_builder(*_): try: assert "_test_custom" in PROFILE_BUILDER_REGISTRY - assert PROFILE_BUILDER_REGISTRY["_test_custom"].builder is _custom_builder - assert PROFILE_BUILDER_REGISTRY["_test_custom"].builder(None, "", None, {}, {}) is sentinel + entry = PROFILE_BUILDER_REGISTRY["_test_custom"] + assert entry.builder is _custom_builder + assert entry.contextual_builder is None + assert entry.builder(None, "", None, {}, {}) is sentinel finally: PROFILE_BUILDER_REGISTRY.pop("_test_custom", None) @@ -81,19 +96,61 @@ def _my_builder(*_): PROFILE_BUILDER_REGISTRY.pop("_test_return", None) +def test_register_profile_builder_context_type_uses_contextual_slot(): + @register_profile_builder("_test_ctx", type="context") + def _ctx_builder(ctx): + return None + + try: + entry = PROFILE_BUILDER_REGISTRY["_test_ctx"] + assert entry.contextual_builder is _ctx_builder + assert entry.builder is None + finally: + PROFILE_BUILDER_REGISTRY.pop("_test_ctx", None) + + +def test_register_profile_builder_legacy_type_uses_builder_slot(): + @register_profile_builder("_test_legacy_kw", type="legacy") + def _legacy_builder(*_): + return None + + try: + entry = PROFILE_BUILDER_REGISTRY["_test_legacy_kw"] + assert entry.builder is _legacy_builder + assert entry.contextual_builder is None + finally: + PROFILE_BUILDER_REGISTRY.pop("_test_legacy_kw", None) + + +def test_dq_profile_builder_rejects_both_callbacks(): + with pytest.raises(ValidationError): + DQProfileBuilder( + name="both", + builder=lambda *_: None, + contextual_builder=lambda _ctx: None, + ) + + +def test_dq_profile_builder_rejects_neither_callback(): + with pytest.raises(ValidationError): + DQProfileBuilder(name="neither") + + # --------------------------------------------------------------------------- -# make_null_or_empty_profile — text types +# make_null_or_empty_profile — text types (contextual dispatch) # --------------------------------------------------------------------------- @pytest.mark.parametrize("column_type", [T.StringType(), T.CharType(10), T.VarcharType(50)]) def test_null_or_empty_text_no_nulls_no_empties_returns_not_null_or_empty(mock_df, column_type): profile = make_null_or_empty_profile( - mock_df, - "col", - column_type, - {"count_null": 0, "empty_count": 0, "count": 10}, - {"max_null_ratio": 0.0, "max_empty_ratio": 0.0}, + _ctx( + mock_df, + "col", + column_type, + {"count_null": 0, "empty_count": 0, "count": 10}, + {"max_null_ratio": 0.0, "max_empty_ratio": 0.0}, + ) ) assert profile == DQProfile( name="is_not_null_or_empty", column="col", description=None, parameters={"trim_strings": True}, filter=None @@ -102,11 +159,13 @@ def test_null_or_empty_text_no_nulls_no_empties_returns_not_null_or_empty(mock_d def test_null_or_empty_text_nulls_and_empties_within_threshold_has_description(mock_df): profile = make_null_or_empty_profile( - mock_df, - "col", - T.StringType(), - {"count_null": 1, "empty_count": 1, "count": 10}, - {"max_null_ratio": 0.2, "max_empty_ratio": 0.2}, + _ctx( + mock_df, + "col", + T.StringType(), + {"count_null": 1, "empty_count": 1, "count": 10}, + {"max_null_ratio": 0.2, "max_empty_ratio": 0.2}, + ) ) assert profile is not None assert profile.name == "is_not_null_or_empty" @@ -116,11 +175,13 @@ def test_null_or_empty_text_nulls_and_empties_within_threshold_has_description(m def test_null_or_empty_text_nulls_exceed_threshold_empty_ok_returns_is_not_empty(mock_df): profile = make_null_or_empty_profile( - mock_df, - "col", - T.StringType(), - {"count_null": 5, "empty_count": 0, "count": 10}, - {"max_null_ratio": 0.3, "max_empty_ratio": 0.0}, + _ctx( + mock_df, + "col", + T.StringType(), + {"count_null": 5, "empty_count": 0, "count": 10}, + {"max_null_ratio": 0.3, "max_empty_ratio": 0.0}, + ) ) assert profile is not None assert profile.name == "is_not_empty" @@ -129,11 +190,13 @@ def test_null_or_empty_text_nulls_exceed_threshold_empty_ok_returns_is_not_empty def test_null_or_empty_text_empties_exceed_threshold_null_ok_returns_is_not_null(mock_df): profile = make_null_or_empty_profile( - mock_df, - "col", - T.StringType(), - {"count_null": 0, "empty_count": 5, "count": 10}, - {"max_null_ratio": 0.0, "max_empty_ratio": 0.3}, + _ctx( + mock_df, + "col", + T.StringType(), + {"count_null": 0, "empty_count": 5, "count": 10}, + {"max_null_ratio": 0.0, "max_empty_ratio": 0.3}, + ) ) assert profile is not None assert profile.name == "is_not_null" @@ -145,22 +208,26 @@ def test_null_or_empty_text_empties_exceed_threshold_null_ok_returns_is_not_null def test_null_or_empty_text_both_exceed_threshold_returns_none(mock_df): profile = make_null_or_empty_profile( - mock_df, - "col", - T.StringType(), - {"count_null": 5, "empty_count": 4, "count": 10}, - {"max_null_ratio": 0.3, "max_empty_ratio": 0.3}, + _ctx( + mock_df, + "col", + T.StringType(), + {"count_null": 5, "empty_count": 4, "count": 10}, + {"max_null_ratio": 0.3, "max_empty_ratio": 0.3}, + ) ) assert profile is None def test_null_or_empty_text_trim_strings_false_propagated(mock_df): profile = make_null_or_empty_profile( - mock_df, - "col", - T.StringType(), - {"count_null": 0, "empty_count": 0, "count": 5}, - {"max_null_ratio": 0.0, "max_empty_ratio": 0.0, "trim_strings": False}, + _ctx( + mock_df, + "col", + T.StringType(), + {"count_null": 0, "empty_count": 0, "count": 5}, + {"max_null_ratio": 0.0, "max_empty_ratio": 0.0, "trim_strings": False}, + ) ) assert profile is not None assert profile.parameters == {"trim_strings": False} @@ -168,11 +235,13 @@ def test_null_or_empty_text_trim_strings_false_propagated(mock_df): def test_null_or_empty_text_filter_propagated(mock_df): profile = make_null_or_empty_profile( - mock_df, - "col", - T.StringType(), - {"count_null": 0, "empty_count": 0, "count": 5}, - {"max_null_ratio": 0.0, "max_empty_ratio": 0.0, "filter": "x > 0"}, + _ctx( + mock_df, + "col", + T.StringType(), + {"count_null": 0, "empty_count": 0, "count": 5}, + {"max_null_ratio": 0.0, "max_empty_ratio": 0.0, "filter": "x > 0"}, + ) ) assert profile is not None assert profile.filter == "x > 0" @@ -180,11 +249,7 @@ def test_null_or_empty_text_filter_propagated(mock_df): def test_null_or_empty_text_empty_dataframe_returns_none(mock_df): profile = make_null_or_empty_profile( - mock_df, - "col", - T.StringType(), - {"count_null": 0, "empty_count": 0, "count": 0}, - {}, + _ctx(mock_df, "col", T.StringType(), {"count_null": 0, "empty_count": 0, "count": 0}, {}) ) assert profile is None @@ -196,22 +261,14 @@ def test_null_or_empty_text_empty_dataframe_returns_none(mock_df): def test_null_or_empty_non_text_no_nulls_returns_is_not_null(mock_df): profile = make_null_or_empty_profile( - mock_df, - "age", - T.IntegerType(), - {"count_null": 0, "count": 10}, - {"max_null_ratio": 0.0}, + _ctx(mock_df, "age", T.IntegerType(), {"count_null": 0, "count": 10}, {"max_null_ratio": 0.0}) ) assert profile == DQProfile(name="is_not_null", column="age", description=None, parameters=None, filter=None) def test_null_or_empty_non_text_nulls_within_threshold_has_description(mock_df): profile = make_null_or_empty_profile( - mock_df, - "age", - T.IntegerType(), - {"count_null": 1, "count": 10}, - {"max_null_ratio": 0.2}, + _ctx(mock_df, "age", T.IntegerType(), {"count_null": 1, "count": 10}, {"max_null_ratio": 0.2}) ) assert profile is not None assert profile.name == "is_not_null" @@ -221,11 +278,7 @@ def test_null_or_empty_non_text_nulls_within_threshold_has_description(mock_df): def test_null_or_empty_non_text_nulls_exceed_threshold_returns_none(mock_df): profile = make_null_or_empty_profile( - mock_df, - "age", - T.IntegerType(), - {"count_null": 5, "count": 10}, - {"max_null_ratio": 0.3}, + _ctx(mock_df, "age", T.IntegerType(), {"count_null": 5, "count": 10}, {"max_null_ratio": 0.3}) ) assert profile is None @@ -233,11 +286,7 @@ def test_null_or_empty_non_text_nulls_exceed_threshold_returns_none(mock_df): @pytest.mark.parametrize("column_type", [T.LongType(), T.DoubleType(), T.DateType(), T.BooleanType()]) def test_null_or_empty_non_text_types_no_nulls_return_is_not_null(mock_df, column_type): profile = make_null_or_empty_profile( - mock_df, - "col", - column_type, - {"count_null": 0, "count": 5}, - {"max_null_ratio": 0.0}, + _ctx(mock_df, "col", column_type, {"count_null": 0, "count": 5}, {"max_null_ratio": 0.0}) ) assert profile is not None assert profile.name == "is_not_null" @@ -263,37 +312,36 @@ def _make_mock_df(columns: list, distinct_values: list) -> DataFrame: @pytest.mark.parametrize("column_type", [T.DoubleType(), T.FloatType(), T.BooleanType(), T.DateType()]) def test_is_in_unsupported_type_returns_none(mock_df, column_type): - assert make_is_in_profile(mock_df, "col", column_type, {"count": 10}, {}) is None + assert make_is_in_profile(_ctx(mock_df, "col", column_type, {"count": 10}, {})) is None @pytest.mark.parametrize("column_type", [T.CharType(10), T.VarcharType(50)]) def test_is_in_char_varchar_type_returns_profile(column_type): df = _make_mock_df(["col"], ["a", "b", "c"]) - profile = make_is_in_profile(df, "col", column_type, {"count": 10}, {"max_in_count": 10, "distinct_ratio": 1.0}) + profile = make_is_in_profile( + _ctx(df, "col", column_type, {"count": 10}, {"max_in_count": 10, "distinct_ratio": 1.0}) + ) assert profile is not None assert profile.name == "is_in" assert set(profile.parameters["in"]) == {"a", "b", "c"} def test_is_in_total_count_zero_returns_none(mock_df): - assert make_is_in_profile(mock_df, "col", T.IntegerType(), {"count": 0}, {}) is None + assert make_is_in_profile(_ctx(mock_df, "col", T.IntegerType(), {"count": 0}, {})) is None def test_is_in_no_distinct_values_returns_none(): df = _make_mock_df(["col"], []) assert ( - make_is_in_profile(df, "col", T.StringType(), {"count": 3}, {"max_in_count": 10, "distinct_ratio": 1.0}) is None + make_is_in_profile(_ctx(df, "col", T.StringType(), {"count": 3}, {"max_in_count": 10, "distinct_ratio": 1.0})) + is None ) def test_is_in_conditions_met_returns_profile(): df = _make_mock_df(["col"], [1, 2, 3]) profile = make_is_in_profile( - df, - "status", - T.IntegerType(), - {"count": 5}, - {"max_in_count": 10, "distinct_ratio": 1.0}, + _ctx(df, "status", T.IntegerType(), {"count": 5}, {"max_in_count": 10, "distinct_ratio": 1.0}) ) assert profile is not None assert profile.name == "is_in" @@ -305,11 +353,7 @@ def test_is_in_distinct_count_exceeds_max_in_count_returns_none(): # 11 distinct values, max_in_count=10 → distinct_count > max_in_count → None df = _make_mock_df(["col"], list(range(11))) profile = make_is_in_profile( - df, - "col", - T.IntegerType(), - {"count": 100}, - {"max_in_count": 10, "distinct_ratio": 1.0}, + _ctx(df, "col", T.IntegerType(), {"count": 100}, {"max_in_count": 10, "distinct_ratio": 1.0}) ) assert profile is None @@ -318,11 +362,7 @@ def test_is_in_distinct_ratio_exceeds_threshold_returns_none(): # 10 distinct values in 10 total → ratio=1.0, threshold=0.5 df = _make_mock_df(["col"], list(range(10))) profile = make_is_in_profile( - df, - "col", - T.StringType(), - {"count": 10}, - {"max_in_count": 20, "distinct_ratio": 0.5}, + _ctx(df, "col", T.StringType(), {"count": 10}, {"max_in_count": 20, "distinct_ratio": 0.5}) ) assert profile is None @@ -330,37 +370,69 @@ def test_is_in_distinct_ratio_exceeds_threshold_returns_none(): def test_is_in_filter_propagated(): df = _make_mock_df(["col"], ["a", "b"]) profile = make_is_in_profile( - df, - "col", - T.StringType(), - {"count": 5}, - {"max_in_count": 10, "distinct_ratio": 1.0, "filter": "x > 0"}, + _ctx(df, "col", T.StringType(), {"count": 5}, {"max_in_count": 10, "distinct_ratio": 1.0, "filter": "x > 0"}) ) assert profile is not None assert profile.filter == "x > 0" +def test_is_in_reuses_enum_values_without_extra_spark_action(mock_df): + """When the enum detector already collected distinct values, the builder must reuse them.""" + mock_df.columns = ["vehicle_type"] + profile = make_is_in_profile( + _ctx( + mock_df, + "vehicle_type", + T.StringType(), + {"count": 100, "count_non_null": 100}, + {"max_in_count": 10, "distinct_ratio": 0.1}, + semantic_type=DQSemanticType(name="enum", properties={"values": ("car", "truck", "van")}), + ) + ) + assert profile is not None + assert profile.name == "is_in" + assert profile.parameters == {"in": ["car", "truck", "van"]} + # No .distinct().collect() should have been called since the values came from the detector. + mock_df.select.assert_not_called() + + +@pytest.mark.parametrize("other_type", ["key", "measurement", "text", "custom_type"]) +def test_is_in_skipped_for_non_enum_semantic_type(mock_df, other_type): + profile = make_is_in_profile( + _ctx( + mock_df, + "col", + T.StringType(), + {"count": 100}, + {"max_in_count": 10, "distinct_ratio": 0.1}, + semantic_type=DQSemanticType(name=other_type), + ) + ) + assert profile is None + + # --------------------------------------------------------------------------- # make_min_max_profile # --------------------------------------------------------------------------- def test_min_max_count_non_null_zero_returns_none(mock_df): - assert make_min_max_profile(mock_df, "col", T.IntegerType(), {"count_non_null": 0}, {}) is None + assert make_min_max_profile(_ctx(mock_df, "col", T.IntegerType(), {"count_non_null": 0}, {})) is None @pytest.mark.parametrize("column_type", [T.StringType(), T.BooleanType(), T.ByteType()]) def test_min_max_unsupported_type_returns_none(mock_df, column_type): - assert make_min_max_profile(mock_df, "col", column_type, {"count_non_null": 5}, {"remove_outliers": False}) is None + assert ( + make_min_max_profile(_ctx(mock_df, "col", column_type, {"count_non_null": 5}, {"remove_outliers": False})) + is None + ) def test_min_max_without_outlier_removal_uses_metrics(mock_df): profile = make_min_max_profile( - mock_df, - "amount", - T.IntegerType(), - {"count_non_null": 5, "min": 1, "max": 100}, - {"remove_outliers": False}, + _ctx( + mock_df, "amount", T.IntegerType(), {"count_non_null": 5, "min": 1, "max": 100}, {"remove_outliers": False} + ) ) assert profile is not None assert profile.name == "min_max" @@ -371,11 +443,9 @@ def test_min_max_without_outlier_removal_uses_metrics(mock_df): def test_min_max_without_outlier_removal_double_type(mock_df): profile = make_min_max_profile( - mock_df, - "score", - T.DoubleType(), - {"count_non_null": 10, "min": 0.5, "max": 9.9}, - {"remove_outliers": False}, + _ctx( + mock_df, "score", T.DoubleType(), {"count_non_null": 10, "min": 0.5, "max": 9.9}, {"remove_outliers": False} + ) ) assert profile is not None assert profile.parameters == {"min": 0.5, "max": 9.9} @@ -383,11 +453,13 @@ def test_min_max_without_outlier_removal_double_type(mock_df): def test_min_max_filter_propagated(mock_df): profile = make_min_max_profile( - mock_df, - "col", - T.IntegerType(), - {"count_non_null": 5, "min": 1, "max": 10}, - {"remove_outliers": False, "filter": "x > 0"}, + _ctx( + mock_df, + "col", + T.IntegerType(), + {"count_non_null": 5, "min": 1, "max": 10}, + {"remove_outliers": False, "filter": "x > 0"}, + ) ) assert profile is not None assert profile.filter == "x > 0" @@ -399,11 +471,7 @@ def test_min_max_filter_propagated(mock_df): ) def test_min_max_supported_numeric_types_return_profile(mock_df, column_type): profile = make_min_max_profile( - mock_df, - "col", - column_type, - {"count_non_null": 5, "min": 1, "max": 10}, - {"remove_outliers": False}, + _ctx(mock_df, "col", column_type, {"count_non_null": 5, "min": 1, "max": 10}, {"remove_outliers": False}) ) assert profile is not None assert profile.name == "min_max" @@ -413,11 +481,13 @@ def test_min_max_with_outlier_removal_stddev_zero_returns_real_min_max(mock_df): # stddev=0 means all values are identical; sigma bounds collapse to mean. # None of the sigma-capping branches fire, so real min/max are used. profile = make_min_max_profile( - mock_df, - "amount", - T.IntegerType(), - {"count_non_null": 10, "min": 5, "max": 5, "mean": 5.0, "stddev": 0.0}, - {"remove_outliers": True, "outlier_columns": ["amount"]}, + _ctx( + mock_df, + "amount", + T.IntegerType(), + {"count_non_null": 10, "min": 5, "max": 5, "mean": 5.0, "stddev": 0.0}, + {"remove_outliers": True, "outlier_columns": ["amount"]}, + ) ) assert profile is not None assert profile.name == "min_max" @@ -429,11 +499,13 @@ def test_min_max_empty_outlier_columns_applies_outlier_removal_to_all_columns(mo # empty outlier_columns with remove_outliers=True must apply to all columns (regression test for issue #1) # mean=50, stddev=10, sigmas=3 → bounds [20, 80] which cap the real range [1, 100] profile = make_min_max_profile( - mock_df, - "amount", - T.IntegerType(), - {"count_non_null": 10, "min": 1, "max": 100, "mean": 50.0, "stddev": 10.0}, - {"remove_outliers": True, "outlier_columns": []}, + _ctx( + mock_df, + "amount", + T.IntegerType(), + {"count_non_null": 10, "min": 1, "max": 100, "mean": 50.0, "stddev": 10.0}, + {"remove_outliers": True, "outlier_columns": []}, + ) ) assert profile is not None assert profile.parameters == {"min": 20, "max": 80} @@ -443,11 +515,13 @@ def test_min_max_empty_outlier_columns_applies_outlier_removal_to_all_columns(mo def test_min_max_column_not_in_outlier_columns_skips_outlier_removal(mock_df): # when outlier_columns is set but does not include this column, use real min/max profile = make_min_max_profile( - mock_df, - "amount", - T.IntegerType(), - {"count_non_null": 10, "min": 1, "max": 100, "mean": 50.0, "stddev": 10.0}, - {"remove_outliers": True, "outlier_columns": ["other_col"]}, + _ctx( + mock_df, + "amount", + T.IntegerType(), + {"count_non_null": 10, "min": 1, "max": 100, "mean": 50.0, "stddev": 10.0}, + {"remove_outliers": True, "outlier_columns": ["other_col"]}, + ) ) assert profile is not None assert profile.parameters == {"min": 1, "max": 100} @@ -458,11 +532,13 @@ def test_min_max_rounding_zero_min_is_not_skipped(mock_df): # regression: falsy check `if not value` would skip rounding when min=0.0, # leaving a float instead of the expected int. Fixed by `if value is None`. profile = make_min_max_profile( - mock_df, - "amount", - T.IntegerType(), - {"count_non_null": 5, "min": 0, "max": 10}, - {"remove_outliers": False, "round": True}, + _ctx( + mock_df, + "amount", + T.IntegerType(), + {"count_non_null": 5, "min": 0, "max": 10}, + {"remove_outliers": False, "round": True}, + ) ) assert profile is not None assert profile.parameters["min"] == 0 @@ -471,11 +547,13 @@ def test_min_max_rounding_zero_min_is_not_skipped(mock_df): def test_min_max_rounding_disabled_returns_float_as_is(mock_df): profile = make_min_max_profile( - mock_df, - "amount", - T.DoubleType(), - {"count_non_null": 5, "min": 1.2, "max": 9.9}, - {"remove_outliers": False, "round": False}, + _ctx( + mock_df, + "amount", + T.DoubleType(), + {"count_non_null": 5, "min": 1.2, "max": 9.9}, + {"remove_outliers": False, "round": False}, + ) ) assert profile is not None assert profile.parameters["min"] == 1.2 @@ -486,11 +564,13 @@ def test_min_max_rounding_enabled_floors_float_min_and_ceils_float_max(mock_df): # regression: when min/max came from summary-stats metrics (fast path), round=True was # silently ignored for float types. Values must be floor/ceil'd just as the Spark fallback does. profile = make_min_max_profile( - mock_df, - "price", - T.DoubleType(), - {"count_non_null": 5, "min": 1.2, "max": 9.9}, - {"remove_outliers": False, "round": True}, + _ctx( + mock_df, + "price", + T.DoubleType(), + {"count_non_null": 5, "min": 1.2, "max": 9.9}, + {"remove_outliers": False, "round": True}, + ) ) assert profile is not None assert profile.parameters["min"] == 1.0 @@ -499,17 +579,49 @@ def test_min_max_rounding_enabled_floors_float_min_and_ceils_float_max(mock_df): def test_min_max_rounding_enabled_for_decimal_type(mock_df): profile = make_min_max_profile( - mock_df, - "amount", - T.DecimalType(10, 2), - {"count_non_null": 5, "min": decimal.Decimal("1.20"), "max": decimal.Decimal("9.90")}, - {"remove_outliers": False, "round": True}, + _ctx( + mock_df, + "amount", + T.DecimalType(10, 2), + {"count_non_null": 5, "min": decimal.Decimal("1.20"), "max": decimal.Decimal("9.90")}, + {"remove_outliers": False, "round": True}, + ) ) assert profile is not None assert profile.parameters["min"] == decimal.Decimal("1") assert profile.parameters["max"] == decimal.Decimal("10") +@pytest.mark.parametrize("other_type", ["enum", "key", "text", "custom_type"]) +def test_min_max_skipped_for_non_measurement_semantic_type(mock_df, other_type): + profile = make_min_max_profile( + _ctx( + mock_df, + "col", + T.IntegerType(), + {"count_non_null": 5, "min": 1, "max": 10}, + {"remove_outliers": False}, + semantic_type=DQSemanticType(name=other_type), + ) + ) + assert profile is None + + +def test_min_max_emitted_when_semantic_type_is_measurement(mock_df): + profile = make_min_max_profile( + _ctx( + mock_df, + "col", + T.IntegerType(), + {"count_non_null": 5, "min": 1, "max": 10}, + {"remove_outliers": False}, + semantic_type=DQSemanticType(name="measurement"), + ) + ) + assert profile is not None + assert profile.name == "min_max" + + # --------------------------------------------------------------------------- # make_has_no_outliers_profile # --------------------------------------------------------------------------- @@ -518,14 +630,14 @@ def test_min_max_rounding_enabled_for_decimal_type(mock_df): @pytest.mark.parametrize("column_type", [T.StringType(), T.BooleanType(), T.DateType(), T.TimestampType()]) def test_has_no_outliers_non_numeric_type_returns_none(mock_df, column_type): profile = make_has_no_outliers_profile( - mock_df, "col", column_type, {"count_non_null": 10}, {"outliers_ratio": 0.01} + _ctx(mock_df, "col", column_type, {"count_non_null": 10}, {"outliers_ratio": 0.01}) ) assert profile is None def test_has_no_outliers_count_non_null_zero_returns_none(mock_df): profile = make_has_no_outliers_profile( - mock_df, "col", T.IntegerType(), {"count_non_null": 0}, {"outliers_ratio": 0.01} + _ctx(mock_df, "col", T.IntegerType(), {"count_non_null": 0}, {"outliers_ratio": 0.01}) ) assert profile is None @@ -537,7 +649,7 @@ def test_has_no_outliers_numeric_no_outliers_returns_profile(mock_df): mock_df.filter.return_value.count.return_value = 0 profile = make_has_no_outliers_profile( - mock_df, "measurement", T.IntegerType(), {"count_non_null": 10}, {"outliers_ratio": 0.05} + _ctx(mock_df, "measurement", T.IntegerType(), {"count_non_null": 10}, {"outliers_ratio": 0.05}) ) assert profile is not None @@ -556,7 +668,7 @@ def test_has_no_outliers_outliers_exceed_threshold_returns_none(mock_df): mock_df.filter.return_value.count.return_value = 2 profile = make_has_no_outliers_profile( - mock_df, "col", T.IntegerType(), {"count_non_null": 4}, {"outliers_ratio": 0.1} + _ctx(mock_df, "col", T.IntegerType(), {"count_non_null": 4}, {"outliers_ratio": 0.1}) ) assert profile is None @@ -568,7 +680,7 @@ def test_has_no_outliers_bounds_none_returns_none(mock_df): mock_df.select.return_value.agg.return_value.collect.return_value = [[None]] profile = make_has_no_outliers_profile( - mock_df, "col", T.IntegerType(), {"count_non_null": 5}, {"outliers_ratio": 0.01} + _ctx(mock_df, "col", T.IntegerType(), {"count_non_null": 5}, {"outliers_ratio": 0.01}) ) assert profile is None @@ -581,7 +693,7 @@ def test_has_no_outliers_bounds_none_returns_none(mock_df): def test_has_no_outliers_disabled_via_option_returns_none(mock_df): profile = make_has_no_outliers_profile( - mock_df, "col", T.IntegerType(), {"count_non_null": 10}, {"has_no_outliers": False} + _ctx(mock_df, "col", T.IntegerType(), {"count_non_null": 10}, {"has_no_outliers": False}) ) assert profile is None @@ -592,11 +704,13 @@ def test_has_no_outliers_column_in_allow_columns_returns_profile(mock_df): mock_df.filter.return_value.count.return_value = 0 profile = make_has_no_outliers_profile( - mock_df, - "measurement", - T.IntegerType(), - {"count_non_null": 10}, - {"has_no_outliers_allow_columns": ["measurement"], "outliers_ratio": 0.05}, + _ctx( + mock_df, + "measurement", + T.IntegerType(), + {"count_non_null": 10}, + {"has_no_outliers_allow_columns": ["measurement"], "outliers_ratio": 0.05}, + ) ) assert profile is not None @@ -606,22 +720,26 @@ def test_has_no_outliers_column_in_allow_columns_returns_profile(mock_df): def test_has_no_outliers_column_not_in_allow_columns_returns_none(mock_df): profile = make_has_no_outliers_profile( - mock_df, - "other_col", - T.IntegerType(), - {"count_non_null": 10}, - {"has_no_outliers_allow_columns": ["measurement"], "outliers_ratio": 0.05}, + _ctx( + mock_df, + "other_col", + T.IntegerType(), + {"count_non_null": 10}, + {"has_no_outliers_allow_columns": ["measurement"], "outliers_ratio": 0.05}, + ) ) assert profile is None def test_has_no_outliers_column_in_deny_columns_returns_none(mock_df): profile = make_has_no_outliers_profile( - mock_df, - "measurement", - T.IntegerType(), - {"count_non_null": 10}, - {"has_no_outliers_deny_columns": ["measurement"], "outliers_ratio": 0.05}, + _ctx( + mock_df, + "measurement", + T.IntegerType(), + {"count_non_null": 10}, + {"has_no_outliers_deny_columns": ["measurement"], "outliers_ratio": 0.05}, + ) ) assert profile is None @@ -632,11 +750,13 @@ def test_has_no_outliers_column_not_in_deny_columns_returns_profile(mock_df): mock_df.filter.return_value.count.return_value = 0 profile = make_has_no_outliers_profile( - mock_df, - "measurement", - T.IntegerType(), - {"count_non_null": 10}, - {"has_no_outliers_deny_columns": ["other_col"], "outliers_ratio": 0.05}, + _ctx( + mock_df, + "measurement", + T.IntegerType(), + {"count_non_null": 10}, + {"has_no_outliers_deny_columns": ["other_col"], "outliers_ratio": 0.05}, + ) ) assert profile is not None @@ -647,14 +767,13 @@ def test_has_no_outliers_column_not_in_deny_columns_returns_profile(mock_df): def test_has_no_outliers_both_allow_and_deny_columns_raises(mock_df): with pytest.raises(InvalidParameterError): make_has_no_outliers_profile( - mock_df, - "measurement", - T.IntegerType(), - {"count_non_null": 10}, - { - "has_no_outliers_allow_columns": ["measurement"], - "has_no_outliers_deny_columns": ["other_col"], - }, + _ctx( + mock_df, + "measurement", + T.IntegerType(), + {"count_non_null": 10}, + {"has_no_outliers_allow_columns": ["measurement"], "has_no_outliers_deny_columns": ["other_col"]}, + ) ) @@ -664,11 +783,7 @@ def test_has_no_outliers_filter_propagated(mock_df): mock_df.filter.return_value.count.return_value = 0 profile = make_has_no_outliers_profile( - mock_df, - "col", - T.IntegerType(), - {"count_non_null": 10}, - {"outliers_ratio": 0.05, "filter": "x > 0"}, + _ctx(mock_df, "col", T.IntegerType(), {"count_non_null": 10}, {"outliers_ratio": 0.05, "filter": "x > 0"}) ) assert profile is not None @@ -683,7 +798,7 @@ def test_has_no_outliers_ratio_equal_to_threshold_emits_profile(mock_df): mock_df.filter.return_value.count.return_value = 1 profile = make_has_no_outliers_profile( - mock_df, "measurement", T.IntegerType(), {"count_non_null": 10}, {"outliers_ratio": 0.1} + _ctx(mock_df, "measurement", T.IntegerType(), {"count_non_null": 10}, {"outliers_ratio": 0.1}) ) assert profile is not None @@ -697,7 +812,7 @@ def test_has_no_outliers_near_degenerate_mad_returns_none(mock_df): mock_df.select.return_value.agg.return_value.collect.return_value = [[1e-9]] profile = make_has_no_outliers_profile( - mock_df, "col", T.IntegerType(), {"count_non_null": 10}, {"outliers_ratio": 0.05} + _ctx(mock_df, "col", T.IntegerType(), {"count_non_null": 10}, {"outliers_ratio": 0.05}) ) assert profile is None From fa10e8bb62aa12d8edc165b721b26fd70d8f0601 Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 29 Aug 2026 14:46:31 +0200 Subject: [PATCH 2/9] linters fixed --- src/databricks/labs/dqx/profiler/common.py | 7 + src/databricks/labs/dqx/profiler/profile.py | 4 +- .../labs/dqx/profiler/profile_builder.py | 34 ++-- src/databricks/labs/dqx/profiler/profiler.py | 56 ++++-- src/databricks/labs/dqx/profiler/semantic.py | 178 +++++++++++------- .../profiler/test_semantic_profiling.py | 25 ++- tests/unit/profiler/test_semantic.py | 40 ++-- tests/unit/test_profile_builder.py | 10 +- 8 files changed, 219 insertions(+), 135 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/common.py b/src/databricks/labs/dqx/profiler/common.py index 07f72442c..3217a6271 100644 --- a/src/databricks/labs/dqx/profiler/common.py +++ b/src/databricks/labs/dqx/profiler/common.py @@ -3,6 +3,13 @@ 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 val_to_str(value: Any, include_sql_quotes: bool = True): """ diff --git a/src/databricks/labs/dqx/profiler/profile.py b/src/databricks/labs/dqx/profiler/profile.py index 0c9d7a264..adf78e8aa 100644 --- a/src/databricks/labs/dqx/profiler/profile.py +++ b/src/databricks/labs/dqx/profiler/profile.py @@ -35,7 +35,7 @@ class DQProfile: # Legacy 5-argument callback shape used by pre-semantic profile builders. Kept for backward # compatibility with user-authored builders registered via @register_profile_builder without -# type="context". Prefer *ContextualProfileBuilder* for new code. +# kind="context". Prefer *ContextualProfileBuilder* for new code. ProfileBuilder = Callable[ [DataFrame, str, DataType, dict[str, Any], dict[str, Any]], DQProfile | None, @@ -58,7 +58,7 @@ class DQProfileBuilder(BaseModel): builder: Legacy 5-argument callback. Left in place for backward compatibility so existing user-authored builders keep working when wrapped in *DQProfileBuilder* or registered via - *@register_profile_builder* without *type="context"*. Does not + *@register_profile_builder* without *kind="context"*. Does not receive semantic-type information — prefer *contextual_builder* for new code. contextual_builder: Preferred single-argument callback receiving a diff --git a/src/databricks/labs/dqx/profiler/profile_builder.py b/src/databricks/labs/dqx/profiler/profile_builder.py index d8c7309df..6d5ed769d 100644 --- a/src/databricks/labs/dqx/profiler/profile_builder.py +++ b/src/databricks/labs/dqx/profiler/profile_builder.py @@ -1,9 +1,9 @@ import datetime import decimal import logging -from collections.abc import Callable +from collections.abc import Callable, Mapping import math -from typing import Any, Literal, Mapping +from typing import Any, Literal from pyspark.sql import DataFrame from pyspark.sql import types as T, functions as F @@ -11,7 +11,8 @@ from databricks.labs.dqx.check_funcs import get_limit_expr from databricks.labs.dqx.errors import InvalidParameterError from databricks.labs.dqx.profiler.profile import DQProfile, DQProfileBuilder -from databricks.labs.dqx.profiler.semantic import DQProfileContext +from databricks.labs.dqx.profiler.common import TEXT_TYPES +from databricks.labs.dqx.profiler.semantic import DQProfileContext, EnumProperties from databricks.labs.dqx.profiling_utils import calculate_median_absolute_deviation_bounds from databricks.labs.dqx.profiler.profile_options import ( PROFILE_OPTION_DISTINCT_RATIO, @@ -31,10 +32,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). @@ -49,13 +46,13 @@ def register_profile_builder( profile_type: str, *, - type: Literal["legacy", "context"] | None = None, + kind: Literal["legacy", "context"] | None = None, ) -> Callable: """Register a profile builder in *PROFILE_BUILDER_REGISTRY*. Args: profile_type: Registry key (e.g. *min_max*). - type: Callback shape. + kind: Callback shape. * *None* (default) or *"legacy"* — builder is a 5-argument callback matching the *ProfileBuilder* type alias. Registered as *DQProfileBuilder(name=..., builder=fn)*. Backward-compatible path. @@ -65,7 +62,7 @@ def register_profile_builder( """ def wrapper(builder_func: Callable) -> Callable: - if type == "context": + if kind == "context": PROFILE_BUILDER_REGISTRY[profile_type] = DQProfileBuilder( name=profile_type, contextual_builder=builder_func ) @@ -76,7 +73,7 @@ def wrapper(builder_func: Callable) -> Callable: return wrapper -@register_profile_builder("null_or_empty", type="context") +@register_profile_builder("null_or_empty", kind="context") def make_null_or_empty_profile(ctx: DQProfileContext) -> DQProfile | None: """ Creates an *is_not_null_or_empty*, *is_not_null*, or *is_not_empty* profile by checking @@ -94,7 +91,7 @@ def make_null_or_empty_profile(ctx: DQProfileContext) -> DQProfile | None: return _make_null_profile(ctx.column_name, ctx.metrics, ctx.options) -@register_profile_builder("is_in", type="context") +@register_profile_builder("is_in", kind="context") def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: """ Creates an *is_in* profile. @@ -125,10 +122,11 @@ def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: return None if semantic_type is not None and semantic_type.name == "enum": - values = semantic_type.properties.get("values", ()) - distinct_values = list(values) if isinstance(values, tuple) else list(values) - if not distinct_values: + enum_props = semantic_type.get_typed_properties(EnumProperties) + if enum_props is None or not enum_props.values: return None + # Sort so the emitted rule is deterministic across runs — EnumProperties.values is a set. + distinct_values = sorted(enum_props.values) return DQProfile( name="is_in", column=ctx.column_name, @@ -160,7 +158,7 @@ def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: return None -@register_profile_builder("min_max", type="context") +@register_profile_builder("min_max", kind="context") def make_min_max_profile(ctx: DQProfileContext) -> DQProfile | None: """ Creates a *min_max* profile. @@ -389,7 +387,7 @@ def _is_profile_enabled( profile_enabled_option_name: str, profile_allow_columns_option_name: str, profile_deny_columns_option_name: str, - profiler_options: dict[str, Any], + profiler_options: Mapping[str, Any], ) -> bool: """ Checks if a profiler builder is enabled for the given column. @@ -800,7 +798,7 @@ def _round_decimal(value: decimal.Decimal, rounding_direction: str) -> decimal.D return value -@register_profile_builder("has_no_outliers", type="context") +@register_profile_builder("has_no_outliers", kind="context") def make_has_no_outliers_profile(ctx: DQProfileContext) -> DQProfile | None: """ Creates a *has_no_outliers* profile using the same MAD method as the *has_no_outliers* check rule. diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 90f876aec..31a285e06 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -14,15 +14,18 @@ from pyspark.errors import AnalysisException from pyspark.sql import DataFrame, SparkSession from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import DatabricksError from databricks.labs.dqx.base import DQEngineBase 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.profile import DQProfile -from databricks.labs.dqx.profiler.profile_builder import PROFILE_BUILDER_REGISTRY, TEXT_TYPES, validate_profile_options +from databricks.labs.dqx.profiler.common import TEXT_TYPES +from databricks.labs.dqx.profiler.profile_builder import PROFILE_BUILDER_REGISTRY, validate_profile_options from databricks.labs.dqx.profiler.semantic import ( DQProfileContext, + DQSemanticType, SemanticRegistry, ) from databricks.labs.dqx.profiler.profile_options import ( @@ -456,14 +459,15 @@ def _fetch_table_metadata(self, location: str) -> dict[str, Any]: """ try: table = self.ws.tables.get(location) - # TODO (IK): Specify exception - except Exception as exc: # noqa: BLE001 — SDK raises many concrete types; degrade gracefully + except DatabricksError as exc: safe_location = str(location).replace("\n", " ").replace("\r", " ") logger.warning(f"Could not fetch table metadata for {safe_location}: {exc}") return {} columns: dict[str, dict[str, Any]] = {} for col in table.columns or []: + if col.name is None: + continue col_entry: dict[str, Any] = {} if col.comment is not None: col_entry["column_comment"] = col.comment @@ -597,23 +601,7 @@ def _build_profiles_for_column( without triggering a second Spark action. """ metadata_map = dict(column_metadata) if column_metadata else {} - detector_ctx = DQProfileContext( - df=column_df, - column_name=field_name, - column_type=field_type, - metrics=metrics, - options=opts, - metadata=metadata_map, - semantic_type=None, - ) - - semantic_type = None - if self._semantic_registry is not None: - for detector in self._semantic_registry.detectors: - match = detector.detect(detector_ctx) - if match is not None: - semantic_type = match - break + semantic_type = self._detect_semantic_type(column_df, field_name, field_type, metrics, opts, metadata_map) # Reconstruct via the constructor (not model_copy) so validators run — same rationale as # SemanticRegistry.prepend/replace. @@ -632,6 +620,8 @@ def _build_profiles_for_column( profile = profile_type.contextual_builder(builder_ctx) elif profile_type.builder is not None: profile = profile_type.builder(column_df, field_name, field_type, dict(metrics), dict(opts)) + else: + continue if not profile: continue @@ -646,6 +636,32 @@ def _build_profiles_for_column( if profile.parameters.get("max") is not None: metrics["max"] = profile.parameters.get("max") + def _detect_semantic_type( + self, + column_df: DataFrame, + field_name: str, + field_type: T.DataType, + metrics: dict[str, Any], + opts: dict[str, Any], + metadata_map: dict[str, Any], + ) -> DQSemanticType | None: + if self._semantic_registry is None: + return None + detector_ctx = DQProfileContext( + df=column_df, + column_name=field_name, + column_type=field_type, + metrics=metrics, + options=opts, + metadata=metadata_map, + semantic_type=None, + ) + for detector in self._semantic_registry.detectors: + match = detector.detect(detector_ctx) + if match is not None: + return match + return None + def _add_llm_primary_key_for_dataframe( self, df: DataFrame, dq_rules: list[DQProfile], summary_stats: dict[str, Any], opts: dict[str, Any] ) -> None: diff --git a/src/databricks/labs/dqx/profiler/semantic.py b/src/databricks/labs/dqx/profiler/semantic.py index ea02922c0..676c9ec05 100644 --- a/src/databricks/labs/dqx/profiler/semantic.py +++ b/src/databricks/labs/dqx/profiler/semantic.py @@ -13,6 +13,8 @@ Public surface: * models: *DQSemanticType*, *DQProfileContext*, *DQSemanticTypeDetector* + * properties: *DQSemanticTypeProperties*, *EnumProperties*, *KeyProperties*, + *MeasurementProperties* * registry: *SemanticRegistry* * detectors: *DEFAULT_ENUM_DETECTOR*, *DEFAULT_KEY_DETECTOR*, *DEFAULT_MEASUREMENT_DETECTOR*, *DEFAULT_TEXT_DETECTOR*, @@ -23,7 +25,7 @@ import logging from collections.abc import Callable, Mapping, Sequence -from typing import Any +from typing import Any, Literal, TypeVar from pydantic import BaseModel, ConfigDict, Field, model_validator from pyspark.sql import DataFrame @@ -31,16 +33,63 @@ from pyspark.sql import types as T from pyspark.sql.types import DataType +from databricks.labs.dqx.profiler.common import TEXT_TYPES from databricks.labs.dqx.profiler.profile_options import PROFILE_OPTION_MAX_IN_COUNT logger = logging.getLogger(__name__) Scalar = str | int | float | bool -PropertyValue = Scalar | tuple[Scalar, ...] -# TODO (IK): Make properties `BaseModel` too. +class DQSemanticTypeProperties(BaseModel): + """Base class for detector-specific properties emitted alongside a *DQSemanticType*. + + Extension pattern: user-defined detectors subclass this and pass an instance + as the *properties* field of the *DQSemanticType* they emit. The model is + frozen and forbids extra fields, matching *DQSemanticType* itself. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + +class EnumProperties(DQSemanticTypeProperties): + """Properties for the built-in *enum* semantic type. + + Attributes: + values: Collected distinct values. Homogeneous by construction — the + enum detector only fires on a single scalar column type, so the + set carries a single scalar type per instance. + """ + + values: set[str] | set[int] + + +class KeyProperties(DQSemanticTypeProperties): + """Properties for the built-in *key* semantic type. + + Attributes: + signal: Which secondary signal fired alongside distinctness — *density* + for numeric keys, *length_stability* for string keys. + """ + + signal: Literal["density", "length_stability"] + + +class MeasurementProperties(DQSemanticTypeProperties): + """Properties for the built-in *measurement* semantic type. + + Attributes: + distribution: Best-effort distribution family guess. *unknown* when + statistics were insufficient to classify. + """ + + distribution: Literal["constant", "uniform", "normal", "exponential", "unknown"] = "unknown" + + +_PropertiesT = TypeVar("_PropertiesT", bound=DQSemanticTypeProperties) + + class DQSemanticType(BaseModel): """Classification of a column's semantic meaning. @@ -48,18 +97,30 @@ class DQSemanticType(BaseModel): name: Globally unique identifier (e.g. *key*, *enum*, *measurement*). Used in generated check metadata and for detector lookup. description: Optional human-readable description. - properties: Detector-specific properties (e.g. distribution family, - enum values). Sequence values must be *tuple* (the model's value - type is *Scalar | tuple[Scalar, ...]*), so individual entries - are immutable; the *frozen* model config blocks reassignment of - the *properties* field itself. + properties: Detector-specific properties as a *DQSemanticTypeProperties* + subclass instance, or *None* when the detector emits no properties. + The *frozen* model config blocks reassignment of the field itself; + each properties subclass is frozen too, so the returned instance + is fully immutable. """ model_config = ConfigDict(frozen=True, extra="forbid") name: str description: str | None = None - properties: Mapping[str, PropertyValue] = Field(default_factory=dict) + properties: DQSemanticTypeProperties | None = None + + def get_typed_properties(self, properties_type: type[_PropertiesT]) -> _PropertiesT | None: + """Return *self.properties* narrowed to *properties_type*, or *None* when the type does not match. + + Args: + properties_type: Concrete *DQSemanticTypeProperties* subclass to cast to. + + Returns: + The properties instance typed as *properties_type* when it is an + instance of that class, else *None*. + """ + return self.properties if isinstance(self.properties, properties_type) else None class DQProfileContext(BaseModel): @@ -130,27 +191,12 @@ class DQSemanticTypeDetector(BaseModel): KEY_MIN_LENGTH_STABILITY_RATIO = 0.95 -# TODO (IK): Use standard PySpark numeric types for this. -_NUMERIC_TYPES: tuple[type[DataType], ...] = ( - T.IntegerType, - T.LongType, - T.ShortType, - T.DoubleType, - T.FloatType, - T.DecimalType, - T.ByteType, -) - -# TODO (IK): Reuse `TEXT_TYPES` from profile_builder.py -_TEXT_TYPES: tuple[type[DataType], ...] = (T.StringType, T.CharType, T.VarcharType) - - def _is_numeric(column_type: DataType) -> bool: - return isinstance(column_type, _NUMERIC_TYPES) + return isinstance(column_type, T.NumericType) def _is_text(column_type: DataType) -> bool: - return isinstance(column_type, _TEXT_TYPES) + return isinstance(column_type, TEXT_TYPES) def _detect_enum(ctx: DQProfileContext) -> DQSemanticType | None: @@ -158,7 +204,7 @@ def _detect_enum(ctx: DQProfileContext) -> DQSemanticType | None: The detector is the sole owner of the enum value collection: when it fires it runs a single distinct-collect on *ctx.df* and hands the result - to downstream builders through *ctx.semantic_type.properties["values"]*. + to downstream builders through *ctx.semantic_type.properties.values*. """ column_type = ctx.column_type if not (_is_text(column_type) or isinstance(column_type, (T.IntegerType, T.LongType, T.ShortType))): @@ -181,13 +227,9 @@ def _detect_enum(ctx: DQProfileContext) -> DQSemanticType | None: col = ctx.df.columns[0] distinct_rows = ctx.df.select(col).distinct().collect() - distinct_values = [row[0] for row in distinct_rows] - try: - sorted_distinct = sorted(distinct_values) - except TypeError: - sorted_distinct = distinct_values + distinct_values = {row[0] for row in distinct_rows} - return DQSemanticType(name="enum", properties={"values": tuple(sorted_distinct)}) + return DQSemanticType(name="enum", properties=EnumProperties(values=distinct_values)) def _detect_key(ctx: DQProfileContext) -> DQSemanticType | None: @@ -216,41 +258,49 @@ def _detect_key(ctx: DQProfileContext) -> DQSemanticType | None: column_type = ctx.column_type if isinstance(column_type, (T.IntegerType, T.LongType, T.ShortType)): - min_value = ctx.metrics.get("min") - max_value = ctx.metrics.get("max") - if min_value is None or max_value is None: - return None - span = max_value - min_value + 1 - if span <= 0: - return None - density = cardinality / span - if density < KEY_MIN_DENSITY_RATIO: - return None - return DQSemanticType(name="key", properties={"signal": "density"}) + return _detect_numeric_key(ctx, cardinality) if _is_text(column_type): - col = ctx.df.columns[0] - agg = ctx.df.select( - F.min(F.length(F.col(col))).alias("min_len"), - F.max(F.length(F.col(col))).alias("max_len"), - ).first() - if agg is None: - return None - min_len = agg["min_len"] - max_len = agg["max_len"] - if min_len is None or max_len is None: - return None - if max_len == 0: - # Empty-string-only column — division by zero would follow. Not a key. - return None - length_stability = min_len / max_len - if length_stability < KEY_MIN_LENGTH_STABILITY_RATIO: - return None - return DQSemanticType(name="key", properties={"signal": "length_stability"}) + return _detect_text_key(ctx) return None +def _detect_numeric_key(ctx: DQProfileContext, cardinality: int) -> DQSemanticType | None: + min_value = ctx.metrics.get("min") + max_value = ctx.metrics.get("max") + if min_value is None or max_value is None: + return None + span = max_value - min_value + 1 + if span <= 0: + return None + density = cardinality / span + if density < KEY_MIN_DENSITY_RATIO: + return None + return DQSemanticType(name="key", properties=KeyProperties(signal="density")) + + +def _detect_text_key(ctx: DQProfileContext) -> DQSemanticType | None: + col = ctx.df.columns[0] + agg = ctx.df.select( + F.min(F.length(F.col(col))).alias("min_len"), + F.max(F.length(F.col(col))).alias("max_len"), + ).first() + if agg is None: + return None + min_len = agg["min_len"] + max_len = agg["max_len"] + if min_len is None or max_len is None: + return None + if max_len == 0: + # Empty-string-only column — division by zero would follow. Not a key. + return None + length_stability = min_len / max_len + if length_stability < KEY_MIN_LENGTH_STABILITY_RATIO: + return None + return DQSemanticType(name="key", properties=KeyProperties(signal="length_stability")) + + def _detect_measurement(ctx: DQProfileContext) -> DQSemanticType | None: """Fallback numeric detector: emit *measurement* with a best-effort distribution guess.""" if not _is_numeric(ctx.column_type): @@ -263,7 +313,7 @@ def _detect_measurement(ctx: DQProfileContext) -> DQSemanticType | None: min_value = ctx.metrics.get("min") max_value = ctx.metrics.get("max") - distribution = "unknown" + distribution: Literal["constant", "uniform", "normal", "exponential", "unknown"] = "unknown" if mean is not None and stddev is not None and min_value is not None and max_value is not None: try: span = float(max_value) - float(min_value) @@ -284,7 +334,7 @@ def _detect_measurement(ctx: DQProfileContext) -> DQSemanticType | None: else: distribution = "normal" - return DQSemanticType(name="measurement", properties={"distribution": distribution}) + return DQSemanticType(name="measurement", properties=MeasurementProperties(distribution=distribution)) def _detect_text(ctx: DQProfileContext) -> DQSemanticType | None: diff --git a/tests/integration/profiler/test_semantic_profiling.py b/tests/integration/profiler/test_semantic_profiling.py index 38b11e580..95b02a721 100644 --- a/tests/integration/profiler/test_semantic_profiling.py +++ b/tests/integration/profiler/test_semantic_profiling.py @@ -83,12 +83,12 @@ def test_default_semantic_registry_classifies_grounded_columns(spark, ws): assert all(p.semantic_type == "enum" for p in by_column.get("vehicle_type", []) if p.name == "is_in") for measurement_col in ("cargo_weight", "deal_value"): - names = {p.name for p in by_column.get(measurement_col, [])} + names = {profile.name for profile in by_column.get(measurement_col, [])} assert "min_max" in names assert "is_in" not in names - for p in by_column.get(measurement_col, []): - if p.name == "min_max": - assert p.semantic_type == "measurement" + for profile in by_column.get(measurement_col, []): + if profile.name == "min_max": + assert profile.semantic_type == "measurement" user_id_names = {p.name for p in by_column.get("user_id", [])} assert "min_max" not in user_id_names @@ -99,18 +99,18 @@ def test_default_semantic_registry_classifies_grounded_columns(spark, ws): assert "min_max" not in order_id_names user_name_profiles = by_column.get("user_name", []) - assert all(p.name != "is_in" for p in user_name_profiles) + assert all(profile.name != "is_in" for profile in user_name_profiles) # user_name has variable length → falls through to text (not key). text emits no additional # rules; only null_or_empty candidates remain. - for p in user_name_profiles: - if p.semantic_type is not None: - assert p.semantic_type == "text" + for profile in user_name_profiles: + if profile.semantic_type is not None: + assert profile.semantic_type == "text" work_desc_profiles = by_column.get("work_description", []) - assert all(p.name != "is_in" for p in work_desc_profiles) - for p in work_desc_profiles: - if p.semantic_type is not None: - assert p.semantic_type == "text" + assert all(profile.name != "is_in" for profile in work_desc_profiles) + for profile in work_desc_profiles: + if profile.semantic_type is not None: + assert profile.semantic_type == "text" def test_registry_without_enum_falls_through_to_measurement(spark, ws): @@ -165,7 +165,6 @@ def test_profile_table_populates_metadata_from_unity_catalog(spark, ws, make_sch def _spy_detect(ctx): captured.append(dict(ctx.metadata)) - return None spy = DQSemanticTypeDetector(name="spy", detect=_spy_detect) registry = SemanticRegistry(detectors=(spy, *default_semantic_detectors())) diff --git a/tests/unit/profiler/test_semantic.py b/tests/unit/profiler/test_semantic.py index 7023a9558..1dd921c64 100644 --- a/tests/unit/profiler/test_semantic.py +++ b/tests/unit/profiler/test_semantic.py @@ -13,6 +13,9 @@ DQProfileContext, DQSemanticType, DQSemanticTypeDetector, + EnumProperties, + KeyProperties, + MeasurementProperties, SemanticRegistry, default_semantic_detectors, ) @@ -23,21 +26,28 @@ # --------------------------------------------------------------------------- -def test_dq_semantic_type_tuple_property_round_trips(): - sem = DQSemanticType(name="enum", properties={"values": (1, 2, 3)}) - assert sem.properties["values"] == (1, 2, 3) - assert isinstance(sem.properties["values"], tuple) +def test_dq_semantic_type_set_property_round_trips(): + sem = DQSemanticType(name="enum", properties=EnumProperties(values={1, 2, 3})) + typed = sem.get_typed_properties(EnumProperties) + assert typed is not None + assert typed.values == {1, 2, 3} + assert isinstance(typed.values, set) + + +def test_dq_semantic_type_get_typed_properties_wrong_type_returns_none(): + sem = DQSemanticType(name="enum", properties=EnumProperties(values={1, 2, 3})) + assert sem.get_typed_properties(KeyProperties) is None def test_dq_semantic_type_frozen_blocks_field_reassignment(): - sem = DQSemanticType(name="enum") + sem = DQSemanticType(name="enum", properties=EnumProperties(values={1})) with pytest.raises(ValidationError): sem.name = "x" # type: ignore[misc] def test_dq_semantic_type_rejects_unsupported_property_value(): with pytest.raises(ValidationError): - DQSemanticType(name="enum", properties={"values": {"nested": "dict"}}) # type: ignore[arg-type] + EnumProperties(values={"nested": "dict"}) # type: ignore[arg-type] # --------------------------------------------------------------------------- @@ -62,7 +72,7 @@ def test_semantic_registry_default_matches_helper(): def test_semantic_registry_empty_constructor_has_empty_detectors(): - assert SemanticRegistry().detectors == () + assert not SemanticRegistry().detectors def test_semantic_registry_prepend_returns_new_instance(): @@ -131,7 +141,7 @@ def test_chain_first_match_wins_and_stops(): calls: list[str] = [] def make_detector(name: str, match: bool): - def _detect(ctx): + def _detect(_ctx): calls.append(name) return DQSemanticType(name=name) if match else None @@ -180,7 +190,8 @@ def test_measurement_detector_positive_numeric_column(): result = DEFAULT_MEASUREMENT_DETECTOR.detect(ctx) assert result is not None assert result.name == "measurement" - assert "distribution" in result.properties + assert isinstance(result.properties, MeasurementProperties) + assert result.properties.distribution in {"constant", "uniform", "normal", "exponential", "unknown"} def test_measurement_detector_rejects_string_column(): @@ -209,7 +220,8 @@ def test_key_detector_numeric_dense_positive(): result = DEFAULT_KEY_DETECTOR.detect(ctx) assert result is not None assert result.name == "key" - assert result.properties.get("signal") == "density" + assert isinstance(result.properties, KeyProperties) + assert result.properties.signal == "density" def test_key_detector_numeric_sparse_negative_falls_through(): @@ -268,8 +280,9 @@ def test_enum_detector_positive_low_cardinality(): result = DEFAULT_ENUM_DETECTOR.detect(ctx) assert result is not None assert result.name == "enum" - assert isinstance(result.properties["values"], tuple) - assert set(result.properties["values"]) == {"car", "truck", "van"} + typed = result.get_typed_properties(EnumProperties) + assert typed is not None + assert typed.values == {"car", "truck", "van"} # --------------------------------------------------------------------------- @@ -304,7 +317,8 @@ def test_key_detector_string_uniform_length_positive(): result = DEFAULT_KEY_DETECTOR.detect(ctx) assert result is not None assert result.name == "key" - assert result.properties.get("signal") == "length_stability" + assert isinstance(result.properties, KeyProperties) + assert result.properties.signal == "length_stability" def test_key_detector_string_variable_length_negative(): diff --git a/tests/unit/test_profile_builder.py b/tests/unit/test_profile_builder.py index 2fd02232c..15478bf45 100644 --- a/tests/unit/test_profile_builder.py +++ b/tests/unit/test_profile_builder.py @@ -17,7 +17,7 @@ register_profile_builder, validate_profile_options, ) -from databricks.labs.dqx.profiler.semantic import DQProfileContext, DQSemanticType +from databricks.labs.dqx.profiler.semantic import DQProfileContext, DQSemanticType, EnumProperties @pytest.fixture @@ -97,8 +97,8 @@ def _my_builder(*_): def test_register_profile_builder_context_type_uses_contextual_slot(): - @register_profile_builder("_test_ctx", type="context") - def _ctx_builder(ctx): + @register_profile_builder("_test_ctx", kind="context") + def _ctx_builder(_ctx): return None try: @@ -110,7 +110,7 @@ def _ctx_builder(ctx): def test_register_profile_builder_legacy_type_uses_builder_slot(): - @register_profile_builder("_test_legacy_kw", type="legacy") + @register_profile_builder("_test_legacy_kw", kind="legacy") def _legacy_builder(*_): return None @@ -386,7 +386,7 @@ def test_is_in_reuses_enum_values_without_extra_spark_action(mock_df): T.StringType(), {"count": 100, "count_non_null": 100}, {"max_in_count": 10, "distinct_ratio": 0.1}, - semantic_type=DQSemanticType(name="enum", properties={"values": ("car", "truck", "van")}), + semantic_type=DQSemanticType(name="enum", properties=EnumProperties(values={"car", "truck", "van"})), ) ) assert profile is not None From bfc04febf264251f945bcc9e3ecf24b5611682c9 Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 29 Aug 2026 15:01:12 +0200 Subject: [PATCH 3/9] integration tests fixes --- tests/integration/profiler/__init__.py | 0 ..._profiling.py => test_profile_semantic.py} | 96 ++++++++++++------- 2 files changed, 62 insertions(+), 34 deletions(-) delete mode 100644 tests/integration/profiler/__init__.py rename tests/integration/{profiler/test_semantic_profiling.py => test_profile_semantic.py} (62%) diff --git a/tests/integration/profiler/__init__.py b/tests/integration/profiler/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/integration/profiler/test_semantic_profiling.py b/tests/integration/test_profile_semantic.py similarity index 62% rename from tests/integration/profiler/test_semantic_profiling.py rename to tests/integration/test_profile_semantic.py index 95b02a721..bd34b101e 100644 --- a/tests/integration/profiler/test_semantic_profiling.py +++ b/tests/integration/test_profile_semantic.py @@ -44,7 +44,7 @@ def _make_demo_df(spark): str(uuid.UUID(int=i)), # Free-form names 3..30 chars — length_stability ~= 0.1 ("A" * (3 + (i % 27))), - "The quick brown fox jumps over the lazy dog " * (1 + i % 3), + "The quick brown fox jumps over the lazy dog " * (1 + i), ) ) return spark.createDataFrame(rows, schema=schema) @@ -66,8 +66,8 @@ def test_profile_without_semantic_registry_matches_pre_feature_output(spark, ws) profiler = DQProfiler(ws) _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) assert profiles, "expected the pre-feature profiler to emit at least one profile" - for profile in profiles: - assert profile.semantic_type is None + tagged = [(p.column, p.name, p.semantic_type) for p in profiles if p.semantic_type is not None] + assert not tagged, f"expected all profiles to have semantic_type=None, got tagged: {tagged}" def test_default_semantic_registry_classifies_grounded_columns(spark, ws): @@ -78,39 +78,52 @@ def test_default_semantic_registry_classifies_grounded_columns(spark, ws): by_column = _profile_by_column(profiles) vehicle_names = {p.name for p in by_column.get("vehicle_type", [])} - assert "is_in" in vehicle_names - assert "min_max" not in vehicle_names - assert all(p.semantic_type == "enum" for p in by_column.get("vehicle_type", []) if p.name == "is_in") + assert "is_in" in vehicle_names, f"expected 'is_in' profile on vehicle_type, got: {vehicle_names}" + assert "min_max" not in vehicle_names, f"unexpected 'min_max' profile on vehicle_type, got: {vehicle_names}" + vehicle_is_in_semantic_types = [p.semantic_type for p in by_column.get("vehicle_type", []) if p.name == "is_in"] + assert all( + st == "enum" for st in vehicle_is_in_semantic_types + ), f"expected all vehicle_type is_in profiles tagged semantic_type='enum', got: {vehicle_is_in_semantic_types}" for measurement_col in ("cargo_weight", "deal_value"): names = {profile.name for profile in by_column.get(measurement_col, [])} - assert "min_max" in names - assert "is_in" not in names - for profile in by_column.get(measurement_col, []): - if profile.name == "min_max": - assert profile.semantic_type == "measurement" + assert "min_max" in names, f"expected 'min_max' profile on {measurement_col}, got: {names}" + assert "is_in" not in names, f"unexpected 'is_in' profile on {measurement_col}, got: {names}" + min_max_semantic_types = [ + profile.semantic_type for profile in by_column.get(measurement_col, []) if profile.name == "min_max" + ] + assert all(st == "measurement" for st in min_max_semantic_types), ( + f"expected all {measurement_col} min_max profiles tagged semantic_type='measurement', " + f"got: {min_max_semantic_types}" + ) user_id_names = {p.name for p in by_column.get("user_id", [])} - assert "min_max" not in user_id_names - assert "is_in" not in user_id_names + assert "min_max" not in user_id_names, f"unexpected 'min_max' profile on user_id key column, got: {user_id_names}" + assert "is_in" not in user_id_names, f"unexpected 'is_in' profile on user_id key column, got: {user_id_names}" order_id_names = {p.name for p in by_column.get("order_id", [])} - assert "is_in" not in order_id_names - assert "min_max" not in order_id_names + assert "is_in" not in order_id_names, f"unexpected 'is_in' profile on order_id key column, got: {order_id_names}" + assert ( + "min_max" not in order_id_names + ), f"unexpected 'min_max' profile on order_id key column, got: {order_id_names}" - user_name_profiles = by_column.get("user_name", []) - assert all(profile.name != "is_in" for profile in user_name_profiles) # user_name has variable length → falls through to text (not key). text emits no additional # rules; only null_or_empty candidates remain. - for profile in user_name_profiles: - if profile.semantic_type is not None: - assert profile.semantic_type == "text" + user_name_profiles = by_column.get("user_name", []) + user_name_names = {p.name for p in user_name_profiles} + assert "is_in" not in user_name_names, f"unexpected 'is_in' profile on user_name, got: {user_name_names}" + user_name_semantic_types = {p.semantic_type for p in user_name_profiles if p.semantic_type is not None} + assert user_name_semantic_types <= { + "text" + }, f"expected user_name profiles tagged only with 'text' semantic_type, got: {user_name_semantic_types}" work_desc_profiles = by_column.get("work_description", []) - assert all(profile.name != "is_in" for profile in work_desc_profiles) - for profile in work_desc_profiles: - if profile.semantic_type is not None: - assert profile.semantic_type == "text" + work_desc_names = {p.name for p in work_desc_profiles} + assert "is_in" not in work_desc_names, f"unexpected 'is_in' profile on work_description, got: {work_desc_names}" + work_desc_semantic_types = {p.semantic_type for p in work_desc_profiles if p.semantic_type is not None} + assert work_desc_semantic_types <= { + "text" + }, f"expected work_description profiles tagged only with 'text' semantic_type, got: {work_desc_semantic_types}" def test_registry_without_enum_falls_through_to_measurement(spark, ws): @@ -122,8 +135,8 @@ def test_registry_without_enum_falls_through_to_measurement(spark, ws): by_column = _profile_by_column(profiles) vehicle_names = {p.name for p in by_column.get("vehicle_type", [])} # vehicle_type is StringType, so it now falls through to text (no min_max, no is_in). - assert "is_in" not in vehicle_names - assert "min_max" not in vehicle_names + assert "is_in" not in vehicle_names, f"unexpected 'is_in' profile on vehicle_type, got: {vehicle_names}" + assert "min_max" not in vehicle_names, f"unexpected 'min_max' profile on vehicle_type, got: {vehicle_names}" def test_prepend_custom_detector_takes_precedence(spark, ws): @@ -140,8 +153,12 @@ def _always_text(_ctx): by_column = _profile_by_column(profiles) vehicle_names = {p.name for p in by_column.get("vehicle_type", [])} # forced_text emits "text" — same gate as DEFAULT_TEXT_DETECTOR — so is_in and min_max are skipped - assert "is_in" not in vehicle_names - assert "min_max" not in vehicle_names + assert ( + "is_in" not in vehicle_names + ), f"unexpected 'is_in' profile on vehicle_type when forced_text prepended, got: {vehicle_names}" + assert ( + "min_max" not in vehicle_names + ), f"unexpected 'min_max' profile on vehicle_type when forced_text prepended, got: {vehicle_names}" def test_profile_table_populates_metadata_from_unity_catalog(spark, ws, make_schema, make_random): @@ -175,9 +192,15 @@ def _spy_detect(ctx): ) vehicle_metadata = next((m for m in captured if m.get("column_comment") == "kind of vehicle"), None) - assert vehicle_metadata is not None - assert vehicle_metadata.get("table_name") == table_name - assert vehicle_metadata.get("table_comment") == "demo table with vehicle types" + assert ( + vehicle_metadata is not None + ), f"expected captured metadata for vehicle_type with column_comment='kind of vehicle', got: {captured}" + assert ( + vehicle_metadata.get("table_name") == table_name + ), f"expected table_name={table_name!r}, got: {vehicle_metadata.get('table_name')!r}" + assert ( + vehicle_metadata.get("table_comment") == "demo table with vehicle types" + ), f"expected table_comment='demo table with vehicle types', got: {vehicle_metadata.get('table_comment')!r}" def test_enum_detector_values_reused_by_is_in_builder(spark, ws): @@ -197,8 +220,13 @@ def _tracing_detect(ctx): profiler = DQProfiler(ws, semantic_registry=registry) _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) - assert invocations == ["vehicle_type"] + assert invocations == [ + "vehicle_type" + ], f"expected enum detector invoked exactly once for vehicle_type, got invocations: {invocations}" is_in_profiles = [p for p in profiles if p.column == "vehicle_type" and p.name == "is_in"] - assert len(is_in_profiles) == 1 + assert ( + len(is_in_profiles) == 1 + ), f"expected exactly one 'is_in' profile for vehicle_type, got {len(is_in_profiles)}: {is_in_profiles}" is_in = is_in_profiles[0] - assert set(is_in.parameters["in"]) == {"car", "truck", "van"} + values = set(is_in.parameters["in"]) + assert values == {"car", "truck", "van"}, f"expected is_in values {{car, truck, van}}, got: {values}" From 8a225c70bb1bd0c7564a200fdffd46a78bc5b942 Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 29 Aug 2026 16:03:53 +0200 Subject: [PATCH 4/9] perfornance tests added --- .../labs/dqx/profiler/profile_builder.py | 11 +-- src/databricks/labs/dqx/profiler/profiler.py | 2 - src/databricks/labs/dqx/profiler/semantic.py | 18 +--- tests/perf/test_profile_semantic.py | 91 +++++++++++++++++++ tests/unit/profiler/test_semantic.py | 9 +- 5 files changed, 102 insertions(+), 29 deletions(-) create mode 100644 tests/perf/test_profile_semantic.py diff --git a/src/databricks/labs/dqx/profiler/profile_builder.py b/src/databricks/labs/dqx/profiler/profile_builder.py index 6d5ed769d..01f5a61f8 100644 --- a/src/databricks/labs/dqx/profiler/profile_builder.py +++ b/src/databricks/labs/dqx/profiler/profile_builder.py @@ -116,12 +116,11 @@ def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: return None semantic_type = ctx.semantic_type - if semantic_type is not None and semantic_type.name != "enum": - # A non-enum semantic type was assigned; do not emit an is_in candidate that would - # contradict the semantic classification. - return None - - if semantic_type is not None and semantic_type.name == "enum": + if semantic_type is not None: + if semantic_type.name != "enum": + # A non-enum semantic type was assigned; do not emit an is_in candidate that would + # contradict the semantic classification. + return None enum_props = semantic_type.get_typed_properties(EnumProperties) if enum_props is None or not enum_props.values: return None diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 31a285e06..d919706f5 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -603,8 +603,6 @@ def _build_profiles_for_column( metadata_map = dict(column_metadata) if column_metadata else {} semantic_type = self._detect_semantic_type(column_df, field_name, field_type, metrics, opts, metadata_map) - # Reconstruct via the constructor (not model_copy) so validators run — same rationale as - # SemanticRegistry.prepend/replace. builder_ctx = DQProfileContext( df=column_df, column_name=field_name, diff --git a/src/databricks/labs/dqx/profiler/semantic.py b/src/databricks/labs/dqx/profiler/semantic.py index 676c9ec05..02933864f 100644 --- a/src/databricks/labs/dqx/profiler/semantic.py +++ b/src/databricks/labs/dqx/profiler/semantic.py @@ -13,8 +13,7 @@ Public surface: * models: *DQSemanticType*, *DQProfileContext*, *DQSemanticTypeDetector* - * properties: *DQSemanticTypeProperties*, *EnumProperties*, *KeyProperties*, - *MeasurementProperties* + * properties: *DQSemanticTypeProperties*, *EnumProperties*, *MeasurementProperties* * registry: *SemanticRegistry* * detectors: *DEFAULT_ENUM_DETECTOR*, *DEFAULT_KEY_DETECTOR*, *DEFAULT_MEASUREMENT_DETECTOR*, *DEFAULT_TEXT_DETECTOR*, @@ -65,17 +64,6 @@ class EnumProperties(DQSemanticTypeProperties): values: set[str] | set[int] -class KeyProperties(DQSemanticTypeProperties): - """Properties for the built-in *key* semantic type. - - Attributes: - signal: Which secondary signal fired alongside distinctness — *density* - for numeric keys, *length_stability* for string keys. - """ - - signal: Literal["density", "length_stability"] - - class MeasurementProperties(DQSemanticTypeProperties): """Properties for the built-in *measurement* semantic type. @@ -277,7 +265,7 @@ def _detect_numeric_key(ctx: DQProfileContext, cardinality: int) -> DQSemanticTy density = cardinality / span if density < KEY_MIN_DENSITY_RATIO: return None - return DQSemanticType(name="key", properties=KeyProperties(signal="density")) + return DQSemanticType(name="key") def _detect_text_key(ctx: DQProfileContext) -> DQSemanticType | None: @@ -298,7 +286,7 @@ def _detect_text_key(ctx: DQProfileContext) -> DQSemanticType | None: length_stability = min_len / max_len if length_stability < KEY_MIN_LENGTH_STABILITY_RATIO: return None - return DQSemanticType(name="key", properties=KeyProperties(signal="length_stability")) + return DQSemanticType(name="key") def _detect_measurement(ctx: DQProfileContext) -> DQSemanticType | None: diff --git a/tests/perf/test_profile_semantic.py b/tests/perf/test_profile_semantic.py new file mode 100644 index 000000000..fa5a6763f --- /dev/null +++ b/tests/perf/test_profile_semantic.py @@ -0,0 +1,91 @@ +"""Benchmarks for semantic-aware profiling. + +Compares *DQProfiler.profile* runs with and without a *SemanticRegistry* to +quantify the overhead of the default semantic-detection chain (enum → key → +measurement → text) on worst-case columns — inputs shaped so the classifier +walks the full chain before matching the trailing fallback detector: + +* **Integer, worst-case measurement** — random ints in a wide range so + distinctness is ~1.0 but density (*count_distinct / (max - min + 1)*) is + well below the key threshold. Chain traversal: enum rejects (cardinality + ≫ *max_in_count*), key rejects (density below threshold), measurement + matches. + +* **String, worst-case text** — variable-length free-form words so + distinctness is ~1.0 but length stability (*min_len / max_len*) is well + below the key threshold. Chain traversal: enum rejects, key runs a Spark + aggregate for min/max length and then rejects on low stability, text + matches. The key detector's aggregate is the primary added cost this + benchmark surfaces. +""" + +import pytest + +from databricks.labs.dqx.profiler.profiler import DQProfiler +from databricks.labs.dqx.profiler.semantic import SemanticRegistry + + +BENCHMARK_ROWS = 10_000_000 # 10M rows keeps benchmark runtime bounded while still exercising Spark aggregates +BENCHMARK_COLUMNS = 4 + +# Wide-range random integers force key-density rejection: cardinality is ~BENCHMARK_ROWS but +# max-min+1 is 2*10^9, so density = BENCHMARK_ROWS / 2*10^9 ≪ KEY_MIN_DENSITY_RATIO and the +# chain falls through to measurement. The upper bound must stay below int32 max (2_147_483_647) +# because dbldatagen samples via double then casts to int; a larger max triggers CAST_OVERFLOW. +INT_WIDE_RANGE_OPTS = {"minValue": 1, "maxValue": 2_000_000_000, "random": True} + +# Dash-separated word tokens (mirroring the DEFAULT_EMAIL_TEMPLATE style in conftest) produce +# variable-length free-form strings, driving length_stability (min_len / max_len) well below +# KEY_MIN_LENGTH_STABILITY_RATIO so key rejects and text matches. +TEXT_VARIABLE_LENGTH_TEMPLATE = r"\w-\w-\w-\w" + +PROFILE_OPTIONS = { + # sample_fraction=None profiles the full dataset so numbers reflect the real cost, + # not sample-driven noise. + "sample_fraction": None, + # LLM primary-key detection is orthogonal to the semantic chain and would dominate the timing. + "llm_primary_key_detection": False, +} + +SEMANTIC_REGISTRY_PARAMS = [ + pytest.param(None, id="no_semantic_registry"), + pytest.param(SemanticRegistry.default(), id="default_semantic_registry"), +] + + +@pytest.mark.parametrize( + "generated_integer_df", + [{"n_rows": BENCHMARK_ROWS, "n_columns": BENCHMARK_COLUMNS, "opts": INT_WIDE_RANGE_OPTS}], + indirect=True, + ids=lambda param: f"n_rows_{param['n_rows']}_n_columns_{param['n_columns']}", +) +@pytest.mark.parametrize("semantic_registry", SEMANTIC_REGISTRY_PARAMS) +@pytest.mark.benchmark(group="test_benchmark_profile_semantic_integer_measurement") +def test_benchmark_profile_semantic_integer_measurement(benchmark, ws, generated_integer_df, semantic_registry): + """Profile wide-range random integers — worst-case measurement path.""" + _columns, df, n_rows = generated_integer_df + profiler = DQProfiler(workspace_client=ws, semantic_registry=semantic_registry) + _stats, profiles = benchmark(lambda: profiler.profile(df, options=PROFILE_OPTIONS)) + assert profiles, f"expected profiler to emit at least one profile for {n_rows} rows" + + +@pytest.mark.parametrize( + "generated_string_df", + [ + { + "n_rows": BENCHMARK_ROWS, + "n_columns": BENCHMARK_COLUMNS, + "template": TEXT_VARIABLE_LENGTH_TEMPLATE, + } + ], + indirect=True, + ids=lambda param: f"n_rows_{param['n_rows']}_n_columns_{param['n_columns']}", +) +@pytest.mark.parametrize("semantic_registry", SEMANTIC_REGISTRY_PARAMS) +@pytest.mark.benchmark(group="test_benchmark_profile_semantic_string_text") +def test_benchmark_profile_semantic_string_text(benchmark, ws, generated_string_df, semantic_registry): + """Profile variable-length random text — worst-case text path (key detector triggers Spark aggregate).""" + _columns, df, n_rows = generated_string_df + profiler = DQProfiler(workspace_client=ws, semantic_registry=semantic_registry) + _stats, profiles = benchmark(lambda: profiler.profile(df, options=PROFILE_OPTIONS)) + assert profiles, f"expected profiler to emit at least one profile for {n_rows} rows" diff --git a/tests/unit/profiler/test_semantic.py b/tests/unit/profiler/test_semantic.py index 1dd921c64..86d498162 100644 --- a/tests/unit/profiler/test_semantic.py +++ b/tests/unit/profiler/test_semantic.py @@ -14,7 +14,6 @@ DQSemanticType, DQSemanticTypeDetector, EnumProperties, - KeyProperties, MeasurementProperties, SemanticRegistry, default_semantic_detectors, @@ -36,7 +35,7 @@ def test_dq_semantic_type_set_property_round_trips(): def test_dq_semantic_type_get_typed_properties_wrong_type_returns_none(): sem = DQSemanticType(name="enum", properties=EnumProperties(values={1, 2, 3})) - assert sem.get_typed_properties(KeyProperties) is None + assert sem.get_typed_properties(MeasurementProperties) is None def test_dq_semantic_type_frozen_blocks_field_reassignment(): @@ -220,8 +219,7 @@ def test_key_detector_numeric_dense_positive(): result = DEFAULT_KEY_DETECTOR.detect(ctx) assert result is not None assert result.name == "key" - assert isinstance(result.properties, KeyProperties) - assert result.properties.signal == "density" + assert result.properties is None def test_key_detector_numeric_sparse_negative_falls_through(): @@ -317,8 +315,7 @@ def test_key_detector_string_uniform_length_positive(): result = DEFAULT_KEY_DETECTOR.detect(ctx) assert result is not None assert result.name == "key" - assert isinstance(result.properties, KeyProperties) - assert result.properties.signal == "length_stability" + assert result.properties is None def test_key_detector_string_variable_length_negative(): From f604aa18184468fe90f235b9e1555bc66c8ba2ea Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 29 Aug 2026 16:32:47 +0200 Subject: [PATCH 5/9] fix performance tests --- tests/perf/test_profile_semantic.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/perf/test_profile_semantic.py b/tests/perf/test_profile_semantic.py index fa5a6763f..5ecff49ca 100644 --- a/tests/perf/test_profile_semantic.py +++ b/tests/perf/test_profile_semantic.py @@ -25,20 +25,23 @@ from databricks.labs.dqx.profiler.semantic import SemanticRegistry -BENCHMARK_ROWS = 10_000_000 # 10M rows keeps benchmark runtime bounded while still exercising Spark aggregates +# 100K rows is smaller than the DEFAULT_ROWS used elsewhere in tests/perf because the legacy +# *is_in* profile builder does an unconditional *df.select(col).distinct().collect()* on every +# column. Our worst-case fixtures deliberately produce ~1.0 distinctness, so the collect returns +# roughly BENCHMARK_ROWS rows to the driver — at DEFAULT_ROWS this exceeds the Spark Connect +# message limit / driver memory and the no-registry baseline fails with PythonException. 100K +# keeps the collect under a few MB while still letting the semantic-detection chain exercise +# its full worst-case path (enum + key aggregates fire, then fall through to the trailing +# fallback detector). +BENCHMARK_ROWS = 100_000 BENCHMARK_COLUMNS = 4 # Wide-range random integers force key-density rejection: cardinality is ~BENCHMARK_ROWS but # max-min+1 is 2*10^9, so density = BENCHMARK_ROWS / 2*10^9 ≪ KEY_MIN_DENSITY_RATIO and the -# chain falls through to measurement. The upper bound must stay below int32 max (2_147_483_647) +# chain falls through to measurement. The upper bound stays below int32 max (2_147_483_647) # because dbldatagen samples via double then casts to int; a larger max triggers CAST_OVERFLOW. INT_WIDE_RANGE_OPTS = {"minValue": 1, "maxValue": 2_000_000_000, "random": True} -# Dash-separated word tokens (mirroring the DEFAULT_EMAIL_TEMPLATE style in conftest) produce -# variable-length free-form strings, driving length_stability (min_len / max_len) well below -# KEY_MIN_LENGTH_STABILITY_RATIO so key rejects and text matches. -TEXT_VARIABLE_LENGTH_TEMPLATE = r"\w-\w-\w-\w" - PROFILE_OPTIONS = { # sample_fraction=None profiles the full dataset so numbers reflect the real cost, # not sample-driven noise. @@ -75,7 +78,6 @@ def test_benchmark_profile_semantic_integer_measurement(benchmark, ws, generated { "n_rows": BENCHMARK_ROWS, "n_columns": BENCHMARK_COLUMNS, - "template": TEXT_VARIABLE_LENGTH_TEMPLATE, } ], indirect=True, From fd4c3eaf7b664f1a577f086207f265574b01cc73 Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 5 Sep 2026 14:12:56 +0200 Subject: [PATCH 6/9] addressed code review feedback --- .agents/2026-09-05-11-42-05-pr-1491-review.md | 164 ++++++++ .agents/2026-09-05-12-48-53-plan.md | 386 ++++++++++++++++++ .agents/2026-09-05-12-48-53-review-round-0.md | 50 +++ docs/dqx/docs/reference/profiler.mdx | 70 ++-- .../labs/dqx/profiler/profile_builder.py | 15 +- src/databricks/labs/dqx/profiler/profiler.py | 99 +---- src/databricks/labs/dqx/profiler/semantic.py | 154 +++++-- tests/integration/test_profile_semantic.py | 118 ++++-- tests/unit/profiler/test_semantic.py | 177 +++++++- tests/unit/test_profile_builder.py | 63 ++- 10 files changed, 1079 insertions(+), 217 deletions(-) create mode 100644 .agents/2026-09-05-11-42-05-pr-1491-review.md create mode 100644 .agents/2026-09-05-12-48-53-plan.md create mode 100644 .agents/2026-09-05-12-48-53-review-round-0.md diff --git a/.agents/2026-09-05-11-42-05-pr-1491-review.md b/.agents/2026-09-05-11-42-05-pr-1491-review.md new file mode 100644 index 000000000..f0d64f6ad --- /dev/null +++ b/.agents/2026-09-05-11-42-05-pr-1491-review.md @@ -0,0 +1,164 @@ +--- +title: Manual PR-state review of feature/semantic_type_classification (PR #1491) +scope: uncommitted changes + branch commits + previously-addressed PR review comments +created_at: 2026-09-05T11:42:05Z +pr: https://github.com/databrickslabs/dqx/pull/1491 +verdict: ship-ready with minor polish +--- + +# Top Critical Issues + +## 1. Enum-gate denominator still differs from legacy `is_in` — residual divergence +`_detect_enum` computes `cardinality / count_non_null` (nulls excluded), but the legacy +`make_is_in_profile` computes `distinct_count / total_count` (nulls included). Even with the +`distinct_ratio` threshold now aligned, a null-heavy column can be classified differently by +the two paths. + +**Example** — 100 rows, 40 non-null, 3 distinct: +- Legacy is_in: 3/100 = 0.03 < 0.05 → emits `is_in`. +- Semantic `_detect_enum`: 3/40 = 0.075 ≥ 0.05 → rejects. + +Reviewer comment #3 asked for alignment; the threshold was aligned but the denominator was not. +Options: (a) switch semantic to `count`, or (b) document the deliberate difference in the +`ENUM_MAX_CARDINALITY_RATIO` comment. + +**Severity**: modest — semantic profiling is opt-in and the divergence is bounded. +Feedback: adjust `is_in` profile to take into count non-null values. +**Resolution (2026-09-05)**: `make_is_in_profile` in `src/databricks/labs/dqx/profiler/profile_builder.py` now divides by `count_non_null` (matching `_detect_enum`) and short-circuits when `count_non_null == 0`. Unit tests in `tests/unit/test_profile_builder.py` updated to supply `count_non_null` alongside `count`; the previously named `test_is_in_total_count_zero_returns_none` is now `test_is_in_count_non_null_zero_returns_none`. `distinct_ratio` docs row updated from "5% of total" → "5% of non-null values". + +## 2. `with_metrics` fires on every builder iteration — wasteful cache-invalidation +`profiler.py:_build_profiles_for_column` calls `builder_ctx = builder_ctx.with_metrics(metrics)` +inside the loop, allocating a new frozen model + copy of `metrics` for every contextual builder +— even the first, and even when no prior builder mutated `metrics`. Only `min_max` writes back +today, so a cheaper pattern is: refresh only *after* a builder returns a `min_max` profile. + +Correctness is fine; the current implementation pays a small allocation tax on every +column × builder. + +**Severity**: low. +Feedback: skip + +## 3. `_fetch_table_metadata` bare `except Exception` — safe but broad +`profiler.py:461` catches everything, per reviewer comment #4. The comment explains the reasoning, +and it does honor the docstring guarantee. But this now swallows bugs in metadata assembly (e.g. +`AttributeError` from an SDK schema change). Consider tightening to `(DatabricksError, ValueError, +TimeoutError)` — the concrete client-side failures actually observed. + +**NOTE (updated 2026-09-05T11:42:05Z)**: superseded — the entire `_fetch_table_metadata` method, +its call chain, and the `DQProfileContext.metadata` field were removed in a subsequent commit. +Finding no longer applies. + +## 4. Naming: `SemanticRegistry.insert(name, detector)` reads as *insert at position `name`* +The method inserts *after* the entry named `name`. Docstring says so, but the naming still +surprises. `insert_after` was rejected as verbose. An alternative would be `after(name, detector)` +— reads naturally and encodes the semantic. + +**Severity**: cosmetic; safe to leave. +Feedback: this method must be already removed, ignore + +## 5. Docs code block for the composition summary lacks a language hint +`profiler.mdx:156` opens a fence with no language marker: +```` +``` +# prepend(detector) → puts one detector at the front (highest priority) +... +``` +```` +Not `python` (uses `→`), not `text`. Renders monospace without syntax highlighting. + +**Severity**: minor — mark as `text` or wrap in prose. +Feedback: fix +**Resolution (2026-09-05)**: `docs/dqx/docs/reference/profiler.mdx` fence now opens `` ```text `` and the `#` comment markers are dropped in favour of plain lines. + +## 6. PR title is a placeholder ("Feature/semantic type classification") +Follows branch-name convention. Before merge, replace with a real title (e.g. "Add opt-in +semantic-aware profiling to DQProfiler"). Also refresh the PR body — some sections still +describe `SemanticRegistry` as offering only `default()`, `prepend(...)`, and `replace([...])`, +which is the old API. + +**Severity**: blocker for merge (metadata only, no code change). +Feedback: fix +**Resolution (2026-09-05)**: PR #1491 title updated on GitHub to *"Add opt-in semantic-aware profiling to DQProfiler"* via `gh pr edit`. + +## 7. Deprecated / broken PR-body claims +The PR body describes `SemanticRegistry.replace([...])` as accepting a list. That's the old +API. Refresh the description to list `of`, `append`, `insert(name, detector)`, +`replace(name, detector)`, `remove(name)` before merge. + +**Severity**: blocker for merge (metadata only). +Feedback: fix +**Resolution (2026-09-05)**: PR #1491 body refreshed on GitHub via `gh pr edit` — now lists the correct composition surface (`prepend`, `append`, `insert`, `replace(name, detector)`, `remove`, `of(...)`), drops the removed `metadata` field from the `DQProfileContext` description, and updates the enum-gate/denominator prose. UC metadata plumbing bullet removed to reflect its removal from the branch. + +## 8. No unit test asserts `with_metrics` is actually observed by a downstream contextual builder +`test_semantic.py` covers `with_metrics` in isolation (snapshot semantics), but there is no test +proving that a contextual builder registered *after* `min_max` actually sees the updated +`min`/`max` in `ctx.metrics`. The bug this fix was meant to prevent is a silent latent contract +break (reviewer comment #2). Add a builder-loop integration test to lock the behavior down — +otherwise a future refactor could reintroduce the bug without failing tests. + +**Severity**: modest — reviewer round-0 finding #2 also flagged this. +Feedback: fix +**Resolution (2026-09-05)**: `test_contextual_builder_after_min_max_observes_resolved_min_max` added in `tests/integration/test_profile_semantic.py`. Registers a spy contextual builder via `@register_profile_builder(kind="context")` (lands after `min_max` in `PROFILE_BUILDER_REGISTRY` insertion order), asserts spy's `ctx.metrics` carries the resolved `min`/`max`, cleans up in `finally`. Uses `remove_outliers=False` for exact expected values. + +## 9. Integration-test fixture inconsistency +`test_default_semantic_registry_classifies_grounded_columns` and +`test_enum_detector_values_reused_by_is_in_builder` pass `"distinct_ratio": 0.1` explicitly; +other tests using `_make_demo_df` do not. This is fine for the specific fixtures used, but a +comment on `_make_demo_df` noting *"3/60 = 0.05 exactly hits the tightened default; callers +who want the enum to fire must widen distinct_ratio"* would prevent future surprise. + +**Severity**: minor doc hygiene. +Feedback: skip + +## 10. Reviewer comment #9 / #10 point 3 — decorator vs value-object asymmetry — remains unaddressed +Two distinct mental models for "add my thing to the chain" (`@register_profile_builder` for +builders vs constructor threading for detectors) still coexist. Explicitly out of scope per +the original prompt. Worth capturing as a follow-up issue so it does not get lost. + +**Severity**: architectural, not a blocker. +Feedback: skip + +--- + +# Final Verdict + +**Ship-ready with minor polish.** + +The five inline reviewer findings (#1–#5) and reviewer #6/#7/#8 doc nits are all addressed. +The semantic-registry ergonomics gaps in #10 (whole-chain naming and missing +`append/insert/remove/replace-one-by-name` operations) are all fixed. Remaining items are +follow-up polish, not blockers. + +## Blockers before merge +- #6 PR title +- #7 PR body accuracy + +Both are metadata-only, no code changes required. + +## Nice-to-have before merge +- #5 doc code-fence language +- #9 fixture comment + +## Recommended follow-up issues +- #1 denominator alignment (or documentation) +- #2 per-iteration allocation +- ~~#3 exception scope~~ (resolved by full removal of `_fetch_table_metadata` plumbing) +- #8 loop-integration test +- #10 API-model asymmetry + +## Assessment of core work +The core architecture (immutable registry, first-match-wins chain, contextual builders reusing +enum values, UC metadata plumbing) is sound. The correctness bugs from the review (silent check +drops on `ShortType`, min/max write-back, permissive enum gate, `DatabricksError`-only catch, +docstring typos) are all fixed with regression coverage. Code quality is production-grade: fully +typed, immutable value objects, Pydantic-validated invariants, and no linting suppressions. + +--- + +# Follow-up (post-review) + +After this review, the user directed removal of the entire UC metadata plumbing +(`_fetch_table_metadata`, `_build_column_metadata`, the `table_metadata`/`column_metadata` +kwargs, the `metadata_map` propagation, and the vestigial `DQProfileContext.metadata` field). +This resolves finding #3 outright and removes the "Table metadata" documentation section along +with its integration test. diff --git a/.agents/2026-09-05-12-48-53-plan.md b/.agents/2026-09-05-12-48-53-plan.md new file mode 100644 index 000000000..ca6a949a4 --- /dev/null +++ b/.agents/2026-09-05-12-48-53-plan.md @@ -0,0 +1,386 @@ +--- +title: Address review comments on PR #1491 (semantic-aware profiling) except the decorator-vs-immutable-registry mental-model feedback. +status: implemented +created_at: 2026-09-05T10:48:53Z +updated_at: 2026-09-05T11:12:57Z +implemented_at: 2026-09-05T11:28:10Z + +input: > + https://github.com/databrickslabs/dqx/pull/1491 - this is PR review for the current branch. + Readout carefully and address comments, except: "They way this is constructed is different than + registry approach like @register_profile_builder(...) . Profile builders are customized with a + decorator against a global registry; semantic detectors are customized by constructing an + immutable value object and threading it through a constructor. Two different mental models for + 'add my thing to the chain.'" and similar feedback on using decorator. + +context: > + Feature branch feature/semantic_type_classification adds opt-in semantic-aware profiling to + `DQProfiler`. Reviewer mwojtyczka posted 10 inline comments on 2026-09-01 with a + CHANGES_REQUESTED verdict. Explicitly out of scope per user: comment #9 (decorator vs immutable + mental-model criticism at profiler.mdx:183) and comment #10 point 3 (same criticism restated at + semantic.py:391). All other structural/correctness/doc feedback is in scope. + + Relevant files: + - src/databricks/labs/dqx/profiler/semantic.py + - src/databricks/labs/dqx/profiler/profile_builder.py + - src/databricks/labs/dqx/profiler/profiler.py + - docs/dqx/docs/reference/profiler.mdx + - tests/unit/profiler/test_semantic.py + - tests/unit/test_profile_builder.py + - tests/integration/test_profile_semantic.py + - tests/perf/test_profile_semantic.py + + Comment → step mapping: + #1 → Step 1 (ShortType enum gap) + #2 → Step 2 (metrics write-back through ctx) + #3 → Step 3 (enum cardinality gate alignment with distinct_ratio) + #4 → Step 4 (broaden UC metadata exception handler) + #5 → Step 5 (docs `type=` → `kind=`) + #6 → Step 5 (docs "iff" wording) + #7 → Step 6 (docs comments on prepend/replace and #8 rename) + #8 → Step 6/7 (rename replace, add append/remove/insert; docs update) + #10 → Step 6/7 (points 1 & 2 only — rename + missing operations) + #9 and #10 point 3 → INTENTIONALLY SKIPPED per user +--- + +# Plan + +### Step - 1 + +**Description** + +Close the `ShortType` gap flagged in comment #1 by widening the `is_in` +builder's supported types (per user feedback: short integers are a valid +enum value range, so keep them in `_detect_enum` and make `is_in` accept +them too). + +- In `src/databricks/labs/dqx/profiler/profile_builder.py`, add + `T.ShortType` to `_supports_distinct` so the tuple becomes + `(T.IntegerType, T.LongType, T.ShortType) + TEXT_TYPES`. Update the + helper's docstring to reflect the extended set. +- Leave `_detect_enum` in `src/databricks/labs/dqx/profiler/semantic.py` + unchanged (`ShortType` stays in its accepted-types tuple). Add a one-line + comment on the `_supports_distinct` tuple noting that its accepted set + must stay in sync with `_detect_enum` so semantic-enum classification + never suppresses `min_max` without producing an `is_in` in return. +- Extend the parametrised type coverage in + `tests/unit/test_profile_builder.py`: + * Add `T.ShortType()` to the parametrise list on + `test_is_in_char_varchar_type_returns_profile` (rename generically to + `test_is_in_supported_types_returns_profile` if it makes the parametrise + read cleaner) so it exercises Short alongside `CharType`/`VarcharType`. + * Remove `T.ShortType()` from any negative parametrise if present (grep + confirms it is not currently in `test_is_in_unsupported_type_returns_none`, + so nothing to remove there — verify at implementation time). +- Add a positive integration case in + `tests/integration/test_profile_semantic.py` (or, if simpler, + `tests/integration/test_profile_builder.py`) covering a low-cardinality + `ShortType` column: assert it is classified as `enum` **and** receives + an `is_in` profile, so the previously-silent gap is regression-guarded + end-to-end. Use the existing `_make_demo_df`-style fixture pattern. +- Confirm `is_in_list` in `src/databricks/labs/dqx/check_funcs.py:524` + accepts a `ShortType` Column at apply time (it uses `.isin(...)` which + is type-agnostic in Spark). No change needed there; call out in the + step feedback if apply-time behaviour differs from expectations. + +**Feedback** + + +### Step - 2 + +**Description** + +Fix the min/max write-back gap flagged in comment #2 +(`src/databricks/labs/dqx/profiler/profiler.py:606`), and encapsulate the +metrics refresh on `DQProfileContext` per user feedback (do not sprinkle +`model_copy(update={...})` at the call site). + +- In `src/databricks/labs/dqx/profiler/semantic.py`, add a new method + on `DQProfileContext`: + ```python + def with_metrics(self, metrics: Mapping[str, Any]) -> "DQProfileContext": + """Return a new context whose `metrics` is a fresh snapshot of *metrics*. + ... + """ + return self.model_copy(update={"metrics": dict(metrics)}) + ``` + Docstring should explain the intended use: because the model is + frozen and pydantic materialises `metrics` as a copy on construction, + contextual builders further down the chain would otherwise not see + write-backs made to the outer mutable `metrics` dict between builder + invocations. `with_metrics` produces a new frozen context wrapping + the current dict snapshot. +- In `src/databricks/labs/dqx/profiler/profiler.py` + (`DQProfiler._build_profiles_for_column`), keep the initial + `builder_ctx` construction, but inside the + `for profile_type in PROFILE_BUILDER_REGISTRY.values():` loop refresh + it before each contextual invocation: + ```python + if profile_type.contextual_builder is not None: + builder_ctx = builder_ctx.with_metrics(metrics) + profile = profile_type.contextual_builder(builder_ctx) + ``` + Leave the legacy-callback branch unchanged (it already receives + `dict(metrics)` per call). +- Add a short reference comment above the refresh line pointing to the + min/max write-back block below (~lines 631–635) so future readers + understand why the refresh is required. +- Do not change per-column initial detector `detector_ctx` construction + in `_detect_semantic_type`. +- Add unit coverage in `tests/unit/profiler/test_semantic.py` for the new + method: assert `with_metrics(new_dict)` returns a new instance with the + snapshot semantics, leaves the original untouched, preserves all other + fields (`df`, `column_name`, `column_type`, `options`, `metadata`, + `semantic_type`), and rejects post-construction mutation of the + original's fields (Pydantic frozen behaviour). + +**Feedback** + + +### Step - 3 + +**Description** + +Align the enum cardinality gate with the legacy `is_in` distinct_ratio flagged +in comment #3 (`src/databricks/labs/dqx/profiler/semantic.py:213`). + +- In `_detect_enum`, replace the standalone `ENUM_MAX_CARDINALITY_RATIO` + gate with a lookup against + `PROFILE_OPTION_DISTINCT_RATIO` from `ctx.options`, falling back to + `ENUM_MAX_CARDINALITY_RATIO` when the option is absent. Use strict `>=` + (mirroring the legacy `is_in` builder's `distinct_ratio < max_distinct_ratio` + emission gate) so semantic profiling and legacy profiling agree on which + columns qualify. +- Lower the default `ENUM_MAX_CARDINALITY_RATIO` in `semantic.py` to `0.05` + so the constant reflects the effective default when options are not merged + (e.g. when semantic detectors are used outside `DQProfiler`, or when a caller + passes `options={}` explicitly). Update the module-level comment that + currently justifies 0.95 as "more permissive" so it explains the alignment + with `PROFILE_OPTION_DISTINCT_RATIO`. +- Import `PROFILE_OPTION_DISTINCT_RATIO` from `profile_options.py` in + `semantic.py` (already imports `PROFILE_OPTION_MAX_IN_COUNT`, so add + alongside). +- No API breakage: `ENUM_MAX_CARDINALITY_RATIO` stays exported but its value + and semantics change — call out in the plan feedback if downstream users + depend on the old 0.95 value. + +**Feedback** + + +### Step - 4 + +**Description** + +Broaden the UC metadata-fetch exception handler flagged in comment #4 +(`src/databricks/labs/dqx/profiler/profiler.py:462`). + +- In `DQProfiler._fetch_table_metadata`, widen the `except DatabricksError` + clause to `except Exception` (bare, not `BaseException`, so `KeyboardInterrupt` + and `SystemExit` still propagate). This honours the docstring's "metadata + errors never block profiling" guarantee for client-side failures like a + `ValueError` raised by SDK name validation on non-UC locations + (e.g. storage paths). +- Keep the log message shape unchanged (`safe_location`, newline scrubbing), + and keep the same "return `{}`" behaviour. +- Remove the now-unused `DatabricksError` import at the top of `profiler.py` + only if no other code path still uses it (verify with grep before deleting). + +**Feedback** + + +### Step - 5 + +**Description** + +Fix the two documentation nits flagged in comments #5 and #6 +(`docs/dqx/docs/reference/profiler.mdx`). + +- Line 271 (prose): replace `an optional \`type="context"\` argument` with + `an optional \`kind="context"\` argument`. +- Line 280 (code block): replace + `@register_profile_builder("my_custom", type="context")` with + `@register_profile_builder("my_custom", kind="context")` so the example + matches the actual decorator signature in + `src/databricks/labs/dqx/profiler/profile_builder.py:46`. +- Line 147: replace `The feature is fully opt-in: it runs iff a`… + with the suggested wording + `The feature is fully opt-in: it runs if a \`semantic_registry\` is + supplied to the`… (accept the reviewer's `suggestion` block verbatim). + +**Feedback** + + +### Step - 6 + +**Description** + +Redesign the `SemanticRegistry` customization surface per comments #7, #8, +and #10 points 1–2 (`src/databricks/labs/dqx/profiler/semantic.py:391`). +Comment #9 and comment #10 point 3 are explicitly out of scope — do not +convert semantic detectors to a decorator/global-registry pattern. + +**Do NOT add a `with_detectors(...)` method** (per user feedback: chaining +`SemanticRegistry.default().with_detectors([...])` reads nonsensically because +the default chain is constructed and then discarded). Whole-chain construction +is served by the existing constructor and a new classmethod for ergonomics +(see below). + +- **Repurpose `replace`.** Change the existing `replace(detectors: Sequence[...])` + method (whole-chain) into a **single-detector swap** with signature + `replace(name: str, detector: DQSemanticTypeDetector) -> "SemanticRegistry"`. + It swaps the entry named `name` with `detector`, preserving position, and + raises `ValueError` when `name` is not present (mirrors the reviewer's + expected `replace(name, detector)` semantics). All call sites that used + the old whole-chain `replace([...])` must migrate to the constructor or the + new `of(...)` classmethod added below — this is covered in Step 7. + +- **Add the following new methods** to `SemanticRegistry`, each returning a new + `SemanticRegistry` instance and routing through the constructor so the + `_validate_unique_names` model-validator runs on the derived instance + (matching the existing `prepend` pattern): + * `append(detector: DQSemanticTypeDetector) -> "SemanticRegistry"` — + lowest-priority fallback: places `detector` at the end of the chain. + * `insert(name: str, detector: DQSemanticTypeDetector) -> "SemanticRegistry"` — + (renamed from `insert_after` per user feedback) inserts `detector` + immediately after the entry whose name equals `name`. Raise `ValueError` + when `name` is not present. The docstring must state the "after" semantics + explicitly so the terser name is unambiguous ("insert this detector after + the one identified by `name`"). + * `remove(name: str) -> "SemanticRegistry"` — returns a new registry with + the detector matching `name` filtered out. Raise `ValueError` when `name` + is not present, to keep the API strict (consistent with the uniqueness + invariant). + +- **Add a classmethod for whole-chain construction from scratch** (per user + feedback: prefer a static/class method over the previously-planned + `with_detectors` instance method). Signature: + ```python + @classmethod + def of(cls, *detectors: DQSemanticTypeDetector) -> "SemanticRegistry": + """Return a registry whose chain is exactly *detectors* (in argument order).""" + return cls(detectors=tuple(detectors)) + ``` + Rationale: varargs reads well when combining with `default_semantic_detectors()` + (e.g. `SemanticRegistry.of(uuid_detector, *default_semantic_detectors())`) and + it complements the existing `.default()` classmethod symmetrically. The + existing constructor `SemanticRegistry(detectors=(...))` continues to work + for callers that want the kwarg form. + +- Add Google-style docstrings on every new/changed method (one paragraph each, + matching the existing tone). Keep `prepend(detector)` unchanged — it is + still the correct primitive for "make mine win first." + +- Update the module docstring header at the top of `semantic.py` so the + registry summary lists the new/renamed methods: `of` (classmethod), + `default` (classmethod, unchanged), `prepend`, `append`, `insert`, + `replace(name, detector)`, `remove`. + +**Feedback** + + +### Step - 7 + +**Description** + +Update the docs and tests to match the new `SemanticRegistry` surface and the +Step 5 fixes. + +- `docs/dqx/docs/reference/profiler.mdx` + * In the "Composing a registry" code block (starting ~line 157), replace the + single `replace(...)` example with usages of the new methods. Illustrate + whole-chain construction via `SemanticRegistry.of(...)` (the new + classmethod) — do NOT show `SemanticRegistry.default().with_detectors([...])` + (that method does not exist per the revised Step 6). Examples to include: + * `SemanticRegistry.default().prepend(uuid_detector)` — highest priority + * `SemanticRegistry.default().append(uuid_detector)` — fallback slot + * `SemanticRegistry.default().insert("enum", uuid_detector)` — after + the enum entry + * `SemanticRegistry.default().replace("enum", custom_enum)` — swap one + * `SemanticRegistry.default().remove("text")` — drop a default + * `SemanticRegistry.of(uuid_detector, *default_semantic_detectors())` — + whole chain from scratch + * `SemanticRegistry()` — empty chain + * Add the comment block requested in comment #7: + ``` + # prepend(detector) → puts one detector at the front (highest priority) + # append(detector) → adds one detector at the back (lowest priority) + # insert(name, detector) → inserts one detector after the named entry + # replace(name, detector) → swaps the entry named `name` + # remove(name) → drops the entry named `name` + # SemanticRegistry.of(...) → build a registry with an explicit chain from scratch + ``` + Place it directly above the code block. + * Rewrite the prose paragraph at lines ~186–189 so it describes all six + composition methods plus the whole-chain constructor/`of(...)` path, + keeps the uniqueness-invariant note, and drops the `replace`-means-whole-chain + wording that prompted comment #8. Explicitly note that `replace` now + swaps a single named detector. + * Regenerate the enum applicability row in the "Built-in detectors" table + (~line 195) so the threshold description matches Step 3: replace the + `ENUM_MAX_CARDINALITY_RATIO (default 0.95)` clause with the aligned + `distinct_ratio` option (default 0.05) and drop the outdated "0.95" number. + +- `tests/unit/profiler/test_semantic.py` + * Delete the existing whole-chain tests that exercised the old `replace([...])` + semantics (`test_semantic_registry_replace_returns_new_instance` and + `test_semantic_registry_replace_duplicate_names_raises`) — the whole-chain + entry point is now `SemanticRegistry(...)` / `SemanticRegistry.of(...)`, + which is already covered by + `test_semantic_registry_direct_constructor_duplicate_names_raises`. + Replace them with: + * `test_semantic_registry_of_returns_new_instance` — asserts + `SemanticRegistry.of(d1, d2)` yields a registry with the given ordered + chain and is a fresh instance. + * `test_semantic_registry_of_duplicate_names_raises` — asserts varargs + duplicates trip the uniqueness validator. + * Add positive + negative tests for each new/renamed method: `append`, + `insert`, `remove(name)`, and `replace(name, detector)`. Cover: + * Ordering — append lands last; insert(name, ...) lands at position n+1 + relative to the named entry; replace(name, ...) preserves position. + * Immutability — original registry unchanged. + * Uniqueness — `append`/`insert`/`replace(name, detector)` reject a + detector whose name clashes with a non-target existing entry (raises + `pydantic.ValidationError`); `replace(name, detector)` allows the new + detector to reuse the target's own name. + * Missing-name errors — `remove`, `insert`, and `replace(name, detector)` + raise `ValueError` for an unknown `name`. + * Adjust the enum-detector positive test (`test_enum_detector_positive_low_cardinality`) + so its `options` dict provides `distinct_ratio` compatible with the new + gate (Step 3). For 3 distinct values / 300 non-null rows, ratio ≈ 0.01, + which passes the default 0.05 threshold. + * Adjust `test_enum_detector_rejects_all_distinct_integer_column` so the + new gate still causes rejection (already does under the default threshold; + verify). + * Add a new test `test_enum_detector_respects_distinct_ratio_option` that + verifies the enum detector suppresses classification when + `distinct_ratio` option is set below the observed ratio (mirrors the + reviewer's 12-rows / 9-distinct example from comment #3). + +- `tests/integration/test_profile_semantic.py` + * Migrate `test_enum_detector_values_reused_by_is_in_builder` (currently + uses the old whole-chain `.replace([...])`) to + `SemanticRegistry.of(tracing, *default_semantic_detectors()[1:])` + (or the plain constructor `SemanticRegistry(detectors=(...))`). + * Grep for any other whole-chain `.replace([...])` call sites in this file + and migrate them the same way. Also grep for `.insert_after(` — should + yield no hits (never released), but confirm and rename if any exist. + * `_make_demo_df` produces `vehicle_type` with 3 distinct values over 60 + rows (ratio 0.05). Under the new default threshold (`distinct_ratio = + 0.05`, strict `>=` rejects), verify semantic-enum still fires by either + (a) tightening the fixture to make cardinality/count_non_null strictly + less than 0.05 (e.g. 3/61 = 0.049) — safer path — or (b) explicitly + passing `options={"distinct_ratio": 0.1, ...}` in the test call. Prefer + (b) to keep the fixture stable. + +- `tests/perf/test_profile_semantic.py` + * Grep for `.replace([`, `.insert_after(`, and `.with_detectors(`. If any + benchmark call uses the old whole-chain `replace`, migrate to + `SemanticRegistry.of(...)` or the constructor. Otherwise no change. + +- `tests/unit/test_profile_builder.py` + * No API change needed here — the profile-builder tests use + `DQProfileContext` directly. Only touch if Step 2's `with_metrics` + change surfaces a test failure (unlikely — no test currently reads + mutated `metrics` through `ctx.metrics`). + +**Feedback** \ No newline at end of file diff --git a/.agents/2026-09-05-12-48-53-review-round-0.md b/.agents/2026-09-05-12-48-53-review-round-0.md new file mode 100644 index 000000000..f55316747 --- /dev/null +++ b/.agents/2026-09-05-12-48-53-review-round-0.md @@ -0,0 +1,50 @@ +--- +title: The implementation largely matches the plan, but the broadened metadata failure path logs unsanitized exception content and the metrics-refresh fix lacks behavioral regression coverage. +status: changes-requested +created_at: 2026-09-05T13:35:38+02:00 +plan: .agents/2026-09-05-12-48-53-plan.md + +summary: > + The registry API, enum threshold alignment, ShortType support, documentation updates, + and associated tests substantially implement the plan. Follow-up is required because + the broadened exception handler can place uncontrolled exception text into logs, the + central builder-chain behavior fixed by with_metrics is not directly regression-tested, + and a new test suppresses a type-checking error contrary to repository policy. +--- + +# Review + +### Finding - 1 +**Location** +`src/databricks/labs/dqx/profiler/profiler.py:461-466` +**Severity** +critical +**Description** +The handler now catches every Exception raised by the SDK boundary but interpolates the raw exception text into the warning. Only the location is stripped of newlines. Exception messages can contain attacker-controlled table identifiers with newlines or control characters, as well as SDK or transport details that should not be exposed in logs. This violates the repository's log-injection and sensitive-content requirements, and widening the caught exception set increases the range of uncontrolled messages reaching this sink. Sanitize the exception text with the shared log-safe helper or omit it from the message and log only a bounded, non-sensitive exception type through structured fields. +**Address** +Y/N +**Comments** + +### Finding - 2 +**Location** +`src/databricks/labs/dqx/profiler/profiler.py:618-641` +**Severity** +modest +**Description** +The tests added for with_metrics establish only the value object's copy behavior; they do not exercise the bug fixed in the profiler loop. A future change could remove or misplace the refresh at line 623 while every new test continues to pass. Add a behavioral test through DQProfiler's public profiling API with a contextual builder ordered after min_max and assert that it observes the resolved min and max values written back by min_max. This directly guards the reviewer-reported failure and follows the project's requirement to test observable behavior rather than helper implementation details. +**Address** +Y +**Comments** +Added `test_contextual_builder_after_min_max_observes_resolved_min_max` in `tests/integration/test_profile_semantic.py`. The test registers a spy contextual builder via `@register_profile_builder(kind="context")` — insertion order in `PROFILE_BUILDER_REGISTRY` guarantees it runs after `min_max` — and asserts the spy's `ctx.metrics` carries the resolved `min`/`max` written back by `min_max`. Registry cleanup happens in a `finally` block. Runs with `remove_outliers=False` so the expected values are exact. + +### Finding - 3 +**Location** +`tests/unit/profiler/test_semantic.py:96-103` +**Severity** +modest +**Description** +The new frozen-context test adds a type-ignore directive to permit an assignment that the type checker correctly rejects. The repository explicitly prohibits adding type-ignore comments to silence lint or type errors. The same public behavior can be tested without suppressing diagnostics by performing the assignment through setattr and asserting that Pydantic raises ValidationError. +**Address** +Y +**Comments** +Replaced `ctx.column_name = "x" # type: ignore[misc]` with `setattr(ctx, "column_name", "x")` in `test_dq_profile_context_frozen_blocks_field_reassignment`. Public behavior (Pydantic raises `ValidationError` on frozen field reassignment) is unchanged; the type-checker suppression is gone. diff --git a/docs/dqx/docs/reference/profiler.mdx b/docs/dqx/docs/reference/profiler.mdx index 03aed5593..1c1e9a8c8 100644 --- a/docs/dqx/docs/reference/profiler.mdx +++ b/docs/dqx/docs/reference/profiler.mdx @@ -55,7 +55,7 @@ The profiler supports extensive configuration options to customize behavior: | --------------------------- | ------------- |--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | `round` | `True` | Round min/max values for cleaner rules | | `max_in_count` | `10` | Generate `is_in` rule if distinct values < this count | -| `distinct_ratio` | `0.05` | Generate `is_in` rule if distinct values < 5% of total | +| `distinct_ratio` | `0.05` | Generate `is_in` rule if distinct values < 5% of non-null values | | `max_null_ratio` | `0.01` | Generate `is_not_null` rule if null values < 1% of total | | `remove_outliers` | `True` | Enable outlier detection for min/max rules | | `outlier_columns` | `[]` | Specific columns for outlier detection (empty = all numeric) | @@ -144,7 +144,7 @@ checks. For example, a *vehicle_type* column classified as an *enum* now receive *is_in* candidate only — not both *is_in* and *min_max* — while a continuous *cargo_weight* column is classified as a *measurement* and receives *min_max* only. -The feature is fully opt-in: it runs iff a `semantic_registry` is supplied to the +The feature is fully opt-in: it runs if a `semantic_registry` is supplied to the `DQProfiler` constructor. When no registry is supplied, the profiler output is byte-identical to the pre-feature behaviour and every `DQProfile.semantic_type` field is `None`. @@ -152,7 +152,16 @@ field is `None`. ### Composing a registry `SemanticRegistry` is an immutable, ordered, name-unique collection of detectors. -Its four composition patterns: +The composition methods (all return a new registry, leaving the original untouched): + +```text +prepend(detector) → puts one detector at the front (highest priority) +append(detector) → adds one detector at the back (lowest priority) +insert(name, detector) → inserts one detector after the named entry +replace(name, detector) → swaps the entry named `name` +remove(name) → drops the entry named `name` +SemanticRegistry.of(...) → build a registry with an explicit chain from scratch +``` ```python from databricks.labs.dqx.profiler.profiler import DQProfiler @@ -174,25 +183,41 @@ def _detect_uuid(ctx): uuid_detector = DQSemanticTypeDetector(name="uuid", detect=_detect_uuid) prepended = SemanticRegistry.default().prepend(uuid_detector) -# 3. Replace the chain entirely -custom = SemanticRegistry.default().replace([uuid_detector, *default_semantic_detectors()]) +# 3. Append as a lowest-priority fallback +appended = SemanticRegistry.default().append(uuid_detector) + +# 4. Insert immediately after the built-in `enum` detector +inserted = SemanticRegistry.default().insert("enum", uuid_detector) + +# 5. Swap a single entry by name (position preserved) +custom_enum_detector = DQSemanticTypeDetector(name="enum", detect=_detect_uuid) +swapped = SemanticRegistry.default().replace("enum", custom_enum_detector) + +# 6. Drop a built-in detector +without_text = SemanticRegistry.default().remove("text") + +# 7. Build an arbitrary chain from scratch +custom = SemanticRegistry.of(uuid_detector, *default_semantic_detectors()) -# 4. Empty registry — semantic types are always None +# 8. Empty registry — semantic types are always None empty = SemanticRegistry() profiler = DQProfiler(ws, semantic_registry=default_registry) ``` -Both `prepend(...)` and `replace(...)` return a **new** `SemanticRegistry` and -route through the constructor so name-uniqueness is enforced on the derived -instance. Attempting to prepend a detector whose name clashes with an existing -one raises `pydantic.ValidationError`. +Every composition method returns a **new** `SemanticRegistry` and routes through +the constructor so name-uniqueness is enforced on the derived instance. Adding +a detector whose name clashes with an existing entry raises +`pydantic.ValidationError`; `insert`, `replace`, and `remove` raise `ValueError` +when the referenced `name` is not present. To construct an arbitrary chain from +scratch, use `SemanticRegistry.of(...)` (varargs) or the +`SemanticRegistry(detectors=(...))` constructor directly. ### Built-in detectors | Detector | Applicability | Produced `properties` | |---------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------| -| `enum` | String or integer columns where `cardinality < max_in_count` AND `cardinality / count_non_null ≤ ENUM_MAX_CARDINALITY_RATIO` (default 0.95). Collects distinct values once and hands them downstream. | `values`: tuple of distinct values | +| `enum` | String or integer columns (including `ShortType`) where `cardinality < max_in_count` AND `cardinality / count_non_null < distinct_ratio` (default 0.05, taken from the `distinct_ratio` profiler option so semantic profiling and the legacy `is_in` builder agree). Collects distinct values once and hands them downstream. | `values`: tuple of distinct values | | `key` | Any column with `distinctness ≥ 0.99` AND a type-specific second signal — numeric density ≥ `KEY_MIN_DENSITY_RATIO` (default 0.99) or string length_stability ≥ `KEY_MIN_LENGTH_STABILITY_RATIO` (default 0.95). | `signal`: `"density"` or `"length_stability"` | | `measurement` | Any numeric column not consumed by an earlier detector. | `distribution`: best-effort guess (`normal` / `uniform` / …) | | `text` | Any string column not consumed by an earlier detector. | (none) | @@ -240,23 +265,6 @@ def _detect_uuid(ctx): uuid_detector = DQSemanticTypeDetector(name="uuid", detect=_detect_uuid) ``` -### Table metadata - -When `profile_table(...)` is used, DQX best-effort fetches table and per-column -metadata from Unity Catalog and exposes it to detectors via `ctx.metadata`. -The following fields are populated when available: - -* *table_name* -* *table_comment* -* *column_comment* - -Tags are not surfaced in this release — they will follow in a subsequent -release. Metadata fetches are best-effort and degrade gracefully: any SDK -failure (missing table, permission denied, non-UC location) is logged as a -warning and treated as an empty mapping, so metadata errors never block -profiling. On the raw-DataFrame path (`.profile(df, ...)`) there is no table -origin, so `ctx.metadata` is always `{}`. - ### Traceability via `DQProfile.semantic_type` When a detector fires, the resulting `DQProfile` records the semantic type name @@ -268,16 +276,16 @@ on every profile. Existing user-authored builders registered via `@register_profile_builder(...)` continue to work unchanged. The decorator now also accepts an optional -`type="context"` argument to opt into the preferred contextual callback shape, +`kind="context"` argument to opt into the preferred contextual callback shape, which receives a `DQProfileContext` (giving access to *ctx.semantic_type*, -*ctx.metadata*, *ctx.metrics*, and *ctx.options*): +*ctx.metrics*, and *ctx.options*): ```python from databricks.labs.dqx.profiler.profile import DQProfile from databricks.labs.dqx.profiler.profile_builder import register_profile_builder from databricks.labs.dqx.profiler.semantic import DQProfileContext -@register_profile_builder("my_custom", type="context") +@register_profile_builder("my_custom", kind="context") def _my_custom_builder(ctx: DQProfileContext) -> DQProfile | None: if ctx.semantic_type and ctx.semantic_type.name != "measurement": return None diff --git a/src/databricks/labs/dqx/profiler/profile_builder.py b/src/databricks/labs/dqx/profiler/profile_builder.py index 01f5a61f8..c8fe36979 100644 --- a/src/databricks/labs/dqx/profiler/profile_builder.py +++ b/src/databricks/labs/dqx/profiler/profile_builder.py @@ -111,8 +111,8 @@ def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: if not _supports_distinct(ctx.column_type): return None - total_count = ctx.metrics.get("count", 0) - if total_count == 0: + count_non_null = ctx.metrics.get("count_non_null", 0) + if count_non_null == 0: return None semantic_type = ctx.semantic_type @@ -144,7 +144,10 @@ def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: # the column is entirely null — no valid values to build an allowlist from. return None - distinct_ratio = (1.0 * distinct_count) / total_count + # Denominator is *count_non_null* so this ratio matches *_detect_enum* in semantic.py. + # A mismatch would let semantic-enum classification and legacy is_in emission disagree on + # low-repetition columns with heavy nulls. + distinct_ratio = (1.0 * distinct_count) / count_non_null if distinct_count < max_in_count and distinct_ratio < max_distinct_ratio: return DQProfile( @@ -316,13 +319,17 @@ def _supports_distinct(column_type: T.DataType) -> bool: """ Validates that the input column type supports distinct operations. + The accepted set (text plus *IntegerType*/*LongType*/*ShortType*) must stay in sync with the + semantic enum detector in *semantic._detect_enum* — a mismatch would let semantic-enum + classification suppress *min_max* without producing an *is_in* in return. + Args: column_type: Input column type Returns: True if the column supports distinct operations, otherwise False """ - return isinstance(column_type, (T.IntegerType, T.LongType) + TEXT_TYPES) + return isinstance(column_type, (T.IntegerType, T.LongType, T.ShortType) + TEXT_TYPES) def _supports_min_max(column_type: T.DataType) -> bool: diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index d919706f5..d1913ea7e 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -2,7 +2,6 @@ import uuid import logging import os -from collections.abc import Mapping from concurrent import futures from decimal import Decimal, Context from difflib import SequenceMatcher @@ -14,7 +13,6 @@ from pyspark.errors import AnalysisException from pyspark.sql import DataFrame, SparkSession from databricks.sdk import WorkspaceClient -from databricks.sdk.errors import DatabricksError from databricks.labs.dqx.base import DQEngineBase from databricks.labs.dqx.config import InputConfig, LLMModelConfig @@ -110,24 +108,15 @@ def profile( A tuple containing a dictionary of summary statistics and a list of data quality profiles. """ - return self._profile_dataframe(df, columns, options, table_metadata=None) + return self._profile_dataframe(df, columns, options) def _profile_dataframe( self, df: DataFrame, columns: list[str] | None, options: dict[str, Any] | None, - *, - table_metadata: Mapping[str, Any] | None = None, ) -> tuple[dict[str, Any], list[DQProfile]]: - """Shared private entry point for *.profile()* and *.profile_table()*. - - Accepts an optional *table_metadata* dict fetched from Unity Catalog (see - *_fetch_table_metadata*) that is threaded down to *_profile* so per-column - *DQProfileContext* instances get the corresponding *metadata* mapping. The - public *.profile()* signature stays unchanged; only the internal path carries - table metadata. - """ + """Shared private entry point for *.profile()* and *.profile_table()*.""" columns = columns or df.columns df_columns = [f for f in df.schema.fields if f.name in columns] df = df.select(*[f.name for f in df_columns]) @@ -145,7 +134,7 @@ def _profile_dataframe( if total_count == 0: return summary_stats, dq_rules - self._profile(df, df_columns, dq_rules, options, summary_stats, total_count, table_metadata=table_metadata) + self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) return summary_stats, dq_rules @@ -172,8 +161,7 @@ def profile_table( logger.info(f"Profiling {input_config.location} with options: {options}") df = read_input_data(spark=self.spark, input_config=input_config) - table_metadata = self._fetch_table_metadata(input_config.location) - return self._profile_dataframe(df=df, columns=columns, options=options, table_metadata=table_metadata) + return self._profile_dataframe(df=df, columns=columns, options=options) @telemetry_logger("profiler", "profile_tables_for_patterns") def profile_tables_for_patterns( @@ -446,41 +434,6 @@ def _stratified_sample( logger.info(f"Stratified sampling on column '{sample_by_column}'") return df.sampleBy(sample_by_column, fractions=sample_fractions, seed=sample_seed) - def _fetch_table_metadata(self, location: str) -> dict[str, Any]: - """Best-effort fetch of Unity Catalog table + column metadata. - - Returns a mapping with keys *table_name*, *table_comment*, *columns* - (a dict keyed by column name, each carrying *column_comment*). Tags - are intentionally excluded — the SDK exposes them via separate - calls and current use cases only need the comment fields. Any SDK - failure (missing table, permission denied, non-UC location) is - logged as a warning and returns an empty dict, so metadata errors - never block profiling. - """ - try: - table = self.ws.tables.get(location) - except DatabricksError as exc: - safe_location = str(location).replace("\n", " ").replace("\r", " ") - logger.warning(f"Could not fetch table metadata for {safe_location}: {exc}") - return {} - - columns: dict[str, dict[str, Any]] = {} - for col in table.columns or []: - if col.name is None: - continue - col_entry: dict[str, Any] = {} - if col.comment is not None: - col_entry["column_comment"] = col.comment - if col_entry: - columns[col.name] = col_entry - - metadata: dict[str, Any] = {"table_name": table.full_name or location} - if table.comment is not None: - metadata["table_comment"] = table.comment - if columns: - metadata["columns"] = columns - return metadata - def _profile( self, df: DataFrame, @@ -489,8 +442,6 @@ def _profile( opts: dict[str, Any], summary_stats: dict[str, Any], total_count: int, - *, - table_metadata: Mapping[str, Any] | None = None, ) -> None: """ Builds a list of DQProfiles by iterating through DQProfileBuilder builders. @@ -506,8 +457,6 @@ def _profile( opts: Dictionary of options for profiling. summary_stats: Summary statistics dictionary to update with profiler results. total_count: Total number of rows in the input DataFrame. - table_metadata: Table metadata fetched from Unity Catalog when available; - surfaced through *ctx.metadata* on each per-column *DQProfileContext*. """ trim_strings = opts.get(PROFILE_OPTION_TRIM_STRINGS, True) @@ -538,8 +487,6 @@ def _profile( else: metrics["empty_count"] = 0 - column_metadata = self._build_column_metadata(field_name, table_metadata) - self._build_profiles_for_column( column_df, field_name, @@ -547,36 +494,10 @@ def _profile( metrics, opts, dq_rules, - column_metadata=column_metadata, ) self._add_llm_primary_key_for_dataframe(df, dq_rules, summary_stats, opts) - @staticmethod - def _build_column_metadata(field_name: str, table_metadata: Mapping[str, Any] | None) -> dict[str, Any]: - """Compose the per-column *metadata* mapping from table-level metadata. - - Fields whose source value is *None* are omitted so detectors can - *metadata.get(...)* cleanly. Returns an empty dict when no table - metadata was supplied. - """ - if not table_metadata: - return {} - metadata: dict[str, Any] = {} - table_name = table_metadata.get("table_name") - if table_name is not None: - metadata["table_name"] = table_name - table_comment = table_metadata.get("table_comment") - if table_comment is not None: - metadata["table_comment"] = table_comment - columns = table_metadata.get("columns") or {} - column_entry = columns.get(field_name) if isinstance(columns, Mapping) else None - if column_entry: - column_comment = column_entry.get("column_comment") - if column_comment is not None: - metadata["column_comment"] = column_comment - return metadata - def _build_profiles_for_column( self, column_df: DataFrame, @@ -585,8 +506,6 @@ def _build_profiles_for_column( metrics: dict[str, Any], opts: dict[str, Any], dq_rules: list[DQProfile], - *, - column_metadata: Mapping[str, Any] | None = None, ) -> None: """Run registered profile builders for a column and append profiles. @@ -600,8 +519,7 @@ def _build_profiles_for_column( *metrics* so that downstream consumers (e.g. LLM primary-key detection) can read them without triggering a second Spark action. """ - metadata_map = dict(column_metadata) if column_metadata else {} - semantic_type = self._detect_semantic_type(column_df, field_name, field_type, metrics, opts, metadata_map) + semantic_type = self._detect_semantic_type(column_df, field_name, field_type, metrics, opts) builder_ctx = DQProfileContext( df=column_df, @@ -609,12 +527,15 @@ def _build_profiles_for_column( column_type=field_type, metrics=metrics, options=opts, - metadata=metadata_map, semantic_type=semantic_type, ) for profile_type in PROFILE_BUILDER_REGISTRY.values(): if profile_type.contextual_builder is not None: + # Refresh the frozen context with the current *metrics* snapshot so contextual + # builders registered after *min_max* observe the resolved min/max values written + # back below (see the write-back block after this loop). + builder_ctx = builder_ctx.with_metrics(metrics) profile = profile_type.contextual_builder(builder_ctx) elif profile_type.builder is not None: profile = profile_type.builder(column_df, field_name, field_type, dict(metrics), dict(opts)) @@ -641,7 +562,6 @@ def _detect_semantic_type( field_type: T.DataType, metrics: dict[str, Any], opts: dict[str, Any], - metadata_map: dict[str, Any], ) -> DQSemanticType | None: if self._semantic_registry is None: return None @@ -651,7 +571,6 @@ def _detect_semantic_type( column_type=field_type, metrics=metrics, options=opts, - metadata=metadata_map, semantic_type=None, ) for detector in self._semantic_registry.detectors: diff --git a/src/databricks/labs/dqx/profiler/semantic.py b/src/databricks/labs/dqx/profiler/semantic.py index 02933864f..f5a532a9c 100644 --- a/src/databricks/labs/dqx/profiler/semantic.py +++ b/src/databricks/labs/dqx/profiler/semantic.py @@ -15,6 +15,10 @@ * models: *DQSemanticType*, *DQProfileContext*, *DQSemanticTypeDetector* * properties: *DQSemanticTypeProperties*, *EnumProperties*, *MeasurementProperties* * registry: *SemanticRegistry* + — construction: *SemanticRegistry.default()*, *SemanticRegistry.of(...)*, + constructor *SemanticRegistry(detectors=(...))* + — composition: *prepend*, *append*, *insert(name, detector)*, + *replace(name, detector)*, *remove(name)* * detectors: *DEFAULT_ENUM_DETECTOR*, *DEFAULT_KEY_DETECTOR*, *DEFAULT_MEASUREMENT_DETECTOR*, *DEFAULT_TEXT_DETECTOR*, *default_semantic_detectors()* @@ -23,7 +27,7 @@ """ import logging -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Mapping from typing import Any, Literal, TypeVar from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -33,7 +37,7 @@ from pyspark.sql.types import DataType from databricks.labs.dqx.profiler.common import TEXT_TYPES -from databricks.labs.dqx.profiler.profile_options import PROFILE_OPTION_MAX_IN_COUNT +from databricks.labs.dqx.profiler.profile_options import PROFILE_OPTION_DISTINCT_RATIO, PROFILE_OPTION_MAX_IN_COUNT logger = logging.getLogger(__name__) @@ -123,11 +127,6 @@ class DQProfileContext(BaseModel): count_non_null, ...). Same shape as *summary_stats[column_name]*. options: Profiler options for this run (max_null_ratio, max_empty_ratio, max_in_count, trim_strings, filter, ...). - metadata: Table/column metadata sourced from Unity Catalog when the - profile is derived from a Delta table. Keys surfaced: - *table_name*, *table_comment*, *column_comment*. Tags are - intentionally excluded. The mapping is empty when profiling a - raw DataFrame with no table origin. semantic_type: Detected semantic type for this column, or *None* if no detector matched. Populated only for profile builders — always *None* for semantic detectors. @@ -140,9 +139,27 @@ class DQProfileContext(BaseModel): column_type: DataType metrics: Mapping[str, Any] = Field(default_factory=dict) options: Mapping[str, Any] = Field(default_factory=dict) - metadata: Mapping[str, Any] = Field(default_factory=dict) semantic_type: DQSemanticType | None = None + def with_metrics(self, metrics: Mapping[str, Any]) -> "DQProfileContext": + """Return a new context whose *metrics* is a fresh snapshot of the supplied mapping. + + The model is frozen and Pydantic materializes *metrics* as a copy on construction, so + contextual profile builders further down the chain would otherwise not observe write-backs + made to the outer mutable metrics dict between builder invocations (e.g. the min/max + write-back performed by *DQProfiler._build_profiles_for_column* after the *min_max* + builder resolves outlier-adjusted bounds). Callers pass the current mutable dict here to + obtain a new frozen context that wraps its snapshot; all other fields are preserved. + + Args: + metrics: Column-level metrics to snapshot into the returned context. + + Returns: + A new *DQProfileContext* instance with *metrics* replaced by a fresh dict copy of + the supplied mapping. + """ + return self.model_copy(update={"metrics": dict(metrics)}) + class DQSemanticTypeDetector(BaseModel): """Named detector that classifies a column's semantic type. @@ -160,12 +177,13 @@ class DQSemanticTypeDetector(BaseModel): detect: Callable[[DQProfileContext], DQSemanticType | None] -# Enum: cardinality / count_non_null must be at or below this ratio for a -# column to be classified as *enum* (in addition to -# cardinality < max_in_count). A higher ratio ceiling is more permissive -# — it lets slightly less-repeated columns still register as enum-like -# while still rejecting near-unique identifier columns. -ENUM_MAX_CARDINALITY_RATIO = 0.95 +# Enum: fallback ceiling for cardinality / count_non_null used only when +# *PROFILE_OPTION_DISTINCT_RATIO* is absent from *ctx.options*. Kept aligned +# with the *is_in* builder's default *distinct_ratio* so semantic profiling +# and legacy profiling agree on which columns qualify as enum-like. The +# effective gate is strict *>=* (mirrors *is_in*'s *distinct_ratio < +# max_distinct_ratio* emission condition). +ENUM_MAX_CARDINALITY_RATIO = 0.05 # Key (numeric): count_distinct / (max - min + 1) must be at or above this # ratio for a numeric column to be classified as *key*. Tolerates small @@ -210,7 +228,8 @@ def _detect_enum(ctx: DQProfileContext) -> DQSemanticType | None: if cardinality >= max_in_count: return None - if (cardinality / count_non_null) > ENUM_MAX_CARDINALITY_RATIO: + distinct_ratio_threshold = ctx.options.get(PROFILE_OPTION_DISTINCT_RATIO, ENUM_MAX_CARDINALITY_RATIO) + if (cardinality / count_non_null) >= distinct_ratio_threshold: return None col = ctx.df.columns[0] @@ -361,13 +380,15 @@ class SemanticRegistry(BaseModel): The registry is immutable — the model is frozen (Pydantic rejects field assignment) and the *detectors* field is a tuple, so external callers cannot mutate the chain by aliasing. All construction paths (direct - instantiation, *default()*, *prepend()*, *replace()*) route through the - model validator, which enforces detector-name uniqueness. + instantiation, *of()*, *default()*, *prepend()*, *append()*, *insert()*, + *replace()*, *remove()*) route through the model validator, which enforces + detector-name uniqueness. Attributes: detectors: Ordered tuple of detectors in first-match-wins order. Defaults to an empty tuple; use *SemanticRegistry.default()* to - obtain the built-in chain. + obtain the built-in chain or *SemanticRegistry.of(...)* to build + an arbitrary chain from scratch. """ model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) @@ -388,8 +409,25 @@ def default(cls) -> "SemanticRegistry": """Return a registry populated with *default_semantic_detectors()*.""" return cls(detectors=default_semantic_detectors()) + @classmethod + def of(cls, *detectors: DQSemanticTypeDetector) -> "SemanticRegistry": + """Return a registry whose chain is exactly *detectors*, in argument order. + + Ergonomic alternative to the ``SemanticRegistry(detectors=(...))`` constructor for + whole-chain construction from scratch. Reads well when composing with the built-in + defaults, e.g. ``SemanticRegistry.of(uuid_detector, *default_semantic_detectors())``. + + Args: + *detectors: Detectors in first-match-wins order. + + Returns: + A new *SemanticRegistry*. Detector-name uniqueness is enforced by the model + validator, so duplicate names raise *pydantic.ValidationError*. + """ + return cls(detectors=tuple(detectors)) + def prepend(self, detector: DQSemanticTypeDetector) -> "SemanticRegistry": - """Return a new registry with *detector* at position 0. + """Return a new registry with *detector* at position 0 (highest priority). The current instance is left unchanged. Re-invokes the constructor so the *_validate_unique_names* model-validator runs on the derived @@ -398,10 +436,78 @@ def prepend(self, detector: DQSemanticTypeDetector) -> "SemanticRegistry": """ return type(self)(detectors=(detector, *self.detectors)) - def replace(self, detectors: Sequence[DQSemanticTypeDetector]) -> "SemanticRegistry": - """Return a new registry whose chain is *detectors*. + def append(self, detector: DQSemanticTypeDetector) -> "SemanticRegistry": + """Return a new registry with *detector* appended at the end (lowest-priority fallback). - The current instance is left unchanged. Re-invokes the constructor - so the uniqueness validator runs on the derived instance. + The current instance is left unchanged. Re-invokes the constructor so the uniqueness + validator runs on the derived instance. + """ + return type(self)(detectors=(*self.detectors, detector)) + + def insert(self, name: str, detector: DQSemanticTypeDetector) -> "SemanticRegistry": + """Return a new registry with *detector* inserted immediately after the entry named *name*. + + Semantic: "insert *detector* after the detector identified by *name*". Position of the + target detector is preserved; every subsequent detector shifts down by one. + + Args: + name: Name of an existing detector; *detector* is placed at the position + immediately after it. + detector: Detector to insert. + + Returns: + A new *SemanticRegistry* with the additional entry. + + Raises: + ValueError: If *name* does not match any detector in the current chain. + pydantic.ValidationError: If *detector*'s name collides with a non-target entry. + """ + index = self._index_of(name) + new_detectors = (*self.detectors[: index + 1], detector, *self.detectors[index + 1 :]) + return type(self)(detectors=new_detectors) + + def replace(self, name: str, detector: DQSemanticTypeDetector) -> "SemanticRegistry": + """Return a new registry with the entry named *name* swapped for *detector*. + + Position is preserved; the rest of the chain is unchanged. To construct a whole + new chain, use ``SemanticRegistry.of(...)`` or the ``SemanticRegistry(detectors=(...))`` + constructor instead. + + Args: + name: Name of the detector to swap out. + detector: Replacement detector. May reuse *name* (position-preserving update) or + take a different name. + + Returns: + A new *SemanticRegistry* with the entry replaced. + + Raises: + ValueError: If *name* does not match any detector in the current chain. + pydantic.ValidationError: If the replacement's name collides with a non-target + entry in the chain. + """ + index = self._index_of(name) + new_detectors = (*self.detectors[:index], detector, *self.detectors[index + 1 :]) + return type(self)(detectors=new_detectors) + + def remove(self, name: str) -> "SemanticRegistry": + """Return a new registry with the entry named *name* filtered out. + + Args: + name: Name of the detector to drop. + + Returns: + A new *SemanticRegistry* without the entry. + + Raises: + ValueError: If *name* does not match any detector in the current chain. """ - return type(self)(detectors=tuple(detectors)) + index = self._index_of(name) + new_detectors = (*self.detectors[:index], *self.detectors[index + 1 :]) + return type(self)(detectors=new_detectors) + + def _index_of(self, name: str) -> int: + for i, detector in enumerate(self.detectors): + if detector.name == name: + return i + raise ValueError(f"Detector {name!r} is not registered") diff --git a/tests/integration/test_profile_semantic.py b/tests/integration/test_profile_semantic.py index bd34b101e..c87e10c24 100644 --- a/tests/integration/test_profile_semantic.py +++ b/tests/integration/test_profile_semantic.py @@ -2,7 +2,10 @@ import pyspark.sql.types as T -from databricks.labs.dqx.config import InputConfig +from databricks.labs.dqx.profiler.profile_builder import ( + PROFILE_BUILDER_REGISTRY, + register_profile_builder, +) from databricks.labs.dqx.profiler.profiler import DQProfiler from databricks.labs.dqx.profiler.semantic import ( DEFAULT_ENUM_DETECTOR, @@ -12,8 +15,6 @@ default_semantic_detectors, ) -from tests.constants import TEST_CATALOG - DEMO_ROW_COUNT = 60 @@ -73,7 +74,11 @@ def test_profile_without_semantic_registry_matches_pre_feature_output(spark, ws) def test_default_semantic_registry_classifies_grounded_columns(spark, ws): df = _make_demo_df(spark) profiler = DQProfiler(ws, semantic_registry=SemanticRegistry.default()) - _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + # vehicle_type has 3 distinct values over 60 rows (ratio 0.05); pass a loose `distinct_ratio` + # so the tightened enum gate still classifies it as enum. + _stats, profiles = profiler.profile( + df, options={"sample_fraction": None, "llm_primary_key_detection": False, "distinct_ratio": 0.1} + ) by_column = _profile_by_column(profiles) @@ -161,46 +166,77 @@ def _always_text(_ctx): ), f"unexpected 'min_max' profile on vehicle_type when forced_text prepended, got: {vehicle_names}" -def test_profile_table_populates_metadata_from_unity_catalog(spark, ws, make_schema, make_random): - """profile_table fetches UC metadata and threads it into ctx.metadata.""" - catalog_name = TEST_CATALOG - schema_name = make_schema(catalog_name=catalog_name).name - table_name = f"{catalog_name}.{schema_name}.t{make_random(10).lower()}" +def test_short_type_low_cardinality_classified_as_enum_and_emits_is_in(spark, ws): + """A low-cardinality ShortType column is classified as `enum` and receives an `is_in` profile. - schema = T.StructType( - [ - T.StructField("vehicle_type", T.StringType(), metadata={"comment": "kind of vehicle"}), - T.StructField("cargo_weight", T.DoubleType()), - ] + Regression guard for the earlier gap where `_detect_enum` accepted ShortType but the `is_in` + builder's `_supports_distinct` rejected it — leaving the column with neither `is_in` nor + `min_max`. + """ + schema = T.StructType([T.StructField("status_code", T.ShortType())]) + rows = [(i % 3 + 1,) for i in range(60)] + df = spark.createDataFrame(rows, schema=schema) + + profiler = DQProfiler(ws, semantic_registry=SemanticRegistry.default()) + _stats, profiles = profiler.profile( + df, options={"sample_fraction": None, "llm_primary_key_detection": False, "distinct_ratio": 0.1} ) - data = [("car", 1.0), ("truck", 2.0), ("van", 3.0)] * 20 - spark.createDataFrame(data, schema).write.format("delta").saveAsTable(table_name) - spark.sql(f"COMMENT ON TABLE {table_name} IS 'demo table with vehicle types'") - spark.sql(f"ALTER TABLE {table_name} ALTER COLUMN vehicle_type COMMENT 'kind of vehicle'") + by_column = _profile_by_column(profiles) - captured: list[dict] = [] + names = {p.name for p in by_column.get("status_code", [])} + assert "is_in" in names, f"expected 'is_in' profile on ShortType status_code column, got: {names}" + is_in_semantic_types = [p.semantic_type for p in by_column.get("status_code", []) if p.name == "is_in"] + assert all( + st == "enum" for st in is_in_semantic_types + ), f"expected ShortType is_in profile tagged semantic_type='enum', got: {is_in_semantic_types}" - def _spy_detect(ctx): - captured.append(dict(ctx.metadata)) - spy = DQSemanticTypeDetector(name="spy", detect=_spy_detect) - registry = SemanticRegistry(detectors=(spy, *default_semantic_detectors())) - profiler = DQProfiler(ws, semantic_registry=registry) - profiler.profile_table( - input_config=InputConfig(location=table_name), - options={"sample_fraction": None, "llm_primary_key_detection": False}, - ) +def test_contextual_builder_after_min_max_observes_resolved_min_max(spark, ws): + """A contextual builder ordered after `min_max` observes resolved `min`/`max` via `ctx.metrics`. - vehicle_metadata = next((m for m in captured if m.get("column_comment") == "kind of vehicle"), None) - assert ( - vehicle_metadata is not None - ), f"expected captured metadata for vehicle_type with column_comment='kind of vehicle', got: {captured}" - assert ( - vehicle_metadata.get("table_name") == table_name - ), f"expected table_name={table_name!r}, got: {vehicle_metadata.get('table_name')!r}" - assert ( - vehicle_metadata.get("table_comment") == "demo table with vehicle types" - ), f"expected table_comment='demo table with vehicle types', got: {vehicle_metadata.get('table_comment')!r}" + Regression guard: `DQProfiler._build_profiles_for_column` calls + `builder_ctx.with_metrics(metrics)` before each contextual builder so the frozen context + reflects write-backs performed by earlier builders. Removing or misplacing that refresh + would silently return stale metrics to downstream builders — a latent contract break for + custom builders that depend on `min_max`'s write-back. + """ + seen_metrics: list[dict] = [] + + @register_profile_builder("_spy_after_min_max", kind="context") + def _spy(ctx): + seen_metrics.append(dict(ctx.metrics)) + return None + + try: + schema = T.StructType([T.StructField("value", T.LongType())]) + rows = [(i,) for i in range(100)] + df = spark.createDataFrame(rows, schema=schema) + + profiler = DQProfiler(ws) + _stats, profiles = profiler.profile( + df, + options={ + "sample_fraction": None, + "llm_primary_key_detection": False, + "remove_outliers": False, + }, + ) + + min_max_profiles = [p for p in profiles if p.column == "value" and p.name == "min_max"] + assert len(min_max_profiles) == 1, f"expected one min_max profile on 'value', got: {min_max_profiles}" + expected_min = min_max_profiles[0].parameters["min"] + expected_max = min_max_profiles[0].parameters["max"] + + assert seen_metrics, "spy contextual builder was not invoked" + spy_seen = seen_metrics[-1] + assert ( + spy_seen.get("min") == expected_min + ), f"expected spy to observe min={expected_min} in ctx.metrics, got: {spy_seen.get('min')}" + assert ( + spy_seen.get("max") == expected_max + ), f"expected spy to observe max={expected_max} in ctx.metrics, got: {spy_seen.get('max')}" + finally: + PROFILE_BUILDER_REGISTRY.pop("_spy_after_min_max", None) def test_enum_detector_values_reused_by_is_in_builder(spark, ws): @@ -216,9 +252,11 @@ def _tracing_detect(ctx): return original(ctx) tracing = DQSemanticTypeDetector(name="enum", detect=_tracing_detect) - registry = SemanticRegistry.default().replace([tracing, *default_semantic_detectors()[1:]]) + registry = SemanticRegistry.of(tracing, *default_semantic_detectors()[1:]) profiler = DQProfiler(ws, semantic_registry=registry) - _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) + _stats, profiles = profiler.profile( + df, options={"sample_fraction": None, "llm_primary_key_detection": False, "distinct_ratio": 0.1} + ) assert invocations == [ "vehicle_type" diff --git a/tests/unit/profiler/test_semantic.py b/tests/unit/profiler/test_semantic.py index 86d498162..a7fa98d96 100644 --- a/tests/unit/profiler/test_semantic.py +++ b/tests/unit/profiler/test_semantic.py @@ -49,6 +49,58 @@ def test_dq_semantic_type_rejects_unsupported_property_value(): EnumProperties(values={"nested": "dict"}) # type: ignore[arg-type] +# --------------------------------------------------------------------------- +# DQProfileContext.with_metrics +# --------------------------------------------------------------------------- + + +def test_dq_profile_context_with_metrics_returns_new_instance_with_snapshot(): + original_metrics = {"count": 10, "count_non_null": 8} + ctx = DQProfileContext( + df=create_autospec(DataFrame), + column_name="c", + column_type=T.IntegerType(), + metrics=original_metrics, + options={"max_in_count": 10}, + semantic_type=DQSemanticType(name="measurement"), + ) + + new_metrics = {"count": 10, "count_non_null": 8, "min": 1, "max": 100} + refreshed = ctx.with_metrics(new_metrics) + + assert refreshed is not ctx + assert dict(refreshed.metrics) == new_metrics + assert dict(ctx.metrics) == original_metrics # original untouched + assert refreshed.column_name == ctx.column_name + assert refreshed.column_type == ctx.column_type + assert refreshed.options == ctx.options + assert refreshed.semantic_type == ctx.semantic_type + assert refreshed.df is ctx.df + + +def test_dq_profile_context_with_metrics_snapshots_supplied_mapping(): + """Later mutations to the source dict must not be reflected in the returned context.""" + source: dict = {"count": 10} + ctx = DQProfileContext( + df=create_autospec(DataFrame), + column_name="c", + column_type=T.IntegerType(), + ) + refreshed = ctx.with_metrics(source) + source["count"] = 999 + assert refreshed.metrics["count"] == 10 + + +def test_dq_profile_context_frozen_blocks_field_reassignment(): + ctx = DQProfileContext( + df=create_autospec(DataFrame), + column_name="c", + column_type=T.IntegerType(), + ) + with pytest.raises(ValidationError): + setattr(ctx, "column_name", "x") + + # --------------------------------------------------------------------------- # default_semantic_detectors chain shape # --------------------------------------------------------------------------- @@ -84,13 +136,19 @@ def test_semantic_registry_prepend_returns_new_instance(): assert original.detectors == snapshot # original unchanged -def test_semantic_registry_replace_returns_new_instance(): - original = SemanticRegistry.default() - replacement = (DQSemanticTypeDetector(name="only", detect=lambda _ctx: None),) - new = original.replace(replacement) - assert new is not original - assert new.detectors == replacement - assert original.detectors == default_semantic_detectors() +def test_semantic_registry_of_returns_new_instance_with_ordered_chain(): + d1 = DQSemanticTypeDetector(name="a", detect=lambda _ctx: None) + d2 = DQSemanticTypeDetector(name="b", detect=lambda _ctx: None) + d3 = DQSemanticTypeDetector(name="c", detect=lambda _ctx: None) + registry = SemanticRegistry.of(d1, d2, d3) + assert [d.name for d in registry.detectors] == ["a", "b", "c"] + + +def test_semantic_registry_of_duplicate_names_raises(): + dup1 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) + dup2 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) + with pytest.raises(ValidationError): + SemanticRegistry.of(dup1, dup2) def test_semantic_registry_assignment_raises(): @@ -105,18 +163,95 @@ def test_semantic_registry_prepend_duplicate_name_raises(): SemanticRegistry.default().prepend(dup) -def test_semantic_registry_replace_duplicate_names_raises(): +def test_semantic_registry_direct_constructor_duplicate_names_raises(): dup1 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) dup2 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) with pytest.raises(ValidationError): - SemanticRegistry().replace([dup1, dup2]) + SemanticRegistry(detectors=(dup1, dup2)) -def test_semantic_registry_direct_constructor_duplicate_names_raises(): - dup1 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) - dup2 = DQSemanticTypeDetector(name="foo", detect=lambda _ctx: None) +# --------------------------------------------------------------------------- +# SemanticRegistry: append / insert / replace / remove +# --------------------------------------------------------------------------- + + +def test_semantic_registry_append_lands_last_and_leaves_original_unchanged(): + original = SemanticRegistry.default() + snapshot = original.detectors + extra = DQSemanticTypeDetector(name="custom", detect=lambda _ctx: None) + new = original.append(extra) + assert new is not original + assert new.detectors[-1] is extra + assert new.detectors[:-1] == snapshot + assert original.detectors == snapshot + + +def test_semantic_registry_append_duplicate_name_raises(): + dup = DQSemanticTypeDetector(name="enum", detect=lambda _ctx: None) with pytest.raises(ValidationError): - SemanticRegistry(detectors=(dup1, dup2)) + SemanticRegistry.default().append(dup) + + +def test_semantic_registry_insert_places_detector_after_named_entry(): + original = SemanticRegistry.default() + extra = DQSemanticTypeDetector(name="custom", detect=lambda _ctx: None) + new = original.insert("enum", extra) + names = [d.name for d in new.detectors] + assert names == ["enum", "custom", "key", "measurement", "text"] + assert original.detectors == default_semantic_detectors() + + +def test_semantic_registry_insert_unknown_name_raises_value_error(): + extra = DQSemanticTypeDetector(name="custom", detect=lambda _ctx: None) + with pytest.raises(ValueError, match="'missing'"): + SemanticRegistry.default().insert("missing", extra) + + +def test_semantic_registry_insert_duplicate_name_raises(): + dup = DQSemanticTypeDetector(name="key", detect=lambda _ctx: None) + with pytest.raises(ValidationError): + SemanticRegistry.default().insert("enum", dup) + + +def test_semantic_registry_replace_swaps_named_entry_preserving_position(): + original = SemanticRegistry.default() + replacement = DQSemanticTypeDetector(name="enum", detect=lambda _ctx: None) + new = original.replace("enum", replacement) + assert new.detectors[0] is replacement + assert [d.name for d in new.detectors] == ["enum", "key", "measurement", "text"] + assert original.detectors == default_semantic_detectors() + + +def test_semantic_registry_replace_allows_renaming_the_target_entry(): + original = SemanticRegistry.default() + replacement = DQSemanticTypeDetector(name="renamed_enum", detect=lambda _ctx: None) + new = original.replace("enum", replacement) + assert [d.name for d in new.detectors] == ["renamed_enum", "key", "measurement", "text"] + + +def test_semantic_registry_replace_unknown_name_raises_value_error(): + replacement = DQSemanticTypeDetector(name="x", detect=lambda _ctx: None) + with pytest.raises(ValueError, match="'missing'"): + SemanticRegistry.default().replace("missing", replacement) + + +def test_semantic_registry_replace_name_clash_with_other_entry_raises(): + # Renaming `enum` → `key` would collide with the existing `key` entry. + replacement = DQSemanticTypeDetector(name="key", detect=lambda _ctx: None) + with pytest.raises(ValidationError): + SemanticRegistry.default().replace("enum", replacement) + + +def test_semantic_registry_remove_drops_named_entry(): + original = SemanticRegistry.default() + new = original.remove("text") + assert [d.name for d in new.detectors] == ["enum", "key", "measurement"] + assert original.detectors == default_semantic_detectors() + + +def test_semantic_registry_remove_unknown_name_raises_value_error(): + with pytest.raises(ValueError, match="'missing'"): + SemanticRegistry.default().remove("missing") # --------------------------------------------------------------------------- @@ -283,6 +418,22 @@ def test_enum_detector_positive_low_cardinality(): assert typed.values == {"car", "truck", "van"} +def test_enum_detector_respects_distinct_ratio_option(): + """When `distinct_ratio` option is set below the observed ratio, semantic-enum is suppressed. + + Mirrors reviewer's 12-rows / 9-distinct example: legacy `is_in` deliberately suppresses + at low repetition, so semantic profiling must too. + """ + ctx = DQProfileContext( + df=_mock_df_with_distinct("c", list("abcdefghi")), + column_name="c", + column_type=T.StringType(), + metrics={"count_non_null": 12, "count_distinct": 9}, + options={"max_in_count": 10, "distinct_ratio": 0.05}, + ) + assert DEFAULT_ENUM_DETECTOR.detect(ctx) is None + + # --------------------------------------------------------------------------- # String-key length stability guard (uses Spark aggregate via autospec) # --------------------------------------------------------------------------- diff --git a/tests/unit/test_profile_builder.py b/tests/unit/test_profile_builder.py index 15478bf45..6f0fd09a4 100644 --- a/tests/unit/test_profile_builder.py +++ b/tests/unit/test_profile_builder.py @@ -312,28 +312,37 @@ def _make_mock_df(columns: list, distinct_values: list) -> DataFrame: @pytest.mark.parametrize("column_type", [T.DoubleType(), T.FloatType(), T.BooleanType(), T.DateType()]) def test_is_in_unsupported_type_returns_none(mock_df, column_type): - assert make_is_in_profile(_ctx(mock_df, "col", column_type, {"count": 10}, {})) is None + assert make_is_in_profile(_ctx(mock_df, "col", column_type, {"count": 10, "count_non_null": 10}, {})) is None -@pytest.mark.parametrize("column_type", [T.CharType(10), T.VarcharType(50)]) -def test_is_in_char_varchar_type_returns_profile(column_type): - df = _make_mock_df(["col"], ["a", "b", "c"]) +@pytest.mark.parametrize("column_type", [T.CharType(10), T.VarcharType(50), T.ShortType()]) +def test_is_in_supported_types_returns_profile(column_type): + df = _make_mock_df(["col"], [1, 2, 3] if isinstance(column_type, T.ShortType) else ["a", "b", "c"]) profile = make_is_in_profile( - _ctx(df, "col", column_type, {"count": 10}, {"max_in_count": 10, "distinct_ratio": 1.0}) + _ctx(df, "col", column_type, {"count": 10, "count_non_null": 10}, {"max_in_count": 10, "distinct_ratio": 1.0}) ) assert profile is not None assert profile.name == "is_in" - assert set(profile.parameters["in"]) == {"a", "b", "c"} + expected = {1, 2, 3} if isinstance(column_type, T.ShortType) else {"a", "b", "c"} + assert set(profile.parameters["in"]) == expected -def test_is_in_total_count_zero_returns_none(mock_df): - assert make_is_in_profile(_ctx(mock_df, "col", T.IntegerType(), {"count": 0}, {})) is None +def test_is_in_count_non_null_zero_returns_none(mock_df): + assert make_is_in_profile(_ctx(mock_df, "col", T.IntegerType(), {"count": 0, "count_non_null": 0}, {})) is None def test_is_in_no_distinct_values_returns_none(): df = _make_mock_df(["col"], []) assert ( - make_is_in_profile(_ctx(df, "col", T.StringType(), {"count": 3}, {"max_in_count": 10, "distinct_ratio": 1.0})) + make_is_in_profile( + _ctx( + df, + "col", + T.StringType(), + {"count": 3, "count_non_null": 3}, + {"max_in_count": 10, "distinct_ratio": 1.0}, + ) + ) is None ) @@ -341,7 +350,13 @@ def test_is_in_no_distinct_values_returns_none(): def test_is_in_conditions_met_returns_profile(): df = _make_mock_df(["col"], [1, 2, 3]) profile = make_is_in_profile( - _ctx(df, "status", T.IntegerType(), {"count": 5}, {"max_in_count": 10, "distinct_ratio": 1.0}) + _ctx( + df, + "status", + T.IntegerType(), + {"count": 5, "count_non_null": 5}, + {"max_in_count": 10, "distinct_ratio": 1.0}, + ) ) assert profile is not None assert profile.name == "is_in" @@ -353,16 +368,28 @@ def test_is_in_distinct_count_exceeds_max_in_count_returns_none(): # 11 distinct values, max_in_count=10 → distinct_count > max_in_count → None df = _make_mock_df(["col"], list(range(11))) profile = make_is_in_profile( - _ctx(df, "col", T.IntegerType(), {"count": 100}, {"max_in_count": 10, "distinct_ratio": 1.0}) + _ctx( + df, + "col", + T.IntegerType(), + {"count": 100, "count_non_null": 100}, + {"max_in_count": 10, "distinct_ratio": 1.0}, + ) ) assert profile is None def test_is_in_distinct_ratio_exceeds_threshold_returns_none(): - # 10 distinct values in 10 total → ratio=1.0, threshold=0.5 + # 10 distinct values over 10 non-null → ratio=1.0, threshold=0.5 df = _make_mock_df(["col"], list(range(10))) profile = make_is_in_profile( - _ctx(df, "col", T.StringType(), {"count": 10}, {"max_in_count": 20, "distinct_ratio": 0.5}) + _ctx( + df, + "col", + T.StringType(), + {"count": 10, "count_non_null": 10}, + {"max_in_count": 20, "distinct_ratio": 0.5}, + ) ) assert profile is None @@ -370,7 +397,13 @@ def test_is_in_distinct_ratio_exceeds_threshold_returns_none(): def test_is_in_filter_propagated(): df = _make_mock_df(["col"], ["a", "b"]) profile = make_is_in_profile( - _ctx(df, "col", T.StringType(), {"count": 5}, {"max_in_count": 10, "distinct_ratio": 1.0, "filter": "x > 0"}) + _ctx( + df, + "col", + T.StringType(), + {"count": 5, "count_non_null": 5}, + {"max_in_count": 10, "distinct_ratio": 1.0, "filter": "x > 0"}, + ) ) assert profile is not None assert profile.filter == "x > 0" @@ -403,7 +436,7 @@ def test_is_in_skipped_for_non_enum_semantic_type(mock_df, other_type): mock_df, "col", T.StringType(), - {"count": 100}, + {"count": 100, "count_non_null": 100}, {"max_in_count": 10, "distinct_ratio": 0.1}, semantic_type=DQSemanticType(name=other_type), ) From 9fd13969c87e9daf867c8ed152816fce369a9561 Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 5 Sep 2026 14:35:05 +0200 Subject: [PATCH 7/9] linter fixes --- tests/integration/test_profile_semantic.py | 9 ++++----- tests/unit/profiler/test_semantic.py | 19 +++++++++++-------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/tests/integration/test_profile_semantic.py b/tests/integration/test_profile_semantic.py index c87e10c24..c6f56c2a3 100644 --- a/tests/integration/test_profile_semantic.py +++ b/tests/integration/test_profile_semantic.py @@ -205,13 +205,11 @@ def test_contextual_builder_after_min_max_observes_resolved_min_max(spark, ws): @register_profile_builder("_spy_after_min_max", kind="context") def _spy(ctx): seen_metrics.append(dict(ctx.metrics)) - return None - try: + def _profile_and_assert(): schema = T.StructType([T.StructField("value", T.LongType())]) rows = [(i,) for i in range(100)] df = spark.createDataFrame(rows, schema=schema) - profiler = DQProfiler(ws) _stats, profiles = profiler.profile( df, @@ -221,12 +219,10 @@ def _spy(ctx): "remove_outliers": False, }, ) - min_max_profiles = [p for p in profiles if p.column == "value" and p.name == "min_max"] assert len(min_max_profiles) == 1, f"expected one min_max profile on 'value', got: {min_max_profiles}" expected_min = min_max_profiles[0].parameters["min"] expected_max = min_max_profiles[0].parameters["max"] - assert seen_metrics, "spy contextual builder was not invoked" spy_seen = seen_metrics[-1] assert ( @@ -235,6 +231,9 @@ def _spy(ctx): assert ( spy_seen.get("max") == expected_max ), f"expected spy to observe max={expected_max} in ctx.metrics, got: {spy_seen.get('max')}" + + try: + _profile_and_assert() finally: PROFILE_BUILDER_REGISTRY.pop("_spy_after_min_max", None) diff --git a/tests/unit/profiler/test_semantic.py b/tests/unit/profiler/test_semantic.py index a7fa98d96..110ffa52b 100644 --- a/tests/unit/profiler/test_semantic.py +++ b/tests/unit/profiler/test_semantic.py @@ -56,8 +56,9 @@ def test_dq_semantic_type_rejects_unsupported_property_value(): def test_dq_profile_context_with_metrics_returns_new_instance_with_snapshot(): original_metrics = {"count": 10, "count_non_null": 8} + df_mock = create_autospec(DataFrame) ctx = DQProfileContext( - df=create_autospec(DataFrame), + df=df_mock, column_name="c", column_type=T.IntegerType(), metrics=original_metrics, @@ -81,8 +82,9 @@ def test_dq_profile_context_with_metrics_returns_new_instance_with_snapshot(): def test_dq_profile_context_with_metrics_snapshots_supplied_mapping(): """Later mutations to the source dict must not be reflected in the returned context.""" source: dict = {"count": 10} + df_mock = create_autospec(DataFrame) ctx = DQProfileContext( - df=create_autospec(DataFrame), + df=df_mock, column_name="c", column_type=T.IntegerType(), ) @@ -92,8 +94,9 @@ def test_dq_profile_context_with_metrics_snapshots_supplied_mapping(): def test_dq_profile_context_frozen_blocks_field_reassignment(): + df_mock = create_autospec(DataFrame) ctx = DQProfileContext( - df=create_autospec(DataFrame), + df=df_mock, column_name="c", column_type=T.IntegerType(), ) @@ -137,10 +140,10 @@ def test_semantic_registry_prepend_returns_new_instance(): def test_semantic_registry_of_returns_new_instance_with_ordered_chain(): - d1 = DQSemanticTypeDetector(name="a", detect=lambda _ctx: None) - d2 = DQSemanticTypeDetector(name="b", detect=lambda _ctx: None) - d3 = DQSemanticTypeDetector(name="c", detect=lambda _ctx: None) - registry = SemanticRegistry.of(d1, d2, d3) + detector_a = DQSemanticTypeDetector(name="a", detect=lambda _ctx: None) + detector_b = DQSemanticTypeDetector(name="b", detect=lambda _ctx: None) + detector_c = DQSemanticTypeDetector(name="c", detect=lambda _ctx: None) + registry = SemanticRegistry.of(detector_a, detector_b, detector_c) assert [d.name for d in registry.detectors] == ["a", "b", "c"] @@ -171,7 +174,7 @@ def test_semantic_registry_direct_constructor_duplicate_names_raises(): # --------------------------------------------------------------------------- -# SemanticRegistry: append / insert / replace / remove +# SemanticRegistry mutation helpers (append, insert, replace, remove) # --------------------------------------------------------------------------- From d45bef38cf06d15c1e3aa92c2dac389dab9f2543 Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Thu, 10 Sep 2026 05:42:31 +0200 Subject: [PATCH 8/9] removed local plans --- .agents/2026-09-05-11-42-05-pr-1491-review.md | 164 -------- .agents/2026-09-05-12-48-53-plan.md | 386 ------------------ .agents/2026-09-05-12-48-53-review-round-0.md | 50 --- 3 files changed, 600 deletions(-) delete mode 100644 .agents/2026-09-05-11-42-05-pr-1491-review.md delete mode 100644 .agents/2026-09-05-12-48-53-plan.md delete mode 100644 .agents/2026-09-05-12-48-53-review-round-0.md diff --git a/.agents/2026-09-05-11-42-05-pr-1491-review.md b/.agents/2026-09-05-11-42-05-pr-1491-review.md deleted file mode 100644 index f0d64f6ad..000000000 --- a/.agents/2026-09-05-11-42-05-pr-1491-review.md +++ /dev/null @@ -1,164 +0,0 @@ ---- -title: Manual PR-state review of feature/semantic_type_classification (PR #1491) -scope: uncommitted changes + branch commits + previously-addressed PR review comments -created_at: 2026-09-05T11:42:05Z -pr: https://github.com/databrickslabs/dqx/pull/1491 -verdict: ship-ready with minor polish ---- - -# Top Critical Issues - -## 1. Enum-gate denominator still differs from legacy `is_in` — residual divergence -`_detect_enum` computes `cardinality / count_non_null` (nulls excluded), but the legacy -`make_is_in_profile` computes `distinct_count / total_count` (nulls included). Even with the -`distinct_ratio` threshold now aligned, a null-heavy column can be classified differently by -the two paths. - -**Example** — 100 rows, 40 non-null, 3 distinct: -- Legacy is_in: 3/100 = 0.03 < 0.05 → emits `is_in`. -- Semantic `_detect_enum`: 3/40 = 0.075 ≥ 0.05 → rejects. - -Reviewer comment #3 asked for alignment; the threshold was aligned but the denominator was not. -Options: (a) switch semantic to `count`, or (b) document the deliberate difference in the -`ENUM_MAX_CARDINALITY_RATIO` comment. - -**Severity**: modest — semantic profiling is opt-in and the divergence is bounded. -Feedback: adjust `is_in` profile to take into count non-null values. -**Resolution (2026-09-05)**: `make_is_in_profile` in `src/databricks/labs/dqx/profiler/profile_builder.py` now divides by `count_non_null` (matching `_detect_enum`) and short-circuits when `count_non_null == 0`. Unit tests in `tests/unit/test_profile_builder.py` updated to supply `count_non_null` alongside `count`; the previously named `test_is_in_total_count_zero_returns_none` is now `test_is_in_count_non_null_zero_returns_none`. `distinct_ratio` docs row updated from "5% of total" → "5% of non-null values". - -## 2. `with_metrics` fires on every builder iteration — wasteful cache-invalidation -`profiler.py:_build_profiles_for_column` calls `builder_ctx = builder_ctx.with_metrics(metrics)` -inside the loop, allocating a new frozen model + copy of `metrics` for every contextual builder -— even the first, and even when no prior builder mutated `metrics`. Only `min_max` writes back -today, so a cheaper pattern is: refresh only *after* a builder returns a `min_max` profile. - -Correctness is fine; the current implementation pays a small allocation tax on every -column × builder. - -**Severity**: low. -Feedback: skip - -## 3. `_fetch_table_metadata` bare `except Exception` — safe but broad -`profiler.py:461` catches everything, per reviewer comment #4. The comment explains the reasoning, -and it does honor the docstring guarantee. But this now swallows bugs in metadata assembly (e.g. -`AttributeError` from an SDK schema change). Consider tightening to `(DatabricksError, ValueError, -TimeoutError)` — the concrete client-side failures actually observed. - -**NOTE (updated 2026-09-05T11:42:05Z)**: superseded — the entire `_fetch_table_metadata` method, -its call chain, and the `DQProfileContext.metadata` field were removed in a subsequent commit. -Finding no longer applies. - -## 4. Naming: `SemanticRegistry.insert(name, detector)` reads as *insert at position `name`* -The method inserts *after* the entry named `name`. Docstring says so, but the naming still -surprises. `insert_after` was rejected as verbose. An alternative would be `after(name, detector)` -— reads naturally and encodes the semantic. - -**Severity**: cosmetic; safe to leave. -Feedback: this method must be already removed, ignore - -## 5. Docs code block for the composition summary lacks a language hint -`profiler.mdx:156` opens a fence with no language marker: -```` -``` -# prepend(detector) → puts one detector at the front (highest priority) -... -``` -```` -Not `python` (uses `→`), not `text`. Renders monospace without syntax highlighting. - -**Severity**: minor — mark as `text` or wrap in prose. -Feedback: fix -**Resolution (2026-09-05)**: `docs/dqx/docs/reference/profiler.mdx` fence now opens `` ```text `` and the `#` comment markers are dropped in favour of plain lines. - -## 6. PR title is a placeholder ("Feature/semantic type classification") -Follows branch-name convention. Before merge, replace with a real title (e.g. "Add opt-in -semantic-aware profiling to DQProfiler"). Also refresh the PR body — some sections still -describe `SemanticRegistry` as offering only `default()`, `prepend(...)`, and `replace([...])`, -which is the old API. - -**Severity**: blocker for merge (metadata only, no code change). -Feedback: fix -**Resolution (2026-09-05)**: PR #1491 title updated on GitHub to *"Add opt-in semantic-aware profiling to DQProfiler"* via `gh pr edit`. - -## 7. Deprecated / broken PR-body claims -The PR body describes `SemanticRegistry.replace([...])` as accepting a list. That's the old -API. Refresh the description to list `of`, `append`, `insert(name, detector)`, -`replace(name, detector)`, `remove(name)` before merge. - -**Severity**: blocker for merge (metadata only). -Feedback: fix -**Resolution (2026-09-05)**: PR #1491 body refreshed on GitHub via `gh pr edit` — now lists the correct composition surface (`prepend`, `append`, `insert`, `replace(name, detector)`, `remove`, `of(...)`), drops the removed `metadata` field from the `DQProfileContext` description, and updates the enum-gate/denominator prose. UC metadata plumbing bullet removed to reflect its removal from the branch. - -## 8. No unit test asserts `with_metrics` is actually observed by a downstream contextual builder -`test_semantic.py` covers `with_metrics` in isolation (snapshot semantics), but there is no test -proving that a contextual builder registered *after* `min_max` actually sees the updated -`min`/`max` in `ctx.metrics`. The bug this fix was meant to prevent is a silent latent contract -break (reviewer comment #2). Add a builder-loop integration test to lock the behavior down — -otherwise a future refactor could reintroduce the bug without failing tests. - -**Severity**: modest — reviewer round-0 finding #2 also flagged this. -Feedback: fix -**Resolution (2026-09-05)**: `test_contextual_builder_after_min_max_observes_resolved_min_max` added in `tests/integration/test_profile_semantic.py`. Registers a spy contextual builder via `@register_profile_builder(kind="context")` (lands after `min_max` in `PROFILE_BUILDER_REGISTRY` insertion order), asserts spy's `ctx.metrics` carries the resolved `min`/`max`, cleans up in `finally`. Uses `remove_outliers=False` for exact expected values. - -## 9. Integration-test fixture inconsistency -`test_default_semantic_registry_classifies_grounded_columns` and -`test_enum_detector_values_reused_by_is_in_builder` pass `"distinct_ratio": 0.1` explicitly; -other tests using `_make_demo_df` do not. This is fine for the specific fixtures used, but a -comment on `_make_demo_df` noting *"3/60 = 0.05 exactly hits the tightened default; callers -who want the enum to fire must widen distinct_ratio"* would prevent future surprise. - -**Severity**: minor doc hygiene. -Feedback: skip - -## 10. Reviewer comment #9 / #10 point 3 — decorator vs value-object asymmetry — remains unaddressed -Two distinct mental models for "add my thing to the chain" (`@register_profile_builder` for -builders vs constructor threading for detectors) still coexist. Explicitly out of scope per -the original prompt. Worth capturing as a follow-up issue so it does not get lost. - -**Severity**: architectural, not a blocker. -Feedback: skip - ---- - -# Final Verdict - -**Ship-ready with minor polish.** - -The five inline reviewer findings (#1–#5) and reviewer #6/#7/#8 doc nits are all addressed. -The semantic-registry ergonomics gaps in #10 (whole-chain naming and missing -`append/insert/remove/replace-one-by-name` operations) are all fixed. Remaining items are -follow-up polish, not blockers. - -## Blockers before merge -- #6 PR title -- #7 PR body accuracy - -Both are metadata-only, no code changes required. - -## Nice-to-have before merge -- #5 doc code-fence language -- #9 fixture comment - -## Recommended follow-up issues -- #1 denominator alignment (or documentation) -- #2 per-iteration allocation -- ~~#3 exception scope~~ (resolved by full removal of `_fetch_table_metadata` plumbing) -- #8 loop-integration test -- #10 API-model asymmetry - -## Assessment of core work -The core architecture (immutable registry, first-match-wins chain, contextual builders reusing -enum values, UC metadata plumbing) is sound. The correctness bugs from the review (silent check -drops on `ShortType`, min/max write-back, permissive enum gate, `DatabricksError`-only catch, -docstring typos) are all fixed with regression coverage. Code quality is production-grade: fully -typed, immutable value objects, Pydantic-validated invariants, and no linting suppressions. - ---- - -# Follow-up (post-review) - -After this review, the user directed removal of the entire UC metadata plumbing -(`_fetch_table_metadata`, `_build_column_metadata`, the `table_metadata`/`column_metadata` -kwargs, the `metadata_map` propagation, and the vestigial `DQProfileContext.metadata` field). -This resolves finding #3 outright and removes the "Table metadata" documentation section along -with its integration test. diff --git a/.agents/2026-09-05-12-48-53-plan.md b/.agents/2026-09-05-12-48-53-plan.md deleted file mode 100644 index ca6a949a4..000000000 --- a/.agents/2026-09-05-12-48-53-plan.md +++ /dev/null @@ -1,386 +0,0 @@ ---- -title: Address review comments on PR #1491 (semantic-aware profiling) except the decorator-vs-immutable-registry mental-model feedback. -status: implemented -created_at: 2026-09-05T10:48:53Z -updated_at: 2026-09-05T11:12:57Z -implemented_at: 2026-09-05T11:28:10Z - -input: > - https://github.com/databrickslabs/dqx/pull/1491 - this is PR review for the current branch. - Readout carefully and address comments, except: "They way this is constructed is different than - registry approach like @register_profile_builder(...) . Profile builders are customized with a - decorator against a global registry; semantic detectors are customized by constructing an - immutable value object and threading it through a constructor. Two different mental models for - 'add my thing to the chain.'" and similar feedback on using decorator. - -context: > - Feature branch feature/semantic_type_classification adds opt-in semantic-aware profiling to - `DQProfiler`. Reviewer mwojtyczka posted 10 inline comments on 2026-09-01 with a - CHANGES_REQUESTED verdict. Explicitly out of scope per user: comment #9 (decorator vs immutable - mental-model criticism at profiler.mdx:183) and comment #10 point 3 (same criticism restated at - semantic.py:391). All other structural/correctness/doc feedback is in scope. - - Relevant files: - - src/databricks/labs/dqx/profiler/semantic.py - - src/databricks/labs/dqx/profiler/profile_builder.py - - src/databricks/labs/dqx/profiler/profiler.py - - docs/dqx/docs/reference/profiler.mdx - - tests/unit/profiler/test_semantic.py - - tests/unit/test_profile_builder.py - - tests/integration/test_profile_semantic.py - - tests/perf/test_profile_semantic.py - - Comment → step mapping: - #1 → Step 1 (ShortType enum gap) - #2 → Step 2 (metrics write-back through ctx) - #3 → Step 3 (enum cardinality gate alignment with distinct_ratio) - #4 → Step 4 (broaden UC metadata exception handler) - #5 → Step 5 (docs `type=` → `kind=`) - #6 → Step 5 (docs "iff" wording) - #7 → Step 6 (docs comments on prepend/replace and #8 rename) - #8 → Step 6/7 (rename replace, add append/remove/insert; docs update) - #10 → Step 6/7 (points 1 & 2 only — rename + missing operations) - #9 and #10 point 3 → INTENTIONALLY SKIPPED per user ---- - -# Plan - -### Step - 1 - -**Description** - -Close the `ShortType` gap flagged in comment #1 by widening the `is_in` -builder's supported types (per user feedback: short integers are a valid -enum value range, so keep them in `_detect_enum` and make `is_in` accept -them too). - -- In `src/databricks/labs/dqx/profiler/profile_builder.py`, add - `T.ShortType` to `_supports_distinct` so the tuple becomes - `(T.IntegerType, T.LongType, T.ShortType) + TEXT_TYPES`. Update the - helper's docstring to reflect the extended set. -- Leave `_detect_enum` in `src/databricks/labs/dqx/profiler/semantic.py` - unchanged (`ShortType` stays in its accepted-types tuple). Add a one-line - comment on the `_supports_distinct` tuple noting that its accepted set - must stay in sync with `_detect_enum` so semantic-enum classification - never suppresses `min_max` without producing an `is_in` in return. -- Extend the parametrised type coverage in - `tests/unit/test_profile_builder.py`: - * Add `T.ShortType()` to the parametrise list on - `test_is_in_char_varchar_type_returns_profile` (rename generically to - `test_is_in_supported_types_returns_profile` if it makes the parametrise - read cleaner) so it exercises Short alongside `CharType`/`VarcharType`. - * Remove `T.ShortType()` from any negative parametrise if present (grep - confirms it is not currently in `test_is_in_unsupported_type_returns_none`, - so nothing to remove there — verify at implementation time). -- Add a positive integration case in - `tests/integration/test_profile_semantic.py` (or, if simpler, - `tests/integration/test_profile_builder.py`) covering a low-cardinality - `ShortType` column: assert it is classified as `enum` **and** receives - an `is_in` profile, so the previously-silent gap is regression-guarded - end-to-end. Use the existing `_make_demo_df`-style fixture pattern. -- Confirm `is_in_list` in `src/databricks/labs/dqx/check_funcs.py:524` - accepts a `ShortType` Column at apply time (it uses `.isin(...)` which - is type-agnostic in Spark). No change needed there; call out in the - step feedback if apply-time behaviour differs from expectations. - -**Feedback** - - -### Step - 2 - -**Description** - -Fix the min/max write-back gap flagged in comment #2 -(`src/databricks/labs/dqx/profiler/profiler.py:606`), and encapsulate the -metrics refresh on `DQProfileContext` per user feedback (do not sprinkle -`model_copy(update={...})` at the call site). - -- In `src/databricks/labs/dqx/profiler/semantic.py`, add a new method - on `DQProfileContext`: - ```python - def with_metrics(self, metrics: Mapping[str, Any]) -> "DQProfileContext": - """Return a new context whose `metrics` is a fresh snapshot of *metrics*. - ... - """ - return self.model_copy(update={"metrics": dict(metrics)}) - ``` - Docstring should explain the intended use: because the model is - frozen and pydantic materialises `metrics` as a copy on construction, - contextual builders further down the chain would otherwise not see - write-backs made to the outer mutable `metrics` dict between builder - invocations. `with_metrics` produces a new frozen context wrapping - the current dict snapshot. -- In `src/databricks/labs/dqx/profiler/profiler.py` - (`DQProfiler._build_profiles_for_column`), keep the initial - `builder_ctx` construction, but inside the - `for profile_type in PROFILE_BUILDER_REGISTRY.values():` loop refresh - it before each contextual invocation: - ```python - if profile_type.contextual_builder is not None: - builder_ctx = builder_ctx.with_metrics(metrics) - profile = profile_type.contextual_builder(builder_ctx) - ``` - Leave the legacy-callback branch unchanged (it already receives - `dict(metrics)` per call). -- Add a short reference comment above the refresh line pointing to the - min/max write-back block below (~lines 631–635) so future readers - understand why the refresh is required. -- Do not change per-column initial detector `detector_ctx` construction - in `_detect_semantic_type`. -- Add unit coverage in `tests/unit/profiler/test_semantic.py` for the new - method: assert `with_metrics(new_dict)` returns a new instance with the - snapshot semantics, leaves the original untouched, preserves all other - fields (`df`, `column_name`, `column_type`, `options`, `metadata`, - `semantic_type`), and rejects post-construction mutation of the - original's fields (Pydantic frozen behaviour). - -**Feedback** - - -### Step - 3 - -**Description** - -Align the enum cardinality gate with the legacy `is_in` distinct_ratio flagged -in comment #3 (`src/databricks/labs/dqx/profiler/semantic.py:213`). - -- In `_detect_enum`, replace the standalone `ENUM_MAX_CARDINALITY_RATIO` - gate with a lookup against - `PROFILE_OPTION_DISTINCT_RATIO` from `ctx.options`, falling back to - `ENUM_MAX_CARDINALITY_RATIO` when the option is absent. Use strict `>=` - (mirroring the legacy `is_in` builder's `distinct_ratio < max_distinct_ratio` - emission gate) so semantic profiling and legacy profiling agree on which - columns qualify. -- Lower the default `ENUM_MAX_CARDINALITY_RATIO` in `semantic.py` to `0.05` - so the constant reflects the effective default when options are not merged - (e.g. when semantic detectors are used outside `DQProfiler`, or when a caller - passes `options={}` explicitly). Update the module-level comment that - currently justifies 0.95 as "more permissive" so it explains the alignment - with `PROFILE_OPTION_DISTINCT_RATIO`. -- Import `PROFILE_OPTION_DISTINCT_RATIO` from `profile_options.py` in - `semantic.py` (already imports `PROFILE_OPTION_MAX_IN_COUNT`, so add - alongside). -- No API breakage: `ENUM_MAX_CARDINALITY_RATIO` stays exported but its value - and semantics change — call out in the plan feedback if downstream users - depend on the old 0.95 value. - -**Feedback** - - -### Step - 4 - -**Description** - -Broaden the UC metadata-fetch exception handler flagged in comment #4 -(`src/databricks/labs/dqx/profiler/profiler.py:462`). - -- In `DQProfiler._fetch_table_metadata`, widen the `except DatabricksError` - clause to `except Exception` (bare, not `BaseException`, so `KeyboardInterrupt` - and `SystemExit` still propagate). This honours the docstring's "metadata - errors never block profiling" guarantee for client-side failures like a - `ValueError` raised by SDK name validation on non-UC locations - (e.g. storage paths). -- Keep the log message shape unchanged (`safe_location`, newline scrubbing), - and keep the same "return `{}`" behaviour. -- Remove the now-unused `DatabricksError` import at the top of `profiler.py` - only if no other code path still uses it (verify with grep before deleting). - -**Feedback** - - -### Step - 5 - -**Description** - -Fix the two documentation nits flagged in comments #5 and #6 -(`docs/dqx/docs/reference/profiler.mdx`). - -- Line 271 (prose): replace `an optional \`type="context"\` argument` with - `an optional \`kind="context"\` argument`. -- Line 280 (code block): replace - `@register_profile_builder("my_custom", type="context")` with - `@register_profile_builder("my_custom", kind="context")` so the example - matches the actual decorator signature in - `src/databricks/labs/dqx/profiler/profile_builder.py:46`. -- Line 147: replace `The feature is fully opt-in: it runs iff a`… - with the suggested wording - `The feature is fully opt-in: it runs if a \`semantic_registry\` is - supplied to the`… (accept the reviewer's `suggestion` block verbatim). - -**Feedback** - - -### Step - 6 - -**Description** - -Redesign the `SemanticRegistry` customization surface per comments #7, #8, -and #10 points 1–2 (`src/databricks/labs/dqx/profiler/semantic.py:391`). -Comment #9 and comment #10 point 3 are explicitly out of scope — do not -convert semantic detectors to a decorator/global-registry pattern. - -**Do NOT add a `with_detectors(...)` method** (per user feedback: chaining -`SemanticRegistry.default().with_detectors([...])` reads nonsensically because -the default chain is constructed and then discarded). Whole-chain construction -is served by the existing constructor and a new classmethod for ergonomics -(see below). - -- **Repurpose `replace`.** Change the existing `replace(detectors: Sequence[...])` - method (whole-chain) into a **single-detector swap** with signature - `replace(name: str, detector: DQSemanticTypeDetector) -> "SemanticRegistry"`. - It swaps the entry named `name` with `detector`, preserving position, and - raises `ValueError` when `name` is not present (mirrors the reviewer's - expected `replace(name, detector)` semantics). All call sites that used - the old whole-chain `replace([...])` must migrate to the constructor or the - new `of(...)` classmethod added below — this is covered in Step 7. - -- **Add the following new methods** to `SemanticRegistry`, each returning a new - `SemanticRegistry` instance and routing through the constructor so the - `_validate_unique_names` model-validator runs on the derived instance - (matching the existing `prepend` pattern): - * `append(detector: DQSemanticTypeDetector) -> "SemanticRegistry"` — - lowest-priority fallback: places `detector` at the end of the chain. - * `insert(name: str, detector: DQSemanticTypeDetector) -> "SemanticRegistry"` — - (renamed from `insert_after` per user feedback) inserts `detector` - immediately after the entry whose name equals `name`. Raise `ValueError` - when `name` is not present. The docstring must state the "after" semantics - explicitly so the terser name is unambiguous ("insert this detector after - the one identified by `name`"). - * `remove(name: str) -> "SemanticRegistry"` — returns a new registry with - the detector matching `name` filtered out. Raise `ValueError` when `name` - is not present, to keep the API strict (consistent with the uniqueness - invariant). - -- **Add a classmethod for whole-chain construction from scratch** (per user - feedback: prefer a static/class method over the previously-planned - `with_detectors` instance method). Signature: - ```python - @classmethod - def of(cls, *detectors: DQSemanticTypeDetector) -> "SemanticRegistry": - """Return a registry whose chain is exactly *detectors* (in argument order).""" - return cls(detectors=tuple(detectors)) - ``` - Rationale: varargs reads well when combining with `default_semantic_detectors()` - (e.g. `SemanticRegistry.of(uuid_detector, *default_semantic_detectors())`) and - it complements the existing `.default()` classmethod symmetrically. The - existing constructor `SemanticRegistry(detectors=(...))` continues to work - for callers that want the kwarg form. - -- Add Google-style docstrings on every new/changed method (one paragraph each, - matching the existing tone). Keep `prepend(detector)` unchanged — it is - still the correct primitive for "make mine win first." - -- Update the module docstring header at the top of `semantic.py` so the - registry summary lists the new/renamed methods: `of` (classmethod), - `default` (classmethod, unchanged), `prepend`, `append`, `insert`, - `replace(name, detector)`, `remove`. - -**Feedback** - - -### Step - 7 - -**Description** - -Update the docs and tests to match the new `SemanticRegistry` surface and the -Step 5 fixes. - -- `docs/dqx/docs/reference/profiler.mdx` - * In the "Composing a registry" code block (starting ~line 157), replace the - single `replace(...)` example with usages of the new methods. Illustrate - whole-chain construction via `SemanticRegistry.of(...)` (the new - classmethod) — do NOT show `SemanticRegistry.default().with_detectors([...])` - (that method does not exist per the revised Step 6). Examples to include: - * `SemanticRegistry.default().prepend(uuid_detector)` — highest priority - * `SemanticRegistry.default().append(uuid_detector)` — fallback slot - * `SemanticRegistry.default().insert("enum", uuid_detector)` — after - the enum entry - * `SemanticRegistry.default().replace("enum", custom_enum)` — swap one - * `SemanticRegistry.default().remove("text")` — drop a default - * `SemanticRegistry.of(uuid_detector, *default_semantic_detectors())` — - whole chain from scratch - * `SemanticRegistry()` — empty chain - * Add the comment block requested in comment #7: - ``` - # prepend(detector) → puts one detector at the front (highest priority) - # append(detector) → adds one detector at the back (lowest priority) - # insert(name, detector) → inserts one detector after the named entry - # replace(name, detector) → swaps the entry named `name` - # remove(name) → drops the entry named `name` - # SemanticRegistry.of(...) → build a registry with an explicit chain from scratch - ``` - Place it directly above the code block. - * Rewrite the prose paragraph at lines ~186–189 so it describes all six - composition methods plus the whole-chain constructor/`of(...)` path, - keeps the uniqueness-invariant note, and drops the `replace`-means-whole-chain - wording that prompted comment #8. Explicitly note that `replace` now - swaps a single named detector. - * Regenerate the enum applicability row in the "Built-in detectors" table - (~line 195) so the threshold description matches Step 3: replace the - `ENUM_MAX_CARDINALITY_RATIO (default 0.95)` clause with the aligned - `distinct_ratio` option (default 0.05) and drop the outdated "0.95" number. - -- `tests/unit/profiler/test_semantic.py` - * Delete the existing whole-chain tests that exercised the old `replace([...])` - semantics (`test_semantic_registry_replace_returns_new_instance` and - `test_semantic_registry_replace_duplicate_names_raises`) — the whole-chain - entry point is now `SemanticRegistry(...)` / `SemanticRegistry.of(...)`, - which is already covered by - `test_semantic_registry_direct_constructor_duplicate_names_raises`. - Replace them with: - * `test_semantic_registry_of_returns_new_instance` — asserts - `SemanticRegistry.of(d1, d2)` yields a registry with the given ordered - chain and is a fresh instance. - * `test_semantic_registry_of_duplicate_names_raises` — asserts varargs - duplicates trip the uniqueness validator. - * Add positive + negative tests for each new/renamed method: `append`, - `insert`, `remove(name)`, and `replace(name, detector)`. Cover: - * Ordering — append lands last; insert(name, ...) lands at position n+1 - relative to the named entry; replace(name, ...) preserves position. - * Immutability — original registry unchanged. - * Uniqueness — `append`/`insert`/`replace(name, detector)` reject a - detector whose name clashes with a non-target existing entry (raises - `pydantic.ValidationError`); `replace(name, detector)` allows the new - detector to reuse the target's own name. - * Missing-name errors — `remove`, `insert`, and `replace(name, detector)` - raise `ValueError` for an unknown `name`. - * Adjust the enum-detector positive test (`test_enum_detector_positive_low_cardinality`) - so its `options` dict provides `distinct_ratio` compatible with the new - gate (Step 3). For 3 distinct values / 300 non-null rows, ratio ≈ 0.01, - which passes the default 0.05 threshold. - * Adjust `test_enum_detector_rejects_all_distinct_integer_column` so the - new gate still causes rejection (already does under the default threshold; - verify). - * Add a new test `test_enum_detector_respects_distinct_ratio_option` that - verifies the enum detector suppresses classification when - `distinct_ratio` option is set below the observed ratio (mirrors the - reviewer's 12-rows / 9-distinct example from comment #3). - -- `tests/integration/test_profile_semantic.py` - * Migrate `test_enum_detector_values_reused_by_is_in_builder` (currently - uses the old whole-chain `.replace([...])`) to - `SemanticRegistry.of(tracing, *default_semantic_detectors()[1:])` - (or the plain constructor `SemanticRegistry(detectors=(...))`). - * Grep for any other whole-chain `.replace([...])` call sites in this file - and migrate them the same way. Also grep for `.insert_after(` — should - yield no hits (never released), but confirm and rename if any exist. - * `_make_demo_df` produces `vehicle_type` with 3 distinct values over 60 - rows (ratio 0.05). Under the new default threshold (`distinct_ratio = - 0.05`, strict `>=` rejects), verify semantic-enum still fires by either - (a) tightening the fixture to make cardinality/count_non_null strictly - less than 0.05 (e.g. 3/61 = 0.049) — safer path — or (b) explicitly - passing `options={"distinct_ratio": 0.1, ...}` in the test call. Prefer - (b) to keep the fixture stable. - -- `tests/perf/test_profile_semantic.py` - * Grep for `.replace([`, `.insert_after(`, and `.with_detectors(`. If any - benchmark call uses the old whole-chain `replace`, migrate to - `SemanticRegistry.of(...)` or the constructor. Otherwise no change. - -- `tests/unit/test_profile_builder.py` - * No API change needed here — the profile-builder tests use - `DQProfileContext` directly. Only touch if Step 2's `with_metrics` - change surfaces a test failure (unlikely — no test currently reads - mutated `metrics` through `ctx.metrics`). - -**Feedback** \ No newline at end of file diff --git a/.agents/2026-09-05-12-48-53-review-round-0.md b/.agents/2026-09-05-12-48-53-review-round-0.md deleted file mode 100644 index f55316747..000000000 --- a/.agents/2026-09-05-12-48-53-review-round-0.md +++ /dev/null @@ -1,50 +0,0 @@ ---- -title: The implementation largely matches the plan, but the broadened metadata failure path logs unsanitized exception content and the metrics-refresh fix lacks behavioral regression coverage. -status: changes-requested -created_at: 2026-09-05T13:35:38+02:00 -plan: .agents/2026-09-05-12-48-53-plan.md - -summary: > - The registry API, enum threshold alignment, ShortType support, documentation updates, - and associated tests substantially implement the plan. Follow-up is required because - the broadened exception handler can place uncontrolled exception text into logs, the - central builder-chain behavior fixed by with_metrics is not directly regression-tested, - and a new test suppresses a type-checking error contrary to repository policy. ---- - -# Review - -### Finding - 1 -**Location** -`src/databricks/labs/dqx/profiler/profiler.py:461-466` -**Severity** -critical -**Description** -The handler now catches every Exception raised by the SDK boundary but interpolates the raw exception text into the warning. Only the location is stripped of newlines. Exception messages can contain attacker-controlled table identifiers with newlines or control characters, as well as SDK or transport details that should not be exposed in logs. This violates the repository's log-injection and sensitive-content requirements, and widening the caught exception set increases the range of uncontrolled messages reaching this sink. Sanitize the exception text with the shared log-safe helper or omit it from the message and log only a bounded, non-sensitive exception type through structured fields. -**Address** -Y/N -**Comments** - -### Finding - 2 -**Location** -`src/databricks/labs/dqx/profiler/profiler.py:618-641` -**Severity** -modest -**Description** -The tests added for with_metrics establish only the value object's copy behavior; they do not exercise the bug fixed in the profiler loop. A future change could remove or misplace the refresh at line 623 while every new test continues to pass. Add a behavioral test through DQProfiler's public profiling API with a contextual builder ordered after min_max and assert that it observes the resolved min and max values written back by min_max. This directly guards the reviewer-reported failure and follows the project's requirement to test observable behavior rather than helper implementation details. -**Address** -Y -**Comments** -Added `test_contextual_builder_after_min_max_observes_resolved_min_max` in `tests/integration/test_profile_semantic.py`. The test registers a spy contextual builder via `@register_profile_builder(kind="context")` — insertion order in `PROFILE_BUILDER_REGISTRY` guarantees it runs after `min_max` — and asserts the spy's `ctx.metrics` carries the resolved `min`/`max` written back by `min_max`. Registry cleanup happens in a `finally` block. Runs with `remove_outliers=False` so the expected values are exact. - -### Finding - 3 -**Location** -`tests/unit/profiler/test_semantic.py:96-103` -**Severity** -modest -**Description** -The new frozen-context test adds a type-ignore directive to permit an assignment that the type checker correctly rejects. The repository explicitly prohibits adding type-ignore comments to silence lint or type errors. The same public behavior can be tested without suppressing diagnostics by performing the assignment through setattr and asserting that Pydantic raises ValidationError. -**Address** -Y -**Comments** -Replaced `ctx.column_name = "x" # type: ignore[misc]` with `setattr(ctx, "column_name", "x")` in `test_dq_profile_context_frozen_blocks_field_reassignment`. Public behavior (Pydantic raises `ValidationError` on frozen field reassignment) is unchanged; the type-checker suppression is gone. From 585722507d9a94b603233cc13f2f4bcf94afed1d Mon Sep 17 00:00:00 2001 From: Ivan Kurchenko Date: Sat, 12 Sep 2026 11:59:29 +0200 Subject: [PATCH 9/9] code review feedback addressed --- docs/dqx/docs/reference/profiler.mdx | 5 +-- .../labs/dqx/profiler/profile_builder.py | 7 +++- src/databricks/labs/dqx/profiler/profiler.py | 14 ++++--- src/databricks/labs/dqx/profiler/semantic.py | 7 ++-- tests/integration/test_profile_semantic.py | 32 +++++++++----- tests/integration/test_profiler.py | 42 +++++++++++++++++++ tests/unit/profiler/test_semantic.py | 30 +++++++++++++ 7 files changed, 113 insertions(+), 24 deletions(-) diff --git a/docs/dqx/docs/reference/profiler.mdx b/docs/dqx/docs/reference/profiler.mdx index 1c1e9a8c8..b9563df88 100644 --- a/docs/dqx/docs/reference/profiler.mdx +++ b/docs/dqx/docs/reference/profiler.mdx @@ -145,9 +145,8 @@ checks. For example, a *vehicle_type* column classified as an *enum* now receive *cargo_weight* column is classified as a *measurement* and receives *min_max* only. The feature is fully opt-in: it runs if a `semantic_registry` is supplied to the -`DQProfiler` constructor. When no registry is supplied, the profiler output is -byte-identical to the pre-feature behaviour and every `DQProfile.semantic_type` -field is `None`. +`DQProfiler` constructor. When no registry is supplied, semantic detection does +not run and every `DQProfile.semantic_type` field is `None`. ### Composing a registry diff --git a/src/databricks/labs/dqx/profiler/profile_builder.py b/src/databricks/labs/dqx/profiler/profile_builder.py index c8fe36979..6f7bc7625 100644 --- a/src/databricks/labs/dqx/profiler/profile_builder.py +++ b/src/databricks/labs/dqx/profiler/profile_builder.py @@ -99,8 +99,11 @@ def make_is_in_profile(ctx: DQProfileContext) -> DQProfile | None: When a semantic registry is configured and the column was classified as *enum*, the distinct values collected by the enum detector are reused (no extra Spark action). Columns classified as other semantic types (*key*, *measurement*, - *text*, user-defined) are skipped. When no semantic type is present (default - path), applicability follows today's byte-identical rules. + *text*, user-defined) are skipped. When no semantic type is present, the + builder applies the shared applicability gate: the type must satisfy + *_supports_distinct* and the *distinct_count / count_non_null* ratio must fall + below the configured *distinct_ratio* threshold — the same denominator used by + the semantic-enum detector so the two paths agree on low-repetition columns. Args: ctx: Profile context (column, type, metrics, options, semantic_type). diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index d1913ea7e..072339a32 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -161,7 +161,10 @@ def profile_table( logger.info(f"Profiling {input_config.location} with options: {options}") df = read_input_data(spark=self.spark, input_config=input_config) - return self._profile_dataframe(df=df, columns=columns, options=options) + # Route through *self.profile* (not *_profile_dataframe*) so the nested + # *profile* telemetry event fires alongside *profile_table*. Downstream + # dashboards key on the *profile* event for per-DataFrame counts. + return self.profile(df=df, columns=columns, options=options) @telemetry_logger("profiler", "profile_tables_for_patterns") def profile_tables_for_patterns( @@ -532,10 +535,6 @@ def _build_profiles_for_column( for profile_type in PROFILE_BUILDER_REGISTRY.values(): if profile_type.contextual_builder is not None: - # Refresh the frozen context with the current *metrics* snapshot so contextual - # builders registered after *min_max* observe the resolved min/max values written - # back below (see the write-back block after this loop). - builder_ctx = builder_ctx.with_metrics(metrics) profile = profile_type.contextual_builder(builder_ctx) elif profile_type.builder is not None: profile = profile_type.builder(column_df, field_name, field_type, dict(metrics), dict(opts)) @@ -554,6 +553,11 @@ def _build_profiles_for_column( metrics["min"] = profile.parameters.get("min") if profile.parameters.get("max") is not None: metrics["max"] = profile.parameters.get("max") + # Refresh the frozen context so contextual builders registered after *min_max* + # observe the resolved min/max values just written back. Pydantic materializes + # *ctx.metrics* as a fresh dict at construction, so contextual builders will not + # see mutations to the outer *metrics* dict without an explicit refresh. + builder_ctx = builder_ctx.with_metrics(metrics) def _detect_semantic_type( self, diff --git a/src/databricks/labs/dqx/profiler/semantic.py b/src/databricks/labs/dqx/profiler/semantic.py index f5a532a9c..6d52b240f 100644 --- a/src/databricks/labs/dqx/profiler/semantic.py +++ b/src/databricks/labs/dqx/profiler/semantic.py @@ -327,9 +327,10 @@ def _detect_measurement(ctx: DQProfileContext) -> DQSemanticType | None: stddev_f = float(stddev) mean_f = float(mean) except (TypeError, ValueError): - span = 0.0 - stddev_f = 0.0 - mean_f = 0.0 + # Stat cast failed; classification is not attempted. Keep *distribution* as *unknown* + # rather than falling through the zero-fallback branches below (which would label a + # highly variable column as *constant*). + return DQSemanticType(name="measurement", properties=MeasurementProperties(distribution="unknown")) if span <= 0: distribution = "constant" elif stddev_f == 0: diff --git a/tests/integration/test_profile_semantic.py b/tests/integration/test_profile_semantic.py index c6f56c2a3..504d901cd 100644 --- a/tests/integration/test_profile_semantic.py +++ b/tests/integration/test_profile_semantic.py @@ -58,15 +58,15 @@ def _profile_by_column(profiles): return grouped -def test_profile_without_semantic_registry_matches_pre_feature_output(spark, ws): - """No registry → profiler output is identical to the pre-feature behaviour. +def test_profile_without_semantic_registry_leaves_semantic_type_unset(spark, ws): + """No registry → semantic detection does not run. Every DQProfile.semantic_type must be None because detection did not run. """ df = _make_demo_df(spark) profiler = DQProfiler(ws) _stats, profiles = profiler.profile(df, options={"sample_fraction": None, "llm_primary_key_detection": False}) - assert profiles, "expected the pre-feature profiler to emit at least one profile" + assert profiles, "expected the profiler to emit at least one profile" tagged = [(p.column, p.name, p.semantic_type) for p in profiles if p.semantic_type is not None] assert not tagged, f"expected all profiles to have semantic_type=None, got tagged: {tagged}" @@ -192,13 +192,15 @@ def test_short_type_low_cardinality_classified_as_enum_and_emits_is_in(spark, ws def test_contextual_builder_after_min_max_observes_resolved_min_max(spark, ws): - """A contextual builder ordered after `min_max` observes resolved `min`/`max` via `ctx.metrics`. - - Regression guard: `DQProfiler._build_profiles_for_column` calls - `builder_ctx.with_metrics(metrics)` before each contextual builder so the frozen context - reflects write-backs performed by earlier builders. Removing or misplacing that refresh - would silently return stale metrics to downstream builders — a latent contract break for - custom builders that depend on `min_max`'s write-back. + """A contextual builder ordered after `min_max` observes the *changed* `min`/`max` via `ctx.metrics`. + + Regression guard: after `min_max` writes its resolved bounds back into *metrics*, + `DQProfiler._build_profiles_for_column` refreshes the frozen builder context with + `builder_ctx.with_metrics(metrics)` so contextual builders registered later see the + new values instead of the pre-`min_max` summary stats. Uses `remove_outliers=True` + with a tight `num_sigmas` so the `min_max` builder emits sigma-capped bounds strictly + inside the raw column range — proving that what the spy observes came from the + write-back, not from the initial summary stats. """ seen_metrics: list[dict] = [] @@ -210,19 +212,27 @@ def _profile_and_assert(): schema = T.StructType([T.StructField("value", T.LongType())]) rows = [(i,) for i in range(100)] df = spark.createDataFrame(rows, schema=schema) + raw_min, raw_max = 0, 99 profiler = DQProfiler(ws) _stats, profiles = profiler.profile( df, options={ "sample_fraction": None, "llm_primary_key_detection": False, - "remove_outliers": False, + "remove_outliers": True, + # Sigma capping with mean≈49.5, stddev≈29 and 0.5 sigmas keeps bounds strictly + # inside [0, 99], so the resolved min/max differ from the raw column bounds. + "num_sigmas": 0.5, + "round": False, }, ) min_max_profiles = [p for p in profiles if p.column == "value" and p.name == "min_max"] assert len(min_max_profiles) == 1, f"expected one min_max profile on 'value', got: {min_max_profiles}" expected_min = min_max_profiles[0].parameters["min"] expected_max = min_max_profiles[0].parameters["max"] + assert ( + raw_min < expected_min < expected_max < raw_max + ), f"expected sigma-capped bounds strictly inside ({raw_min}, {raw_max}), got ({expected_min}, {expected_max})" assert seen_metrics, "spy contextual builder was not invoked" spy_seen = seen_metrics[-1] assert ( diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index cdd7fc799..0af8babb2 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1,6 +1,7 @@ import dataclasses from datetime import date, datetime, timezone from decimal import Decimal +import logging import pytest import pyspark.sql.types as T @@ -9,6 +10,7 @@ from databricks.labs.dqx.config import InputConfig, LLMModelConfig from databricks.labs.dqx.errors import InvalidConfigError from databricks.labs.dqx.profiler.profiler import DQProfiler, DQProfile +from databricks.labs.dqx import telemetry from tests.constants import TEST_CATALOG @@ -923,6 +925,46 @@ def test_profile_table(spark, ws, make_schema, make_random): assert profiles == expected_profiles +def test_profile_table_emits_nested_profile_telemetry(spark, ws, make_schema, make_random, caplog): + """*profile_table* must fire both the outer *profile_table* and the nested *profile* signals. + + Downstream dashboards key on the *profile* event for per-DataFrame counts; if + *profile_table* were to bypass the decorated *profile* entry point, those counts + would silently drop. Observed at the SDK boundary via the debug log + *"Added User-Agent extra ="* that *log_telemetry* emits immediately + before stamping the header on the workspace call — a legitimate external observation + that does not patch any DQX symbol. + """ + catalog_name = TEST_CATALOG + schema_name = make_schema(catalog_name=catalog_name).name + table_name = f"{catalog_name}.{schema_name}.t{make_random(10).lower()}" + + input_schema = T.StructType([T.StructField("id", T.IntegerType())]) + spark.createDataFrame([[1], [2], [3]], schema=input_schema).write.format("delta").saveAsTable(table_name) + + # Reset the per-process dedup cache so previously-sent signals do not suppress the ones + # this test asserts on. + telemetry.reset_telemetry_cache() + try: + with caplog.at_level(logging.DEBUG, logger="databricks.labs.dqx.telemetry"): + profiler = DQProfiler(ws) + profiler.profile_table( + input_config=InputConfig(location=table_name), + options={"sample_fraction": None, "llm_primary_key_detection": False}, + ) + + prefix = "Added User-Agent extra " + emitted = { + record.getMessage()[len(prefix) :] for record in caplog.records if record.getMessage().startswith(prefix) + } + assert ( + "profiler=profile_table" in emitted + ), f"expected 'profiler=profile_table' User-Agent extra, got: {emitted}" + assert "profiler=profile" in emitted, f"expected nested 'profiler=profile' User-Agent extra, got: {emitted}" + finally: + telemetry.reset_telemetry_cache() + + def test_profile_table_non_default_opts(spark, ws, make_schema, make_random): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name diff --git a/tests/unit/profiler/test_semantic.py b/tests/unit/profiler/test_semantic.py index 110ffa52b..b4efbddb8 100644 --- a/tests/unit/profiler/test_semantic.py +++ b/tests/unit/profiler/test_semantic.py @@ -336,6 +336,36 @@ def test_measurement_detector_rejects_string_column(): assert DEFAULT_MEASUREMENT_DETECTOR.detect(ctx) is None +@pytest.mark.parametrize( + "bad_stat", + [ + pytest.param("not-a-number", id="string-triggers-value-error"), + pytest.param(complex(1, 2), id="complex-triggers-type-error"), + ], +) +def test_measurement_detector_malformed_stats_yields_unknown_distribution(bad_stat): + # min/max/mean/stddev present but not castable to float — classification cannot proceed. + # Regression guard: previously the zero-fallback path drove *distribution* to *"constant"* + # for a highly variable column; it must be *"unknown"* instead. Parametrized so both + # exception paths caught by the detector (ValueError from strings, TypeError from complex) + # exercise the *"unknown"* return. + ctx = _fake_ctx( + column_type=T.IntegerType(), + metrics={ + "count_non_null": 100, + "min": bad_stat, + "max": bad_stat, + "mean": bad_stat, + "stddev": bad_stat, + }, + ) + result = DEFAULT_MEASUREMENT_DETECTOR.detect(ctx) + assert result is not None + assert result.name == "measurement" + assert isinstance(result.properties, MeasurementProperties) + assert result.properties.distribution == "unknown" + + def test_text_detector_positive_string_column(): ctx = _fake_ctx(column_type=T.StringType(), metrics={"count_non_null": 5}) result = DEFAULT_TEXT_DETECTOR.detect(ctx)