diff --git a/docs/dqx/docs/reference/profiler.mdx b/docs/dqx/docs/reference/profiler.mdx index 5fe8e7a99..3a1bd5710 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 @@ -54,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) | @@ -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 @@ -173,4 +175,165 @@ See the [Data Profiling Guide](/docs/guide/data_profiling#extending-the-profiler 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 if a `semantic_registry` is supplied to the +`DQProfiler` constructor. When no registry is supplied, semantic detection does +not run and every `DQProfile.semantic_type` field is `None`. + +### Composing a registry + +`SemanticRegistry` is an immutable, ordered, name-unique collection of detectors. +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 +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. 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()) + +# 8. Empty registry — semantic types are always None +empty = SemanticRegistry() + +profiler = DQProfiler(ws, semantic_registry=default_registry) +``` + +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 (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) | + +`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) +``` + +### 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 +`kind="context"` argument to opt into the preferred contextual callback shape, +which receives a `DQProfileContext` (giving access to *ctx.semantic_type*, +*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", kind="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..adf78e8aa 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 +# kind="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 *kind="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 471385ae6..92c547d97 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 +from typing import Any, Literal from pyspark.sql import DataFrame from pyspark.sql import types as T, functions as F @@ -12,6 +12,7 @@ from databricks.labs.dqx.errors import InvalidParameterError from databricks.labs.dqx.profiler.common import TEXT_TYPES, is_text from databricks.labs.dqx.profiler.profile import DQProfile, DQProfileBuilder +from databricks.labs.dqx.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, @@ -42,9 +43,31 @@ logger = logging.getLogger(__name__) -def register_profile_builder(profile_type: str) -> Callable: +def register_profile_builder( + profile_type: str, + *, + kind: Literal["legacy", "context"] | None = None, +) -> Callable: + """Register a profile builder in *PROFILE_BUILDER_REGISTRY*. + + Args: + profile_type: Registry key (e.g. *min_max*). + 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. + * *"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 kind == "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 @@ -61,124 +84,131 @@ def deregister_profile_builder(profile_type: str) -> None: PROFILE_BUILDER_REGISTRY.pop(profile_type, None) -@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", 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 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", kind="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, 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: - 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) - if total_count == 0: + count_non_null = ctx.metrics.get("count_non_null", 0) + if count_non_null == 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: + 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 + # 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, + 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, # 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( 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", kind="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(column_type): + if not _supports_min_max(ctx.column_type): return None - if _remove_outliers(column_name, profiler_options): + if ctx.semantic_type is not None and ctx.semantic_type.name != "measurement": + return None + + 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) ) 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. @@ -251,7 +281,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. @@ -290,13 +320,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: @@ -314,7 +348,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. @@ -335,7 +369,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. @@ -360,7 +394,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. @@ -771,45 +805,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", 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. 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" @@ -833,7 +865,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( @@ -848,7 +880,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 9a88667ee..7c053726a 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -1,6 +1,7 @@ -import uuid +import dataclasses import logging import os +import uuid from concurrent import futures from decimal import Decimal, Context from difflib import SequenceMatcher @@ -11,13 +12,15 @@ import pyspark.sql.types as T from pyspark.errors import AnalysisException from pyspark.sql import DataFrame, SparkSession + from databricks.sdk import WorkspaceClient 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.common import TEXT_TYPES, is_text +from databricks.labs.dqx.profiler.common import TEXT_TYPES +from databricks.labs.dqx.profiler.common import is_text from databricks.labs.dqx.profiler.profile import DQProfile from databricks.labs.dqx.profiler.profile_builder import PROFILE_BUILDER_REGISTRY, validate_profile_options from databricks.labs.dqx.profiler.profile_options import ( @@ -34,8 +37,13 @@ from databricks.labs.dqx.profiler.profiler_column_metrics import ( build_registered_metric_aggregations, ) -from databricks.labs.dqx.utils import list_tables +from databricks.labs.dqx.profiler.semantic import ( + DQProfileContext, + DQSemanticType, + SemanticRegistry, +) from databricks.labs.dqx.telemetry import telemetry_logger +from databricks.labs.dqx.utils import list_tables try: from databricks.labs.dqx.llm.llm_pk_engine import DQLLMPrimaryKeyEngine @@ -55,6 +63,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 @@ -63,6 +73,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]: @@ -102,6 +113,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) + + def _profile_dataframe( + self, + df: DataFrame, + columns: list[str] | None, + options: dict[str, Any] | None, + ) -> tuple[dict[str, Any], list[DQProfile]]: + """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]) @@ -144,6 +164,9 @@ def profile_table( logger.info(f"Profiling {input_config.location} with options: {options}") df = read_input_data(spark=self.spark, input_config=input_config) + # 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") @@ -520,10 +543,28 @@ 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. """ + semantic_type = self._detect_semantic_type(column_df, field.name, field.dataType, metrics, opts) + + builder_ctx = DQProfileContext( + df=column_df, + column_name=field.name, + column_type=field.dataType, + metrics=metrics, + options=opts, + semantic_type=semantic_type, + ) + for profile_type in PROFILE_BUILDER_REGISTRY.values(): - profile = profile_type.builder(column_df, field.name, field.dataType, 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.dataType, dict(metrics), dict(opts)) + else: + continue 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. @@ -532,6 +573,35 @@ 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, + column_df: DataFrame, + field_name: str, + field_type: T.DataType, + metrics: dict[str, Any], + opts: 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, + 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] diff --git a/src/databricks/labs/dqx/profiler/semantic.py b/src/databricks/labs/dqx/profiler/semantic.py new file mode 100644 index 000000000..6d52b240f --- /dev/null +++ b/src/databricks/labs/dqx/profiler/semantic.py @@ -0,0 +1,514 @@ +"""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* + * 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()* + * thresholds: *ENUM_MAX_CARDINALITY_RATIO*, *KEY_MIN_DENSITY_RATIO*, + *KEY_MIN_LENGTH_STABILITY_RATIO* +""" + +import logging +from collections.abc import Callable, Mapping +from typing import Any, Literal, TypeVar + +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.common import TEXT_TYPES +from databricks.labs.dqx.profiler.profile_options import PROFILE_OPTION_DISTINCT_RATIO, PROFILE_OPTION_MAX_IN_COUNT + +logger = logging.getLogger(__name__) + + +Scalar = str | int | float | bool + + +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 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. + + 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 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: 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): + """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, ...). + 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) + 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. + + 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: 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 +# 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 + + +def _is_numeric(column_type: DataType) -> bool: + return isinstance(column_type, T.NumericType) + + +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 + + 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] + distinct_rows = ctx.df.select(col).distinct().collect() + distinct_values = {row[0] for row in distinct_rows} + + return DQSemanticType(name="enum", properties=EnumProperties(values=distinct_values)) + + +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)): + return _detect_numeric_key(ctx, cardinality) + + if _is_text(column_type): + 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") + + +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") + + +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: 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) + stddev_f = float(stddev) + mean_f = float(mean) + except (TypeError, ValueError): + # 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: + 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=MeasurementProperties(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, *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 or *SemanticRegistry.of(...)* to build + an arbitrary chain from scratch. + """ + + 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()) + + @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 (highest priority). + + 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 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. + """ + 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. + """ + 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_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/integration/test_profile_semantic.py b/tests/integration/test_profile_semantic.py new file mode 100644 index 000000000..504d901cd --- /dev/null +++ b/tests/integration/test_profile_semantic.py @@ -0,0 +1,279 @@ +import uuid + +import pyspark.sql.types as T + +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, + DQSemanticType, + DQSemanticTypeDetector, + SemanticRegistry, + default_semantic_detectors, +) + + +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), + ) + ) + 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_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 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}" + + +def test_default_semantic_registry_classifies_grounded_columns(spark, ws): + df = _make_demo_df(spark) + profiler = DQProfiler(ws, semantic_registry=SemanticRegistry.default()) + # 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) + + vehicle_names = {p.name for p in by_column.get("vehicle_type", [])} + 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, 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, 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, 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 has variable length → falls through to text (not key). text emits no additional + # rules; only null_or_empty candidates remain. + 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", []) + 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): + 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, 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): + """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 + ), 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_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. + + 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} + ) + by_column = _profile_by_column(profiles) + + 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 test_contextual_builder_after_min_max_observes_resolved_min_max(spark, ws): + """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] = [] + + @register_profile_builder("_spy_after_min_max", kind="context") + def _spy(ctx): + seen_metrics.append(dict(ctx.metrics)) + + 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": 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 ( + 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')}" + + try: + _profile_and_assert() + finally: + PROFILE_BUILDER_REGISTRY.pop("_spy_after_min_max", None) + + +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.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, "distinct_ratio": 0.1} + ) + + 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 + ), f"expected exactly one 'is_in' profile for vehicle_type, got {len(is_in_profiles)}: {is_in_profiles}" + is_in = is_in_profiles[0] + values = set(is_in.parameters["in"]) + assert values == {"car", "truck", "van"}, f"expected is_in values {{car, truck, van}}, got: {values}" diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index efcb89ff5..173116852 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 @@ -19,6 +20,7 @@ PROFILE_BUILDER_REGISTRY, register_profile_builder, ) +from databricks.labs.dqx import telemetry from tests.constants import TEST_CATALOG @@ -1109,6 +1111,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/perf/test_profile_semantic.py b/tests/perf/test_profile_semantic.py new file mode 100644 index 000000000..5ecff49ca --- /dev/null +++ b/tests/perf/test_profile_semantic.py @@ -0,0 +1,93 @@ +"""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 + + +# 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 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} + +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, + } + ], + 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/__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..b4efbddb8 --- /dev/null +++ b/tests/unit/profiler/test_semantic.py @@ -0,0 +1,514 @@ +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, + EnumProperties, + MeasurementProperties, + SemanticRegistry, + default_semantic_detectors, +) + + +# --------------------------------------------------------------------------- +# DQSemanticType immutability (Pydantic v2 guarantees) +# --------------------------------------------------------------------------- + + +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(MeasurementProperties) is None + + +def test_dq_semantic_type_frozen_blocks_field_reassignment(): + 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): + 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} + df_mock = create_autospec(DataFrame) + ctx = DQProfileContext( + df=df_mock, + 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} + df_mock = create_autospec(DataFrame) + ctx = DQProfileContext( + df=df_mock, + 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(): + df_mock = create_autospec(DataFrame) + ctx = DQProfileContext( + df=df_mock, + column_name="c", + column_type=T.IntegerType(), + ) + with pytest.raises(ValidationError): + setattr(ctx, "column_name", "x") + + +# --------------------------------------------------------------------------- +# 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 not 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_of_returns_new_instance_with_ordered_chain(): + 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"] + + +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(): + 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_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)) + + +# --------------------------------------------------------------------------- +# SemanticRegistry mutation helpers (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.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") + + +# --------------------------------------------------------------------------- +# 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 isinstance(result.properties, MeasurementProperties) + assert result.properties.distribution in {"constant", "uniform", "normal", "exponential", "unknown"} + + +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 + + +@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) + 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 is None + + +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" + typed = result.get_typed_properties(EnumProperties) + assert typed is not None + 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) +# --------------------------------------------------------------------------- + + +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 is None + + +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 d2a020a9c..12cb94753 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, deregister_profile_builder, @@ -17,6 +18,7 @@ register_profile_builder, validate_profile_options, ) +from databricks.labs.dqx.profiler.semantic import DQProfileContext, DQSemanticType, EnumProperties @pytest.fixture @@ -39,6 +41,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 # --------------------------------------------------------------------------- @@ -60,8 +73,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) @@ -115,19 +130,61 @@ def test_deregister_profile_builder_missing_key_is_noop(restore_profile_builder_ assert "_never_registered_key" not in PROFILE_BUILDER_REGISTRY +def test_register_profile_builder_context_type_uses_contextual_slot(): + @register_profile_builder("_test_ctx", kind="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", kind="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 @@ -136,11 +193,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" @@ -150,11 +209,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" @@ -163,11 +224,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" @@ -179,22 +242,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} @@ -202,11 +269,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" @@ -214,11 +283,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 @@ -230,22 +295,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" @@ -255,11 +312,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 @@ -267,11 +320,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" @@ -297,37 +346,51 @@ 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, "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"]) - profile = make_is_in_profile(df, "col", column_type, {"count": 10}, {"max_in_count": 10, "distinct_ratio": 1.0}) +@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, "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(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(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, "count_non_null": 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, "count_non_null": 5}, + {"max_in_count": 10, "distinct_ratio": 1.0}, + ) ) assert profile is not None assert profile.name == "is_in" @@ -339,24 +402,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( - 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( - 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 @@ -364,37 +431,75 @@ 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, "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" +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=EnumProperties(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, "count_non_null": 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" @@ -405,11 +510,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} @@ -417,11 +520,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" @@ -433,11 +538,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" @@ -447,11 +548,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" @@ -463,11 +566,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} @@ -477,11 +582,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} @@ -492,11 +599,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 @@ -505,11 +614,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 @@ -520,11 +631,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 @@ -533,17 +646,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 # --------------------------------------------------------------------------- @@ -552,14 +697,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 @@ -571,7 +716,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 @@ -590,7 +735,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 @@ -602,7 +747,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 @@ -615,7 +760,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 @@ -626,11 +771,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 @@ -640,22 +787,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 @@ -666,11 +817,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 @@ -681,14 +834,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"]}, + ) ) @@ -698,11 +850,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 @@ -717,7 +865,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 @@ -731,7 +879,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