diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 497191b50..7fc2ceb44 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1963,6 +1963,7 @@ You can also define your own custom dataset-level checks (see [Creating custom c | `has_no_gaps_per_time_window` | Dataset check that flags gaps in a time series, i.e. time windows of a given size that contain no rows between windows that do. The violation is reported on the boundary row before each interior gap. | `column`: timestamp or date column (can be a string column name or a column expression); `window_minutes`: size of the time window in minutes that defines the expected data grain (for example 1440 for daily); `group_by`: optional list of columns or column expressions to detect gaps independently within each group; `trailing_gap`: (optional) if `true`, also flags the last present window (per group) when it ends more than one window before `curr_timestamp`, so missing recent data is caught at the tail of the series (defaults to `false`); `curr_timestamp`: (optional) current timestamp column used to anchor trailing-gap detection, only used when `trailing_gap` is `true` (if not provided, current_timestamp() function is used) | | `has_valid_schema` | Schema check that validates whether the DataFrame schema matches an expected schema. In non-strict mode, validates that all expected columns exist with compatible types (allows extra columns). In strict mode, validates exact schema match (same columns, same order, same types) for all columns by default or for all columns specified in `columns`. This check is applied at the dataset level and reports schema violations for all rows in the DataFrame when incompatibilities are detected. All columns in the `exclude_columns` list will be ignored even if the column is present in the `columns` list. | `expected_schema`: (optional) expected schema as a DDL string (e.g., "id INT, name STRING") or StructType object; `ref_df_name`: (optional) name of the reference DataFrame to load the schema from (dictionary of DataFrames can be passed when applying checks); `ref_table`: (optional) fully qualified reference table name to load the schema from (e.g. "catalog.schema.table"); exactly one of `expected_schema`, `ref_df_name`, or `ref_table` must be provided; `columns`: (optional) list of columns to validate (if not provided, all columns are considered); `strict`: (optional) whether to perform strict schema validation (default: False) - False: validates that all expected columns exist with compatible types, True: validates exact schema match; `exclude_columns`: (optional) list of columns to ignore during validation (if not provided, all columns are considered); | | `has_no_outliers` | Checks whether the values in the input column contain any outliers. This function implements a median absolute deviation (MAD) algorithm to find outliers. | `column`: column of type numeric to check (can be a string column name or a column expression); | +| `is_in_distribution` | Checks whether the discrete value distribution of the column is within a given distance of an expected distribution, using the **Total Variation Distance (TVD)** of probability measures. Supplements `is_in_list` at the dataset level by validating not only that values stay within an enumerated set, but also that they follow the expected proportions. Supported column types: Boolean, Char, String, Byte, Short, Integer, Long, and Date. NULL values in the checked column are skipped. When the sum of the expected distribution values is strictly less than 1, the missing mass is placed in an internal `residual` bucket, and any column values not listed as explicit keys are aggregated into the same bucket. | `column`: column to check (can be a string column name or a column expression); `distribution`: expected distribution as a dict mapping literal values to their expected proportions; `distance`: maximum allowed TVD between the actual and expected distributions; `case_sensitive`: (optional) whether string comparisons are case-sensitive (default: `True`); `impute`: (optional) if `True` (default), keys present in the expected distribution but missing from the actual distribution are imputed with `0`; if `False`, the check fails and reports the missing keys | | `has_no_row_anomalies` | Flags rows that are anomalous according to a trained ML model. The model learns "normal" patterns from your training data; at check time each row is scored (severity percentile 0–100) and optionally enriched with SHAP contributions. Requires a model trained with the anomaly engine first. See [Row Anomaly Detection](#row-anomaly-detection) below for training, full parameters, and usage. | `model_name`: fully qualified model name (e.g. catalog.schema.model_name); `registry_table`: fully qualified registry table (e.g. catalog.schema.model_registry); `threshold`: (optional) severity percentile threshold (default 95); `drift_threshold`: (optional) warn when score distribution drifts from training (None = off); `enable_contributions`: (optional) add SHAP per-feature contributions to `_dq_info` (default True; set False to skip the SHAP cost); `enable_confidence_std`: (optional) add ensemble score std to `_dq_info` (default False); `enable_ai_explanation`: (optional) add an LLM-generated explanation to `_dq_info` (default True; degrades to null if contributions are off or no serving endpoint is reachable); `ai_explanation_llm_model_config`: (optional) Databricks Model Serving endpoint config for the explanation; `redact_columns`: (optional) feature/segment names to keep out of the LLM prompt; `max_groups`: (optional) cap on LLM calls per run (default 500). See [Row Anomaly Detection](/docs/reference/quality_checks#row-anomaly-detection) section for full parameter details. | | `are_polygons_mutually_disjoint` | Checks whether the polygons in a geometry column are mutually disjoint. Polygons sharing an edge or boundary are considered intersecting. Nulls and invalid geometries are excluded from the check. Requires Databricks runtime 17.1 or above. | `column`: column to check (can be a string column name or a column expression), must contain polygon or multipolygon geometries | | `is_geo_contains` | Checks if the reference geometry contains each column geometry using `st_contains` with meter-level precision. A geometry A *contains* B when B lies entirely within the interior of A with no boundary points of B on the boundary of A. Points on the shared boundary are not considered contained — use `is_geo_covers` for boundary-inclusive checks. When a convert flag is set to `True`, `try_to_geometry` is applied to parse the input from any supported format (WKT, WKB, EWKT, EWKB). Null values are skipped. Requires Databricks runtime 17.1 or above. | `column`: column to check (can be a string column name or a column expression); `reference_geometry`: reference geometry as a literal WKT/WKB/EWKT/EWKB string or bytes value, or a `Column` expression (e.g. `F.col('col_name')`) — a plain string is always treated as a literal, not a column name; `convert_column`: when `True`, applies `try_to_geometry` to convert the column values to GEOMETRY (default `False`); `convert_reference_geometry`: when `True`, applies `try_to_geometry` to convert the reference geometry to GEOMETRY (default `False`) | @@ -2487,6 +2488,50 @@ Complex data types are supported as well. arguments: column: col1 +# is_in_distribution check — validates the actual value distribution of a categorical +# column against an expected distribution using Total Variation Distance (TVD). +# Fails when TVD(actual, expected) > distance. +- criticality: error + check: + function: is_in_distribution + arguments: + column: status + distribution: + active: 0.7 + inactive: 0.2 + pending: 0.1 + distance: 0.05 # max allowed TVD + +# is_in_distribution check — expected distribution intentionally covers only a subset +# of the possible values (sum < 1); case-insensitive comparison. The missing 0.1 mass +# is placed in an internal `residual` bucket; any column values other than A/B/C +# are aggregated into `residual` and compared against that 0.1 expected mass. +- criticality: warn + check: + function: is_in_distribution + arguments: + column: category + distribution: + A: 0.6 + B: 0.2 + C: 0.1 # remaining 0.1 goes to an internal `residual` bucket for all other values + distance: 0.1 + case_sensitive: false + +# is_in_distribution check — fail (and enumerate missing keys) if any expected key +# is absent from the actual distribution, instead of imputing it with 0. +- criticality: error + check: + function: is_in_distribution + arguments: + column: tier + distribution: + gold: 0.2 + silver: 0.3 + bronze: 0.5 + distance: 0.05 + impute: false # fail and report missing keys instead of imputing with 0 + # are_polygons_mutually_disjoint check (geo, requires runtime 17.1+) - criticality: error check: @@ -3094,6 +3139,47 @@ checks = [ column="col1" # or as expr: F.col("col1") ), + # is_in_distribution check — validates the actual value distribution of a categorical + # column against an expected distribution using Total Variation Distance (TVD). + # Fails when TVD(actual, expected) > distance. + DQDatasetRule( + criticality="error", + check_func=check_funcs.is_in_distribution, + column="status", # or as expr: F.col("status") + check_func_kwargs={ + "distribution": {"active": 0.7, "inactive": 0.2, "pending": 0.1}, + "distance": 0.05, # max allowed TVD + } + ), + + # is_in_distribution check — expected distribution intentionally covers only a subset + # of the possible values (sum < 1); case-insensitive comparison. The missing 0.1 mass + # is placed in an internal `residual` bucket; any column values other than A/B/C + # are aggregated into `residual` and compared against that 0.1 expected mass. + DQDatasetRule( + criticality="warn", + check_func=check_funcs.is_in_distribution, + column="category", # or as expr: F.col("category") + check_func_kwargs={ + "distribution": {"A": 0.6, "B": 0.2, "C": 0.1}, # remaining 0.1 goes to internal `residual` bucket + "distance": 0.1, + "case_sensitive": False, + } + ), + + # is_in_distribution check — fail (and enumerate missing keys) if any expected key + # is absent from the actual distribution, instead of imputing it with 0. + DQDatasetRule( + criticality="error", + check_func=check_funcs.is_in_distribution, + column="tier", # or as expr: F.col("tier") + check_func_kwargs={ + "distribution": {"gold": 0.2, "silver": 0.3, "bronze": 0.5}, + "distance": 0.05, + "impute": False, # fail and report missing keys instead of imputing with 0 + } + ), + # are_polygons_mutually_disjoint check (geo, requires runtime 17.1+) DQDatasetRule( criticality="error", diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index ad8d9d39d..b8c714499 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1,10 +1,11 @@ import datetime +import math import re import warnings import ipaddress import uuid from decimal import Decimal -from functools import lru_cache +from functools import lru_cache, reduce from importlib.resources import files from collections.abc import Callable, Sequence from enum import Enum @@ -15,6 +16,7 @@ import pyspark.sql.functions as F from pyspark.sql import types from pyspark.sql import Column, DataFrame, SparkSession +from pyspark.sql.types import DataType from pyspark.sql.window import Window from databricks.labs.dqx.profiling_utils import calculate_median_absolute_deviation_bounds @@ -83,6 +85,18 @@ # Future: Add other aggregates that don't work with windows (e.g., collect_set with DISTINCT) } +_IS_IN_DISTRIBUTION_SUPPORTED_KEY_TYPES: tuple[type, ...] = (bool, str, int, datetime.date) +_IS_IN_DISTRIBUTION_SUPPORTED_SPARK_TYPES: tuple[type, ...] = ( + types.BooleanType, + types.StringType, + types.CharType, + types.ByteType, + types.ShortType, + types.IntegerType, + types.LongType, + types.DateType, +) + class DQPattern(Enum): """Enum class to represent DQ patterns used to match data in columns.""" @@ -502,6 +516,374 @@ def is_in_list(column: str | Column, allowed: list, case_sensitive: bool = True) ) +@register_rule("dataset") +def is_in_distribution( + column: str | Column, + distribution: dict[bool, float] | dict[str, float] | dict[int, float] | dict[datetime.date, float], + distance: float, + case_sensitive: bool = True, + impute: bool = True, + row_filter: str | None = None, +) -> tuple[Column, Callable]: + """ + Check whether the discrete value distribution of the column is within a distance less than or equal to the given + distance, compared to the given values distribution. The distance is calculated using the + [Total variation distance of probability measures](https://en.wikipedia.org/wiki/Total_variation_distance_of_probability_measures) + method. + + This check aims to supplement *is_in_list* at the dataset level by verifying that the column not only stays within + the boundaries of enumerated values, but also that these values follow the expected distribution. + + Supported column types: Boolean, Char, String, Byte, Short, Integer, Long, and Date. These types have been chosen + because they are the most suitable for representing categorical data and produce stable, well-defined distributions + when aggregated. Floating-point types (Float, Double) are excluded because equality-based grouping is unstable due + to precision issues; complex types (Array, Struct, Variant) are excluded because they lack a natural + equality-based grouping semantics; and potentially large types (Binary) are excluded because they can lead to + performance and correctness issues when used as grouping keys. For any other column type a validation error is + raised. + + The sum of the supplied distribution values must be less than or equal to 1. The distance must also be between 0 + and 1 (inclusive). + The distribution is not allowed to contain *None* — neither as a key nor as a value; please consider using the + corresponding completeness functions for that purpose. + The distribution dict is not allowed to contain *inf*, *-inf*, or *nan* among its values, nor negative values. + The distribution dict must be homogeneous: all keys must be of the same supported primitive type (matching the + column type). Mixing key types (e.g., a dict containing both string and integer keys) raises a validation error. + The values distribution has no size limitation, but it is highly recommended to keep it reasonably small as the + check is intended to be used for enumerated data. + + The actual distribution is calculated over all non-null values in the column (equivalent to + *df.groupBy(col).count()*). *NULL* values in the checked column are skipped, so the distribution is computed over + non-null values only. + + A sum strictly less than 1 is allowed when the intent is to check only a subset of the possible values (e.g., + some enumerated values are intentionally omitted from the expected distribution). To calculate the TVD properly in + this case, the function creates an internal *residual* bucket in the expected distribution to hold the missing + mass (1 - sum), and the actual distribution is computed with the same bucketing: all non-null column values not + listed as explicit keys in the given distribution are aggregated into the *residual* bucket. For example, given a + distribution of {A: 0.5, B: 0.3}, an internal *residual* bucket holds the remaining 0.2; the actual distribution + is then calculated over *A*, *B*, and *residual*, where *residual* counts every non-null column value other than + *A* and *B*. When the sum equals 1, the expected *residual* mass is 0 and any actual values not listed in the + given distribution will contribute to the distance. + + Handling of keys present in the given distribution but missing from the actual distribution depends on the + *impute* boolean parameter. + If it is *True*, such keys are imputed with a value of 0 in the actual distribution, so that the distance can be + calculated across the union of keys. + If it is *False*, the check fails with an error specifying which keys from the given distribution are missing in + the actual distribution. Dataset-level checks apply to the whole dataset (not per-row), so the failure applies to + all rows. + *NOTE*: values present in the actual distribution but not listed as explicit keys in the given distribution are + aggregated into the *residual* bucket rather than compared individually. Consider using *is_in_list* at the row + level for checking individual values. + + Case normalisation is applied when *case_sensitive* is *False*. The same normalisation is applied to the keys of + the given *distribution* before the check runs, so both sides are compared consistently. If two distribution keys + collide after normalisation a validation error is raised. + + Note: + This check is not supported for streaming DataFrames. It performs an ungrouped aggregation over the whole + dataset to compute the actual distribution, which Structured Streaming does not allow on unbounded sources + without a watermark and a compatible output mode. Use it on batch DataFrames only. + + Args: + column: column to check; can be a string column name or a column expression. + distribution: expected distribution of literal values to compare with. The dict must be homogeneous — all + keys must be of the same supported primitive type matching the column type; values must be non-negative + and sum to 1 or less. + distance: max distance between the actual and given values distribution; must be between 0 and 1 (inclusive). + case_sensitive: whether to perform a case-sensitive comparison (default: True). + impute: whether to substitute keys missing from the actual distribution with 0 (True) or fail the check + (False). + row_filter: Optional SQL expression for filtering rows before the distribution is computed. Auto-injected + from the check filter. + + Returns: + A tuple of: + - A Spark Column representing the condition for distribution violations. + - A closure that applies the distribution check and adds the necessary condition/count columns. + + Raises: + MissingParameterError: If the distribution or distance is not provided. + InvalidParameterError: If the distribution parameter is not a dict or is empty; if the distribution dict + contains a *None* key or *None* value; if the distribution values contain *inf*, *-inf*, or *nan*; if any distribution + value is negative; if the sum of the distribution values is greater than 1; if the distance parameter + value is not between 0 and 1 (inclusive); if the column type is not one of the supported primitive + types; if the distribution dict is not homogeneous (contains keys of more than one type); or if two keys + in the given distribution collide after case normalisation when *case_sensitive* is *False*. + """ + _is_in_distribution_validate_distribution(distribution, case_sensitive) + _is_in_distribution_validate_distance(distance) + + col_str_norm, col_expr_str, col_expr = get_normalized_column_and_expr(column) + + unique_str = uuid.uuid4().hex + condition_col = f"__condition_{col_str_norm}_{unique_str}" + message_col = f"__message_{col_str_norm}_{unique_str}" + + def apply(df: DataFrame) -> DataFrame: + # Resolve the type off the expression itself + column_type = _is_in_distribution_get_data_type(df, col_expr, col_expr_str) + + filtered = df.filter(safe_filter_expr(row_filter)) if row_filter else df + + is_string_column = isinstance(column_type, (types.StringType, types.CharType)) + group_expr = col_expr + normalized_distribution: dict[Any, float] = distribution + if is_string_column and not case_sensitive: + group_expr = F.lower(col_expr) + normalized_distribution = {k.lower() if isinstance(k, str) else k: v for k, v in distribution.items()} + + expected_keys_ordered = list(normalized_distribution.keys()) + per_key_aliases = [f"__cnt_{i}" for i in range(len(expected_keys_ordered))] + + # Explicit cast pins the literal to the column type so byte/short/date keys don't rely + # on implicit Spark upcasting — a future type addition with different equality semantics + # would otherwise silently miscount into the residual bucket instead of erroring. + key_lits = [F.lit(k).cast(column_type) for k in expected_keys_ordered] + + # Single-pass conditional aggregation bounded by length of distribution: + # one COUNT(when col == key) per expected key + one COUNT(col) for the total. + # F.count(group_expr) natively skips nulls, matching the docstring's contract that NULLs are + # excluded from the actual distribution. F.when(col == key, 1) also returns NULL when col + # is NULL, so null rows contribute to neither the per-key nor the total counts. + # This avoids groupBy(col).collect() which would pull an unbounded number of distinct + # values to the driver on high-cardinality columns. + aggregations = [ + F.count(F.when(group_expr == key_lit, F.lit(1))).alias(alias) + for key_lit, alias in zip(key_lits, per_key_aliases) + ] + aggregations.append(F.count(group_expr).alias("__total")) + + # limit(1) mirrors _is_aggr_compare — the ungrouped agg already returns one row, but the + # limit is informational for the planner. Broadcasting via crossJoin keeps the whole + # pipeline lazy so apply_checks doesn't force an eager job. + agg_df = filtered.agg(*aggregations).limit(1) + + total_col = F.col("__total") + deviation_terms = _is_in_distribution_get_deviation_terms( + normalized_distribution, + per_key_aliases, + total_col, + ) + tvd_col = F.lit(0.5) * reduce(py_operator.add, deviation_terms) + + # math.isclose (rel_tol=1e-09, abs_tol=0.0) as a Column expression: neutralises IEEE-754 + # noise at the boundary so a "supposed to pass" case with tvd mathematically equal to + # distance can't trip when float rounding nudges the computed tvd just above distance + # (e.g. abs(0.15-0.2) == 0.05000000000000002). The docstring's "less than or equal to the + # given distance" already promises the boundary passes, so equality-within-precision must + # not fire the check. + distance_lit = F.lit(distance) + is_close_col = F.abs(tvd_col - distance_lit) <= F.lit(1e-9) * F.greatest(F.abs(tvd_col), F.abs(distance_lit)) + tvd_violation_col = (tvd_col > distance_lit) & ~is_close_col + + tvd_message_col = F.concat( + F.lit(f"Column '{col_expr_str}' actual distribution deviates from expected by TVD="), + F.format_string("%.6f", tvd_col), + F.lit(f", which exceeds the allowed distance={distance}."), + ) + + stats_df = _is_in_distribution_get_stats_df( + agg_df, + impute, + total_col, + tvd_violation_col, + tvd_message_col, + expected_keys_ordered, + per_key_aliases, + col_expr_str, + condition_col, + message_col, + ) + + return df.crossJoin(stats_df) + + condition = make_condition( + condition=F.col(condition_col), + message=F.col(message_col), + alias=f"{col_str_norm}_is_not_in_distribution", + ) + + return condition, apply + + +def _is_in_distribution_get_data_type(df: DataFrame, column: Column, column_expression: str) -> DataType: + column_type = df.select(column).schema[0].dataType + if not isinstance(column_type, _IS_IN_DISTRIBUTION_SUPPORTED_SPARK_TYPES): + raise InvalidParameterError( + f"Column '{column_expression}' has unsupported type '{column_type.simpleString()}' for " + f"'is_in_distribution'; expected one of: boolean, string, char, byte, short, integer, " + f"long, date." + ) + return column_type + + +def _is_in_distribution_get_deviation_terms( + distribution: dict[Any, float], + per_key_aliases: list[str], + total_column: Column, +) -> list[Column]: + expected_values = list(distribution.values()) + # fsum keeps residual precision stable across many keys. + expected_residual = 1.0 - math.fsum(expected_values) + + listed_count_column = reduce(py_operator.add, (F.col(alias) for alias in per_key_aliases)) + residual_count_column = total_column - listed_count_column + actual_residual_column = residual_count_column / total_column + + deviation_terms = [ + F.abs(F.lit(expected_value) - F.col(alias) / total_column) + for expected_value, alias in zip(expected_values, per_key_aliases) + ] + deviation_terms.append(F.abs(F.lit(expected_residual) - actual_residual_column)) + return deviation_terms + + +def _is_in_distribution_get_stats_df( + aggregate_df: DataFrame, + impute: bool, + total_column: Column, + tvd_violation_column: Column, + tvd_message_column: Column, + expected_keys: list[Any], + per_key_aliases: list[str], + column_expression: str, + condition_column_name: str, + message_column_name: str, +) -> DataFrame: + if impute: + is_violation_column = F.when(total_column == 0, F.lit(False)).otherwise(tvd_violation_column) + message_expression = F.when(total_column == 0, F.lit("")).otherwise(tvd_message_column) + else: + # Missing keys are those whose per-key count is zero; sort for stable output. + missing_key_names = [ + F.when(F.col(alias) == 0, F.lit(repr(key))) for key, alias in zip(expected_keys, per_key_aliases) + ] + missing_array = F.array_sort(F.filter(F.array(*missing_key_names), lambda value: value.isNotNull())) + has_missing_column = F.size(missing_array) > 0 + missing_message_column = F.concat( + F.lit(f"Column '{column_expression}' distribution is missing expected keys ["), + F.array_join(missing_array, ", "), + F.lit("] and impute=False."), + ) + is_violation_column = ( + F.when(total_column == 0, F.lit(False)) + .when(has_missing_column, F.lit(True)) + .otherwise(tvd_violation_column) + ) + message_expression = ( + F.when(total_column == 0, F.lit("")) + .when(has_missing_column, missing_message_column) + .otherwise(tvd_message_column) + ) + + return aggregate_df.select( + is_violation_column.alias(condition_column_name), + message_expression.alias(message_column_name), + ) + + +def _is_in_distribution_validate_distribution( + distribution: dict[bool, float] | dict[str, float] | dict[int, float] | dict[datetime.date, float], + case_sensitive: bool, +) -> None: + """Validate the *distribution* argument of *is_in_distribution*.""" + if distribution is None: + raise MissingParameterError("'distribution' is not provided.") + if not isinstance(distribution, dict): + raise InvalidParameterError(f"'distribution' must be a dict, got {type(distribution).__name__} instead.") + if not distribution: + raise InvalidParameterError("'distribution' must not be empty.") + + _is_in_distribution_validate_items(distribution) + _is_in_distribution_validate_values(distribution) + _is_in_distribution_validate_key_types(distribution) + _is_in_distribution_validate_key_collisions(case_sensitive, distribution) + + +def _is_in_distribution_validate_values( + distribution: dict[bool, float] | dict[str, float] | dict[int, float] | dict[datetime.date, float], +) -> None: + total = math.fsum(distribution.values()) + if total > 1: + raise InvalidParameterError(f"'distribution' values sum ({total}) is greater than 1.") + + +def _is_in_distribution_validate_key_collisions( + case_sensitive: bool, + distribution: dict[bool, float] | dict[str, float] | dict[int, float] | dict[datetime.date, float], +) -> None: + if not case_sensitive: + str_keys = [k for k in distribution.keys() if isinstance(k, str)] + normalized_to_original: dict[str, list[str]] = {} + for key in str_keys: + normalized_to_original.setdefault(key.lower(), []).append(key) + colliding = {norm: originals for norm, originals in normalized_to_original.items() if len(originals) > 1} + if colliding: + raise InvalidParameterError( + f"'distribution' keys collide after case-insensitive normalisation: {colliding}." + ) + + +def _is_in_distribution_validate_key_types( + distribution: dict[bool, float] | dict[str, float] | dict[int, float] | dict[datetime.date, float], +) -> None: + key_types = {type(k) for k in distribution.keys()} + if len(key_types) > 1: + type_names = sorted(t.__name__ for t in key_types) + raise InvalidParameterError( + f"'distribution' keys must be homogeneous (all of the same type), got mixed types: {type_names}." + ) + + key_type = next(iter(key_types)) + if not issubclass(key_type, _IS_IN_DISTRIBUTION_SUPPORTED_KEY_TYPES): + raise InvalidParameterError( + f"'distribution' keys of type {key_type.__name__!r} are not supported; " + f"expected one of: bool, str, int, datetime.date." + ) + + +def _is_in_distribution_validate_items( + distribution: dict[bool, float] | dict[str, float] | dict[int, float] | dict[datetime.date, float], +) -> None: + for key, value in distribution.items(): + if key is None: + raise InvalidParameterError("'distribution' must not contain None as a key.") + if value is None: + raise InvalidParameterError(f"'distribution' must not contain None as a value (key={key!r}).") + if not _is_in_distribution_value_is_numeric(value): + raise InvalidParameterError( + f"'distribution' value for key {key!r} must be a number, got {type(value).__name__} instead." + ) + if not math.isfinite(value): + raise InvalidParameterError( + f"'distribution' value {value} for key {key!r} must be finite (no inf, -inf, or nan)." + ) + if value < 0: + raise InvalidParameterError(f"'distribution' value {value} for key {key!r} must be non-negative.") + + +def _is_in_distribution_validate_distance(distance: float) -> None: + """Validate the *distance* argument of *is_in_distribution*.""" + if distance is None: + raise MissingParameterError("'distance' is not provided.") + if not _is_in_distribution_value_is_numeric(distance): + raise InvalidParameterError(f"'distance' must be a number, got {type(distance).__name__} instead.") + if not 0 <= distance <= 1: + raise InvalidParameterError(f"'distance' must be between 0 and 1 (inclusive), got {distance}.") + + +def _is_in_distribution_value_is_numeric(value: object) -> bool: + """True for numeric arguments accepted by *is_in_distribution*. + + Rejects bool explicitly so that ``True``/``False`` (which are valid distribution *keys* + and would silently coerce to 1/0 elsewhere) never slip through as probabilities or as a + TVD threshold. + """ + return isinstance(value, (int, float)) and not isinstance(value, bool) + + @register_rule("row") def is_not_in_list(column: str | Column, forbidden: list, case_sensitive: bool = True) -> Column: """Checks whether the values in the input column are NOT present in the list of forbidden values diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 3626f6d5b..4f163d7e2 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -5283,7 +5283,7 @@ def test_apply_checks_with_sql_expression(ws, spark): checks = [ { "criticality": "error", - "check": {"function": "sql_expression", "arguments": {"expression": "col1 not like \"val%\""}}, + "check": {"function": "sql_expression", "arguments": {"expression": 'col1 not like "val%"'}}, }, { "criticality": "error", @@ -5327,7 +5327,7 @@ def test_apply_checks_with_sql_expression(ws, spark): [ { "name": "not_col1_not_like_val", - "message": "Value is not matching expression: col1 not like \"val%\"", + "message": 'Value is not matching expression: col1 not like "val%"', "columns": None, "filter": None, "function": "sql_expression", @@ -5337,7 +5337,7 @@ def test_apply_checks_with_sql_expression(ws, spark): }, { "name": "not_col2_not_like_val", - "message": "Value is not matching expression: col2 \nnot \n like \"val%\"", + "message": 'Value is not matching expression: col2 \nnot \n like "val%"', "columns": None, "filter": None, "function": "sql_expression", @@ -5383,7 +5383,7 @@ def test_apply_checks_with_sql_expression_using_classes(ws, spark): DQRowRule( criticality="error", check_func=check_funcs.sql_expression, - check_func_kwargs={"expression": "col1 not like \"val%\""}, + check_func_kwargs={"expression": 'col1 not like "val%"'}, ), DQRowRule( criticality="error", @@ -5418,7 +5418,7 @@ def test_apply_checks_with_sql_expression_using_classes(ws, spark): [ { "name": "not_col1_not_like_val", - "message": "Value is not matching expression: col1 not like \"val%\"", + "message": 'Value is not matching expression: col1 not like "val%"', "columns": None, "filter": None, "function": "sql_expression", @@ -7249,6 +7249,25 @@ def test_apply_checks_all_checks_using_classes(ws, spark): column="*", check_func_kwargs={"aggr_type": "count", "ref_df_name": "ref_df_key"}, ), + # is_in_distribution check — TVD-based categorical distribution check on col10 + # (constant "2" here, so {2: 1.0} matches exactly with TVD=0). + DQDatasetRule( + criticality="error", + check_func=check_funcs.is_in_distribution, + column="col10", + check_func_kwargs={"distribution": {2: 1.0}, "distance": 1.0}, + ), + # Second is_in_distribution rule exercises non-trivial TVD math with a residual bucket + # (sum<1) and a listed key not present in the data (imputed as 0). distance=1.0 keeps + # it safely passing on both fixtures, but a regression that inverts the comparison + # (tvd <= distance) would flag every row and fail the test — the trivial {2: 1.0} rule + # above cannot catch that. + DQDatasetRule( + criticality="error", + check_func=check_funcs.is_in_distribution, + column="col10", + check_func_kwargs={"distribution": {2: 0.5, 3: 0.3}, "distance": 1.0}, + ), # is_valid_json check DQRowRule( criticality="error", @@ -9397,7 +9416,7 @@ def test_compare_datasets_check(ws, spark, set_utc_timezone): "score": {"df": "26.7", "ref": "26.9"}, }, }, - separators=(',', ':'), + separators=(",", ":"), ), "columns": pk_columns, "filter": "id1 != 2", @@ -9499,7 +9518,7 @@ def test_compare_datasets_check_missing_records(ws, spark, set_utc_timezone): "dt": {"df": "2017-01-01", "ref": "2018-01-01"}, }, }, - separators=(',', ':'), + separators=(",", ":"), ), "columns": pk_columns, "filter": None, @@ -9535,7 +9554,7 @@ def test_compare_datasets_check_missing_records(ws, spark, set_utc_timezone): "active": {"df": "true"}, }, }, - separators=(',', ':'), + separators=(",", ":"), ), "columns": pk_columns, "filter": None, @@ -9572,7 +9591,7 @@ def test_compare_datasets_check_missing_records(ws, spark, set_utc_timezone): "active": {"ref": "true"}, }, }, - separators=(',', ':'), + separators=(",", ":"), ), "columns": pk_columns, "filter": None, @@ -9720,7 +9739,7 @@ def test_compare_datasets_check_missing_records_with_partial_filter( "name": {"ref": "Marcin"}, }, }, - separators=(',', ':'), + separators=(",", ":"), ), "columns": pk_columns, "filter": filter_str, @@ -9746,7 +9765,7 @@ def test_compare_datasets_check_missing_records_with_partial_filter( "name": {"df": "Marcin"}, }, }, - separators=(',', ':'), + separators=(",", ":"), ), "columns": pk_columns, "filter": filter_str, @@ -10734,3 +10753,150 @@ def test_apply_checks_by_metadata_skip_checks_with_missing_columns(ws, spark): SCHEMA + complex_cols_schema + REPORTING_COLUMNS, ) assert_df_equality(checked, expected, ignore_nullable=True) + + +# --------------------------------------------------------------------------- +# is_in_distribution — apply_checks_by_metadata (YAML-equivalent) scenarios +# --------------------------------------------------------------------------- + + +def test_apply_checks_by_metadata_is_in_distribution_matches(ws, spark): + """YAML/metadata path: actual distribution matches expected within distance → no violations.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + schema = "id: int, value: string" + test_df = spark.createDataFrame( + [[i + 1, v] for i, v in enumerate(["A"] * 7 + ["B"] * 2 + ["C"])], + schema, + ) + checks = [ + { + "criticality": "error", + "check": { + "function": "is_in_distribution", + "arguments": { + "column": "value", + "distribution": {"A": 0.75, "B": 0.15, "C": 0.10}, + "distance": 0.05, + }, + }, + }, + ] + + checked = dq_engine.apply_checks_by_metadata(test_df, checks) + + expected = spark.createDataFrame( + [[i + 1, v, None, None] for i, v in enumerate(["A"] * 7 + ["B"] * 2 + ["C"])], + schema + REPORTING_COLUMNS, + ) + assert_df_equality(checked.sort("id"), expected, ignore_nullable=True) + + +def test_apply_checks_by_metadata_is_in_distribution_fails_when_distance_too_small(ws, spark): + """YAML/metadata path: distance=0 exposes a TVD=0.05 gap → every row is flagged.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + schema = "id: int, value: string" + test_df = spark.createDataFrame( + [[i + 1, v] for i, v in enumerate(["A"] * 7 + ["B"] * 2 + ["C"])], + schema, + ) + checks_yaml = yaml.safe_load( + """ + - criticality: error + check: + function: is_in_distribution + arguments: + column: value + distribution: + A: 0.75 + B: 0.15 + C: 0.10 + distance: 0.0 + """ + ) + + checked = dq_engine.apply_checks_by_metadata(test_df, checks_yaml) + + violation = build_quality_violation( + name="value_is_not_in_distribution", + message=( + "Column 'value' actual distribution deviates from expected by TVD=0.050000, " + "which exceeds the allowed distance=0.0." + ), + columns=["value"], + function="is_in_distribution", + ) + expected = spark.createDataFrame( + [[i + 1, v, [violation], None] for i, v in enumerate(["A"] * 7 + ["B"] * 2 + ["C"])], + schema + REPORTING_COLUMNS, + ) + assert_df_equality(checked.sort("id"), expected, ignore_nullable=True) + + +def test_apply_checks_by_metadata_is_in_distribution_case_insensitive_normalisation(ws, spark): + """YAML/metadata path with case_sensitive=false lowercases both the column and expected keys.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + schema = "id: int, value: string" + test_df = spark.createDataFrame( + [[i + 1, v] for i, v in enumerate(["a", "A", "A", "B", "b"])], + schema, + ) + checks = [ + { + "criticality": "error", + "check": { + "function": "is_in_distribution", + "arguments": { + "column": "value", + "distribution": {"A": 0.6, "B": 0.4}, + "distance": 0.001, + "case_sensitive": False, + }, + }, + }, + ] + + checked = dq_engine.apply_checks_by_metadata(test_df, checks) + + expected = spark.createDataFrame( + [[i + 1, v, None, None] for i, v in enumerate(["a", "A", "A", "B", "b"])], + schema + REPORTING_COLUMNS, + ) + assert_df_equality(checked.sort("id"), expected, ignore_nullable=True) + + +def test_apply_checks_by_metadata_is_in_distribution_impute_false_missing_keys(ws, spark): + """YAML/metadata path: impute=false flags every row with a message enumerating missing keys.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + schema = "id: int, value: string" + test_df = spark.createDataFrame( + [[i + 1, v] for i, v in enumerate(["A", "A", "A", "B"])], + schema, + ) + checks = [ + { + "criticality": "error", + "check": { + "function": "is_in_distribution", + "arguments": { + "column": "value", + "distribution": {"A": 0.5, "B": 0.4, "C": 0.1}, + "distance": 0.5, + "impute": False, + }, + }, + }, + ] + + checked = dq_engine.apply_checks_by_metadata(test_df, checks) + + violation = build_quality_violation( + name="value_is_not_in_distribution", + message="""Column 'value' distribution is missing expected keys ['C'] and impute=False.""", + columns=["value"], + function="is_in_distribution", + ) + expected = spark.createDataFrame( + [[i + 1, v, [violation], None] for i, v in enumerate(["A", "A", "A", "B"])], + schema + REPORTING_COLUMNS, + ) + assert_df_equality(checked.sort("id"), expected, ignore_nullable=True) diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index a8ce29920..86b0520ce 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -22,9 +22,11 @@ is_data_fresh_per_time_window, has_no_gaps_per_time_window, has_valid_schema, + is_in_distribution, sql_query, aggr_matches_dataset, ) +from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.utils import get_column_name_or_alias from databricks.labs.dqx.errors import InvalidParameterError, MissingParameterError, UnsafeSqlQueryError @@ -4441,3 +4443,294 @@ def test_has_valid_schema_with_exclude_columns_as_expression(spark: SparkSession "a string, b int, c double, d string, has_invalid_schema string", ) assertDataFrameEqual(actual_condition_df, expected_condition_df) + + +# --------------------------------------------------------------------------- +# is_in_distribution — Total Variation Distance dataset check +# --------------------------------------------------------------------------- + + +def _apply_and_collect_violations(condition: Column, apply_method: Callable, df: DataFrame) -> list[Any]: + """Run the dataset-level check and return the violation message per row.""" + return [row["violation"] for row in apply_method(df).select(condition.alias("violation")).collect()] + + +def test_is_in_distribution_matches_within_distance(spark: SparkSession): + """Case 1: A:7 B:2 C:1 (10 rows) vs expected {A:0.75, B:0.15, C:0.10}, distance 0.05 → pass. + + Actual proportions A=0.7, B=0.2, C=0.1 differ from expected by TVD=0.05 which is not greater + than the allowed 0.05, so no rows are flagged.""" + df = spark.createDataFrame([(v,) for v in ["A"] * 7 + ["B"] * 2 + ["C"]], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.75, "B": 0.15, "C": 0.10}, distance=0.05) + assert all(v is None for v in _apply_and_collect_violations(condition, apply_method, df)) + + +def test_is_in_distribution_expected_omits_key_falls_into_residual_bucket(spark: SparkSession): + """Case 2: same 10 rows, expected {A:0.75, B:0.15} — implicit residual mass 0.10 absorbs C.""" + values = ["A"] * 7 + ["B"] * 2 + ["C"] + df = spark.createDataFrame([(v,) for v in values], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.75, "B": 0.15}, distance=0.05) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_fails_when_distance_too_small(spark: SparkSession): + """Case 3: same 10 rows and exact expected, but distance=0 → TVD=0.05 flags every row.""" + df = spark.createDataFrame([(v,) for v in ["A"] * 7 + ["B"] * 2 + ["C"]], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.75, "B": 0.15, "C": 0.10}, distance=0.0) + violations = _apply_and_collect_violations(condition, apply_method, df) + expected_message = ( + "Column 'value' actual distribution deviates from expected by TVD=0.050000, " + "which exceeds the allowed distance=0.0." + ) + assert violations == [expected_message] * len(violations) + + +def test_is_in_distribution_fails_for_residual_bucket_when_distance_too_small(spark: SparkSession): + """Case 4: same 10 rows, expected {A:0.75, B:0.15} (residual 0.10 vs actual 0.10) but distance=0.001 + → violation is driven by A/B deviations (0.05 each), not the residual.""" + df = spark.createDataFrame([(v,) for v in ["A"] * 7 + ["B"] * 2 + ["C"]], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.75, "B": 0.15}, distance=0.001) + violations = _apply_and_collect_violations(condition, apply_method, df) + expected_message = ( + "Column 'value' actual distribution deviates from expected by TVD=0.050000, " + "which exceeds the allowed distance=0.001." + ) + assert violations == [expected_message] * len(violations) + + +def test_is_in_distribution_exact_match_passes_at_distance_zero(spark: SparkSession): + """Exact match: actual distribution equals expected → TVD=0, passes even at distance=0.""" + values = ["A", "A", "B", "B"] + df = spark.createDataFrame([(v,) for v in values], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.5, "B": 0.5}, distance=0.0) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_extra_value_aggregated_into_residual(spark: SparkSession): + """Case 5: A:6 B:2 C:1 D:1 (10 rows) with expected {A:0.7, B:0.2, C:0.1} (sum=1, residual=0). + Actual: A=0.6, B=0.2, C=0.1, residual=0.1 → TVD = 0.5*(0.1+0+0+0.1) = 0.1 → passes at distance=0.15.""" + values = ["A"] * 6 + ["B"] * 2 + ["C"] + ["D"] + df = spark.createDataFrame([(v,) for v in values], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.7, "B": 0.2, "C": 0.1}, distance=0.15) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_skips_null_values(spark: SparkSession): + """Case 6: null values in the target column are excluded from the actual distribution.""" + values = ["A", "A", "A", "B", None, None] + df = spark.createDataFrame([(v,) for v in values], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.75, "B": 0.25}, distance=0.001) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_case_insensitive_normalisation(spark: SparkSession): + """Case 7: case_sensitive=False lowercases both the column and expected keys before comparing.""" + values = ["a", "A", "A", "B", "b"] + df = spark.createDataFrame([(v,) for v in values], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.6, "B": 0.4}, distance=0.001, case_sensitive=False) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +@pytest.mark.parametrize( + "spark_type, arrow_type, values, distribution", + [ + ("boolean", "boolean", [True, True, True, False], {True: 0.75, False: 0.25}), + ("string", "string", ["x", "x", "x", "y"], {"x": 0.75, "y": 0.25}), + ("char(3)", "string", ["abc", "abc", "abc", "xyz"], {"abc": 0.75, "xyz": 0.25}), + ("byte", "byte", [1, 1, 1, 2], {1: 0.75, 2: 0.25}), + ("short", "short", [1, 1, 1, 2], {1: 0.75, 2: 0.25}), + ("int", "int", [1, 1, 1, 2], {1: 0.75, 2: 0.25}), + ("long", "long", [1, 1, 1, 2], {1: 0.75, 2: 0.25}), + ( + "date", + "date", + [date(2024, 1, 1)] * 3 + [date(2024, 2, 1)], + {date(2024, 1, 1): 0.75, date(2024, 2, 1): 0.25}, + ), + ], +) +def test_is_in_distribution_supported_column_types( + spark: SparkSession, + spark_type: str, + arrow_type: str, + values: list, + distribution: dict, +): + """Case 8: every supported Spark type must be accepted and evaluated correctly.""" + df = spark.createDataFrame([(v,) for v in values], f"value: {arrow_type}") + df = df.withColumn("value", F.col("value").cast(spark_type)) + condition, apply_method = is_in_distribution("value", distribution, distance=0.001) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + f"value: {arrow_type}, value_is_not_in_distribution: string", + ).withColumn("value", F.col("value").cast(spark_type)) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +@pytest.mark.parametrize( + "spark_type, values, spark_type_display", + [ + ("float", [1.0, 1.0, 2.0, 2.0], "float"), + ("double", [1.0, 1.0, 2.0, 2.0], "double"), + ("array", [[1], [2]], "array"), + ("binary", [bytes([1]), bytes([2])], "binary"), + ( + "timestamp", + [datetime(2024, 1, 1), datetime(2024, 2, 1)], + "timestamp", + ), + ("decimal(10,2)", [Decimal("1.00"), Decimal("2.00")], "decimal(10,2)"), + ], +) +def test_is_in_distribution_unsupported_column_types( + spark: SparkSession, + spark_type: str, + values: list, + spark_type_display: str, +): + """Case 9: unsupported column types must be rejected at apply time with a clear error.""" + df = spark.createDataFrame([(v,) for v in values], f"value: {spark_type}") + _, apply_method = is_in_distribution("value", {"A": 0.5, "B": 0.5}, distance=0.5) + with pytest.raises(InvalidParameterError) as exc_info: + apply_method(df) + assert str(exc_info.value) == ( + f"Column 'value' has unsupported type '{spark_type_display}' for 'is_in_distribution'; " + "expected one of: boolean, string, char, byte, short, integer, long, date." + ) + + +def test_is_in_distribution_impute_false_flags_missing_expected_keys(spark: SparkSession): + """impute=False fails the check when a key from the expected distribution is absent from the column.""" + df = spark.createDataFrame([("A",), ("A",), ("A",), ("B",)], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.5, "B": 0.4, "C": 0.1}, distance=0.5, impute=False) + violations = _apply_and_collect_violations(condition, apply_method, df) + expected_message = "Column 'value' distribution is missing expected keys ['C'] and impute=False." + assert violations == [expected_message] * len(violations) + + +def test_is_in_distribution_impute_true_treats_missing_expected_key_as_zero(spark: SparkSession): + """impute=True (default) treats a missing expected key as a zero-probability observation and + still computes a TVD (rather than short-circuiting to a violation).""" + values = ["A", "A", "A", "B"] + df = spark.createDataFrame([(v,) for v in values], "value: string") + # actual A=0.75, B=0.25, C=0 (imputed) → TVD = 0.5*(|0.5-0.75|+|0.4-0.25|+|0.1-0|) = 0.25 + condition, apply_method = is_in_distribution("value", {"A": 0.5, "B": 0.4, "C": 0.1}, distance=0.3, impute=True) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_row_filter_applied_before_distribution(spark: SparkSession): + """row_filter narrows the dataset before the actual distribution is computed.""" + rows = [("A", 1), ("A", 1), ("A", 1), ("B", 1), ("X", 2), ("Y", 2)] + df = spark.createDataFrame(rows, "value: string, keep: int") + condition, apply_method = is_in_distribution("value", {"A": 0.75, "B": 0.25}, distance=0.001, row_filter="keep = 1") + actual = apply_method(df).select("value", "keep", condition) + expected = spark.createDataFrame( + [(v, k, None) for v, k in rows], + "value: string, keep: int, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_empty_dataframe(spark: SparkSession): + """An empty dataset (no rows at all) produces no violations.""" + df = spark.createDataFrame([], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.5, "B": 0.5}, distance=0.01) + assert _apply_and_collect_violations(condition, apply_method, df) == [] + + +def test_is_in_distribution_all_nulls_dataframe(spark: SparkSession): + """A dataset where every row is NULL yields no actual distribution and therefore no violation.""" + df = spark.createDataFrame([(None,), (None,), (None,)], "value: string") + condition, apply_method = is_in_distribution("value", {"A": 0.5, "B": 0.5}, distance=0.01) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(None, None), (None, None), (None, None)], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_accepts_column_expression(spark: SparkSession): + """The check accepts a Spark Column expression in addition to a plain column name.""" + values = ["A"] * 7 + ["B"] * 2 + ["C"] + df = spark.createDataFrame([(v,) for v in values], "value: string") + condition, apply_method = is_in_distribution(F.col("value"), {"A": 0.75, "B": 0.15, "C": 0.10}, distance=0.05) + actual = apply_method(df).select("value", condition) + expected = spark.createDataFrame( + [(v, None) for v in values], + "value: string, value_is_not_in_distribution: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_accepts_wrapped_column_expression(spark: SparkSession): + """A wrapped Column expression like F.upper(F.col('v')) is supported — the check must resolve + the type off the expression itself, not by looking up the rendered expression string in the + input schema (which would KeyError since 'upper(v)' is not a field name). + """ + values = ["a", "A", "a", "b"] + df = spark.createDataFrame([(v,) for v in values], "v: string") + # After UPPER, 'a' → 'A' (3 rows) and 'b' → 'B' (1 row) → matches {A: 0.75, B: 0.25} exactly. + condition, apply_method = is_in_distribution(F.upper(F.col("v")), {"A": 0.75, "B": 0.25}, distance=0.0) + actual = apply_method(df).select("v", condition.alias("violation")) + expected = spark.createDataFrame( + [(v, None) for v in values], + "v: string, violation: string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_is_in_distribution_flags_via_metadata_api(ws, spark): + """End-to-end dict/metadata → apply_checks_by_metadata path must catch a real distribution + mismatch, not just verify the check runs (the shared 'apply every check' YAML fixture uses + distance=1.0 which cannot fail). + """ + df = spark.createDataFrame([("A",), ("A",), ("A",), ("A",), ("B",)], "value: string") + # actual A=0.8, B=0.2 vs expected A=0.5, B=0.5 → TVD = 0.5*(0.3+0.3) = 0.3 > distance=0.1. + checks = [ + { + "criticality": "error", + "check": { + "function": "is_in_distribution", + "arguments": { + "column": "value", + "distribution": {"A": 0.5, "B": 0.5}, + "distance": 0.1, + }, + }, + } + ] + checked = DQEngine(ws).apply_checks_by_metadata(df, checks) + errors = checked.select(F.col("_errors")).collect() + assert all(row["_errors"] is not None for row in errors) diff --git a/tests/perf/conftest.py b/tests/perf/conftest.py index 2ff07dde1..334a03d01 100644 --- a/tests/perf/conftest.py +++ b/tests/perf/conftest.py @@ -458,3 +458,26 @@ def generated_language_code_df(spark): gen = gen.withColumnSpec(col, values=values) return gen.build() + + +# Number of distinct categorical values used by generated_distribution_df. +# is_in_distribution is designed for enumerated data, so 20 values approximates a +# realistic "bounded enum" workload while keeping the driver-side aggregation bounded. +DISTRIBUTION_VALUE_COUNT = 20 + + +@pytest.fixture +def generated_distribution_df(spark): + """DEFAULT_ROWS rows with a single categorical column drawn from 20 distinct int values. + + dbldatagen produces a uniform-like distribution over the supplied ``values`` list, so the + actual proportions come out close to 1/20 per key — matched by the paired benchmark's + expected distribution. + """ + schema = _parse_datatype_string("col_categorical: int") + _, gen = make_data_gen(spark, n_rows=DEFAULT_ROWS, n_columns=1, partitions=DEFAULT_PARTITIONS) + gen = gen.withSchema(schema).withColumnSpec( + "col_categorical", + values=list(range(1, DISTRIBUTION_VALUE_COUNT + 1)), + ) + return gen.build() diff --git a/tests/perf/test_apply_checks.py b/tests/perf/test_apply_checks.py index 8239e9876..f4cbca867 100644 --- a/tests/perf/test_apply_checks.py +++ b/tests/perf/test_apply_checks.py @@ -6,7 +6,7 @@ import pyspark.sql.functions as F from databricks.labs.dqx import check_funcs from databricks.labs.dqx.geo import check_funcs as geo_check_funcs -from tests.perf.conftest import DEFAULT_ROWS +from tests.perf.conftest import DEFAULT_ROWS, DISTRIBUTION_VALUE_COUNT RUN_TIME = datetime(2025, 1, 1, 0, 0, 0, 0, tzinfo=timezone.utc) RUN_ID = "2f9120cf-e9f2-446a-8278-12d508b00639" @@ -2417,3 +2417,25 @@ def test_benchmark_is_valid_language_code(benchmark, ws, generated_language_code checked = dq_engine.apply_checks(generated_language_code_df, checks) actual_count = benchmark(lambda: checked.count()) assert actual_count == EXPECTED_ROWS + + +def test_benchmark_is_in_distribution(benchmark, ws, generated_distribution_df): + """Benchmark is_in_distribution against a 20-value categorical column. + + Expected distribution is uniform over the same 20 int values dbldatagen draws from, and + distance=1.0 (the TVD upper bound) keeps the check passing regardless of small sampling + skew introduced by the generator, so the benchmark measures the full aggregation path + without failing on synthetic noise.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + uniform_distribution = {i: 1.0 / DISTRIBUTION_VALUE_COUNT for i in range(1, DISTRIBUTION_VALUE_COUNT + 1)} + checks = [ + DQDatasetRule( + criticality="warn", + check_func=check_funcs.is_in_distribution, + column="col_categorical", + check_func_kwargs={"distribution": uniform_distribution, "distance": 1.0}, + ) + ] + checked = dq_engine.apply_checks(generated_distribution_df, checks) + actual_count = benchmark(lambda: checked.count()) + assert actual_count == EXPECTED_ROWS diff --git a/tests/resources/all_dataset_checks.yaml b/tests/resources/all_dataset_checks.yaml index 16bb81e5a..e783d4072 100644 --- a/tests/resources/all_dataset_checks.yaml +++ b/tests/resources/all_dataset_checks.yaml @@ -264,4 +264,33 @@ time_column: col5 aggr_type: avg lookback_num_intervals: 2 - warmup_num_intervals: 2 \ No newline at end of file + warmup_num_intervals: 2 + +# is_in_distribution check — validates that the actual value distribution of a categorical +# column stays within a Total Variation Distance of an expected distribution. +# Uses col10 (int) which is constant (all 2s) in the integration fixture — matches {2: 1.0} +# exactly (TVD=0). distance=1.0 keeps the check safely passing on the perf fixture too, +# where col10 carries arbitrary integer values. +- criticality: error + check: + function: is_in_distribution + arguments: + column: col10 + distribution: + 2: 1.0 + distance: 1.0 +# Second entry exercises non-trivial TVD math with a residual bucket (sum<1) and a listed key +# not present in the actual data (impute defaults to True → counted as 0). For col10 all 2s in +# the integration fixture: actual 2=1.0, residual=0, expected 2=0.5, 3=0.3, residual=0.2 → +# TVD = 0.5*(0.5 + 0.3 + 0.2) = 0.5. distance=1.0 keeps this safely passing on both fixtures, +# but any regression that inverts the comparison (tvd <= distance) would flag every row and +# fail the test — the {2: 1.0}/distance=1.0 entry above cannot catch that. +- criticality: error + check: + function: is_in_distribution + arguments: + column: col10 + distribution: + 2: 0.5 + 3: 0.3 + distance: 1.0 \ No newline at end of file diff --git a/tests/unit/test_check_func_signatures.py b/tests/unit/test_check_func_signatures.py index 13c6fbe7e..1147140ae 100644 --- a/tests/unit/test_check_func_signatures.py +++ b/tests/unit/test_check_func_signatures.py @@ -54,6 +54,7 @@ "is_ipv6_address_in_cidr": ("column", "cidr_block"), "is_data_fresh": ("column", "max_age_minutes", "base_timestamp"), "has_no_outliers": ("column", "row_filter"), + "is_in_distribution": ("column", "distribution", "distance", "case_sensitive", "impute", "row_filter"), "is_unique": ("columns", "nulls_distinct", "row_filter"), "foreign_key": ("columns", "ref_columns", "ref_df_name", "ref_table", "negate", "row_filter", "null_safe"), "sql_query": ( diff --git a/tests/unit/test_dataset_checks.py b/tests/unit/test_dataset_checks.py index 3c44cb3aa..9062c9dd8 100644 --- a/tests/unit/test_dataset_checks.py +++ b/tests/unit/test_dataset_checks.py @@ -1,8 +1,15 @@ +import math + import pytest import pyspark.sql.functions as F from databricks.labs.dqx import check_funcs -from databricks.labs.dqx.check_funcs import sql_query, is_data_fresh_per_time_window, has_no_gaps_per_time_window +from databricks.labs.dqx.check_funcs import ( + sql_query, + is_data_fresh_per_time_window, + has_no_gaps_per_time_window, + is_in_distribution, +) from databricks.labs.dqx.rule import DQDatasetRule from databricks.labs.dqx.errors import InvalidParameterError, UnsafeSqlQueryError, MissingParameterError @@ -346,6 +353,226 @@ def test_aggr_matches_dataset_invalid_tolerance_exceptions(abs_tolerance, rel_to ) +# --------------------------------------------------------------------------- +# is_in_distribution — validation rules +# --------------------------------------------------------------------------- +# Column-type validation (unsupported Spark types) requires a real DataFrame +# and belongs in integration tests. The unit tests below cover every input +# validation performed by the outer function call. + + +_VALID_DISTRIBUTION = {"A": 0.5, "B": 0.5} +_VALID_DISTANCE = 0.1 + + +def _call(**overrides): + """Invoke is_in_distribution with valid defaults and per-test overrides.""" + kwargs = {"column": "col1", "distribution": _VALID_DISTRIBUTION, "distance": _VALID_DISTANCE} + kwargs.update(overrides) + return is_in_distribution(**kwargs) + + +@pytest.mark.parametrize( + "overrides, expected_message", + [ + ({"distribution": None}, "'distribution' is not provided."), + ({"distance": None}, "'distance' is not provided."), + ], +) +def test_is_in_distribution_missing_required_params(overrides, expected_message): + with pytest.raises(MissingParameterError) as excinfo: + _call(**overrides) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "distribution, expected_message", + [ + ([("A", 0.5), ("B", 0.5)], "'distribution' must be a dict, got list instead."), + ("A=0.5,B=0.5", "'distribution' must be a dict, got str instead."), + (42, "'distribution' must be a dict, got int instead."), + ((("A", 0.5),), "'distribution' must be a dict, got tuple instead."), + ], +) +def test_is_in_distribution_distribution_not_a_dict(distribution, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution=distribution) + assert str(excinfo.value) == expected_message + + +def test_is_in_distribution_empty_distribution(): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution={}) + assert str(excinfo.value) == "'distribution' must not be empty." + + +def test_is_in_distribution_none_key(): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution={None: 0.5, "A": 0.5}) + assert str(excinfo.value) == "'distribution' must not contain None as a key." + + +def test_is_in_distribution_none_value(): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution={"A": None, "B": 0.5}) + assert str(excinfo.value) == "'distribution' must not contain None as a value (key='A')." + + +@pytest.mark.parametrize( + "bad_value, expected_message", + [ + ("0.5", "'distribution' value for key 'A' must be a number, got str instead."), + ([0.5], "'distribution' value for key 'A' must be a number, got list instead."), + ({"nested": 0.5}, "'distribution' value for key 'A' must be a number, got dict instead."), + (True, "'distribution' value for key 'A' must be a number, got bool instead."), + ], +) +def test_is_in_distribution_non_numeric_values(bad_value, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution={"A": bad_value, "B": 0.5}) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "bad_distance, expected_message", + [ + ("0.5", "'distance' must be a number, got str instead."), + ([0.5], "'distance' must be a number, got list instead."), + ({"nested": 0.5}, "'distance' must be a number, got dict instead."), + (True, "'distance' must be a number, got bool instead."), + ], +) +def test_is_in_distribution_non_numeric_distance(bad_distance, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distance=bad_distance) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "bad_value, expected_message", + [ + (math.inf, "'distribution' value inf for key 'A' must be finite (no inf, -inf, or nan)."), + (-math.inf, "'distribution' value -inf for key 'A' must be finite (no inf, -inf, or nan)."), + (math.nan, "'distribution' value nan for key 'A' must be finite (no inf, -inf, or nan)."), + ], +) +def test_is_in_distribution_non_finite_values(bad_value, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution={"A": bad_value, "B": 0.5}) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "bad_value, expected_message", + [ + (-0.1, "'distribution' value -0.1 for key 'A' must be non-negative."), + (-1.0, "'distribution' value -1.0 for key 'A' must be non-negative."), + (-0.0001, "'distribution' value -0.0001 for key 'A' must be non-negative."), + ], +) +def test_is_in_distribution_negative_values(bad_value, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution={"A": bad_value, "B": 0.5}) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "distribution, expected_message", + [ + ( + {"A": 0.7, "B": 0.5}, + "'distribution' values sum (1.2) is greater than 1.", + ), + ( + {"A": 1.0, "B": 0.5, "C": 0.5}, + "'distribution' values sum (2.0) is greater than 1.", + ), + ( + {"A": 1.01}, + "'distribution' values sum (1.01) is greater than 1.", + ), + ], +) +def test_is_in_distribution_sum_greater_than_one(distribution, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution=distribution) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "distance, expected_message", + [ + (-0.0001, "'distance' must be between 0 and 1 (inclusive), got -0.0001."), + (-0.1, "'distance' must be between 0 and 1 (inclusive), got -0.1."), + (-1.0, "'distance' must be between 0 and 1 (inclusive), got -1.0."), + (1.0001, "'distance' must be between 0 and 1 (inclusive), got 1.0001."), + (1.1, "'distance' must be between 0 and 1 (inclusive), got 1.1."), + (2.0, "'distance' must be between 0 and 1 (inclusive), got 2.0."), + ], +) +def test_is_in_distribution_distance_out_of_range(distance, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distance=distance) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "distribution, expected_message", + [ + ( + {"A": 0.5, 1: 0.5}, # str + int + "'distribution' keys must be homogeneous (all of the same type), got mixed types: ['int', 'str'].", + ), + ( + {"A": 0.5, True: 0.5}, # str + bool + "'distribution' keys must be homogeneous (all of the same type), got mixed types: ['bool', 'str'].", + ), + ( + {1: 0.5, 2.0: 0.5}, # int + float (also unsupported key type — but heterogeneity fires first) + "'distribution' keys must be homogeneous (all of the same type), got mixed types: ['float', 'int'].", + ), + ], +) +def test_is_in_distribution_heterogeneous_keys(distribution, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution=distribution) + assert str(excinfo.value) == expected_message + + +@pytest.mark.parametrize( + "distribution, expected_message", + [ + ( + {"US": 0.5, "us": 0.5}, + "'distribution' keys collide after case-insensitive normalisation: {'us': ['US', 'us']}.", + ), + ( + {"Hello": 0.5, "HELLO": 0.5}, + "'distribution' keys collide after case-insensitive normalisation: {'hello': ['Hello', 'HELLO']}.", + ), + ( + {"a": 0.3, "A": 0.3, "b": 0.4}, + "'distribution' keys collide after case-insensitive normalisation: {'a': ['a', 'A']}.", + ), + ], +) +def test_is_in_distribution_case_insensitive_key_collision(distribution, expected_message): + with pytest.raises(InvalidParameterError) as excinfo: + _call(distribution=distribution, case_sensitive=False) + assert str(excinfo.value) == expected_message + + +def test_is_in_distribution_case_sensitive_keys_do_not_collide(): + """When case_sensitive is True (default), keys differing only in case must not raise.""" + # This exercises the negative side of the collision rule — no error should be raised + # for these keys under the default case-sensitive semantics. The call may still return + # something falsy (until the implementation lands), but must not raise. + try: + _call(distribution={"US": 0.5, "us": 0.5}, case_sensitive=True) + except InvalidParameterError: + pytest.fail("case_sensitive=True must not treat 'US' and 'us' as colliding keys") + + @pytest.mark.parametrize("column", ["*", F.expr("*"), F.col("*")]) def test_resolve_aggregate_column_canonicalizes_star_forms(column): """Contract for #1435: every "*" form (the string "*", F.expr("*"), F.col("*")) must resolve to the