diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 497191b50..a7a99dce2 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1961,6 +1961,7 @@ You can also define your own custom dataset-level checks (see [Creating custom c | `compare_datasets` | Compares two DataFrames at both row and column levels, providing detailed information about differences, including new or missing rows and column-level changes. Only columns present in both the source and reference DataFrames are compared. Use with caution if `check_missing_records` is enabled, as this may increase the number of rows in the output beyond the original input DataFrame. The comparison does not support Map types (any column comparison on map type is skipped automatically). Comparing datasets is valuable for validating data during migrations, detecting drift, performing regression testing, or verifying synchronization between source and target systems. | `columns`: columns to use for row matching with the reference DataFrame (can be a list of string column names or column expressions, but only simple column expressions are allowed such as 'F.col("col1")'), if not having primary keys or wanting to match against all columns you can pass 'df.columns'; `ref_columns`: list of columns in the reference DataFrame or Table to row match against the source DataFrame (can be a list of string column names or column expressions, but only simple column expressions are allowed such as 'F.col("col1")'), if not having primary keys or wanting to match against all columns you can pass 'ref_df.columns'; note that `columns` are matched with `ref_columns` by position, so the order of the provided columns in both lists must be exactly aligned; `exclude_columns`: (optional) list of columns to exclude from the value comparison but not from row matching (can be a list of string column names or column expressions, but only simple column expressions are allowed such as 'F.col("col1")'); the `exclude_columns` field does not alter the list of columns used to determine row matches (columns), it only controls which columns are skipped during the value comparison; `ref_df_name`: (optional) name of the reference DataFrame (dictionary of DataFrames can be passed when applying checks); `ref_table`: (optional) fully qualified reference table name; either `ref_df_name` or `ref_table` must be provided but never both; the number of passed `columns` and `ref_columns` must match and keys are checks in the given order; `check_missing_records`: perform a FULL OUTER JOIN to identify records that are missing from source or reference DataFrames, default is False; use with caution as this may increase the number of rows in the output, as unmatched rows from both sides are included; `null_safe_row_matching`: (optional) treat NULLs as equal when matching rows using `columns` and `ref_columns` (default: True); `null_safe_column_value_matching`: (optional) treat NULLs as equal when comparing column values (default: True); `abs_tolerance`: (optional) numeric values are considered equal if the absolute difference is less than or equal to the tolerance (formula: `abs(a - b) <= tolerance`); `rel_tolerance`: differences in numeric values within this relative tolerance are ignored (formula: `abs(a - b) <= rel_tolerance * max(abs(a), abs(b))`) | | `is_data_fresh_per_time_window` | Freshness check that validates whether at least X records arrive within every Y-minute time window. | `column`: timestamp column (can be a string column name or a column expression); `window_minutes`: time window in minutes to check for data arrival; `min_records_per_window`: minimum number of records expected per time window; `lookback_windows`: (optional) number of time windows to look back from `curr_timestamp`, it filters records to include only those within the specified number of time windows from `curr_timestamp` (if no lookback is provided, the check is applied to the entire dataset); `curr_timestamp`: (optional) current timestamp column (if not provided, current_timestamp() function is used) | | `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_no_sequence_gaps` | Dataset check that flags gaps in a numeric sequence, i.e. expected values that are missing between values that are present (for example no invoice numbered 1002 while 1001 and 1003 are present). The violation is reported on the boundary row before each interior gap. This is the numeric counterpart of `has_no_gaps_per_time_window`. | `column`: numeric column to check (can be a string column name or a column expression); `step`: (optional) spacing of the expected sequence, i.e. the size of one bucket on the fixed grid aligned to zero (for example 1 for consecutive integers, or 10 for values expected every 10), must be a positive number (defaults to 1); `group_by`: optional list of columns or column expressions to detect gaps independently within each group, where each group is bounded by its own lowest and highest present value | | `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); | | `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. | @@ -2405,6 +2406,24 @@ Complex data types are supported as well. column: col6 window_minutes: 1440 # daily grain +# has_no_sequence_gaps check +- criticality: error + check: + function: has_no_sequence_gaps + arguments: + column: col2 + step: 1 # consecutive integers + +# has_no_sequence_gaps check per group, with values expected every 10 +- criticality: error + check: + function: has_no_sequence_gaps + arguments: + column: col2 + step: 10 + group_by: + - col1 + # has_valid_schema check (non-strict mode) - criticality: error check: @@ -3068,6 +3087,16 @@ checks = [ } ), + # check for gaps in a numeric sequence (missing values) + DQDatasetRule( + criticality="error", + check_func=check_funcs.has_no_sequence_gaps, + column="col2", + check_func_kwargs={ + "step": 1 # consecutive integers + } + ), + # has_valid_schema check (non-strict mode) DQDatasetRule( criticality="error", diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index ad8d9d39d..677ae52d0 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -3309,6 +3309,120 @@ def apply(df: DataFrame) -> DataFrame: return condition, apply +@register_rule("dataset") +def has_no_sequence_gaps( + column: str | Column, + step: float = 1, + group_by: list[str | Column] | None = None, +) -> tuple[Column, Callable]: + """Checks whether a numeric sequence column has gaps, i.e. expected values that are missing between + values that are present (for example no invoice numbered 1002 while 1001 and 1003 are present). + + This is the numeric counterpart of *has_no_gaps_per_time_window*. A missing value has no row to + attach a violation to, so the gap is reported on every row holding the last present value before + the gap. Distinct values of *column* are bucketed onto a fixed grid of *step* aligned to zero, so + with the default *step* of 1 the grid is the integers and detection is exact sequence-gap + detection. A gap is flagged whenever the next present bucket starts more than one step after the + current one, measured against that fixed grid rather than the distance between consecutive values. + + Only interior gaps are detected: the bounds of the sequence are its own lowest and highest present + values, so missing values beyond either end are not reported because there is no row to anchor them + to. When *group_by* is provided, gaps are detected independently within each group (for example per + customer or per ledger), each group is bounded by its own lowest and highest present value, and the + work partitions by the group key. When it is omitted, the whole column is treated as a single + global sequence. Null values are ignored and pass with no violation reported. + + A fractional *step* is bucketed in double precision, so values that are not exactly representable + in binary floating point (for example a *step* of 0.1) may bucket to a neighbouring grid position; + prefer an integer *step*, or scale the column, when exactness matters. + + In streaming, gaps are detected within individual micro-batches only. + + Args: + column: numeric column to check; can be a string column name or a column expression + step: spacing of the expected sequence, i.e. the size of one bucket on the fixed grid + (for example 1 for consecutive integers, or 10 for values expected every 10); must be a + positive number; defaults to 1 + group_by: optional list of column names or Column expressions to detect gaps independently + within each group; when omitted, the whole column is treated as a single global sequence + + Returns: + A tuple of: + - A Spark Column representing the gap condition. + - A closure that applies the gap detection and adds the necessary condition columns. + + Raises: + InvalidParameterError: if *step* is not a positive number, or if *group_by* is not a list. + """ + if isinstance(step, bool) or not isinstance(step, (int, float)) or step <= 0: + raise InvalidParameterError("step must be a positive number") + if group_by is not None and not isinstance(group_by, list): + raise InvalidParameterError("group_by must be a list of column names or column expressions") + + col_str_norm, _, col_expr = get_normalized_column_and_expr(column) + # Resolve group_by to plain column-name strings so a Column expression (e.g. F.col("region")) is + # used consistently for partitioning, selection, and the join below; get_column_name_or_alias would + # otherwise yield a rendered expression string that is not a real column and breaks the join. + group_by_names = get_columns_as_strings(group_by, allow_simple_expressions_only=True) if group_by else [] + + unique_str = uuid.uuid4().hex + bucket_col = f"__seq_gap_bucket_{col_str_norm}_{unique_str}" + next_bucket_col = f"__seq_gap_next_bucket_{col_str_norm}_{unique_str}" + condition_col = f"__seq_gap_condition_{col_str_norm}_{unique_str}" + + def apply(df: DataFrame) -> DataFrame: + """Bucket rows onto the fixed sequence grid and flag the boundary row before each gap.""" + input_columns = df.columns + + # Bucket each row onto the fixed grid aligned to zero. Nulls stay null and are excluded from gap + # detection below (distinct_buckets filters them out); they never match a gap bucket, so they + # pass through with no violation reported and are not dropped from the output. + numeric_col_expr = col_expr.cast("double") + df = df.withColumn(bucket_col, F.floor(numeric_col_expr / F.lit(step)) * F.lit(step)) + + # For each present bucket (within each group) find the next present bucket over the ordered + # buckets. Partitioning by the group key keeps per-group detection independent and lets the work + # scale across partitions. When group_by is omitted there is no partitionBy, so the sort runs in + # a single partition, but distinct() has already collapsed rows to the occupied-bucket count, so + # the sort scales with that count rather than with row count. + distinct_buckets = df.filter(numeric_col_expr.isNotNull()).select(*group_by_names, bucket_col).distinct() + ordered_buckets = ( + Window.partitionBy(*group_by_names).orderBy(bucket_col) if group_by_names else Window.orderBy(bucket_col) + ) + gaps = distinct_buckets.withColumn(next_bucket_col, F.lead(bucket_col).over(ordered_buckets)) + + # A gap exists when the next present bucket starts more than one step after the current one. + # lead() is null for the highest bucket in each group, which leaves that boundary unflagged and + # keeps detection to interior gaps. + gaps = gaps.withColumn( + condition_col, + F.col(next_bucket_col).isNotNull() & ((F.col(next_bucket_col) - F.col(bucket_col)) > F.lit(step)), + ) + + # Attach the per-bucket gap flag back to every row of the boundary bucket, keeping column order. + joined = _join_results_on_null_safe_columns( + df, + gaps, + [*group_by_names, bucket_col], + [next_bucket_col, condition_col], + ) + return joined.select(*input_columns, bucket_col, next_bucket_col, condition_col) + + condition = make_condition( + condition=F.col(condition_col), + message=F.concat_ws( + "", + F.lit("Gap in sequence: no data between the value at "), + F.col(bucket_col).cast("string"), + F.lit(" and the next present value at "), + F.col(next_bucket_col).cast("string"), + ), + alias=f"{col_str_norm}_has_no_sequence_gaps", + ) + + return condition, apply + + @register_for_original_columns_preselection() @register_rule("dataset") def has_valid_schema( diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 3626f6d5b..225d4aa46 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -139,6 +139,53 @@ def test_apply_checks_and_split_has_no_gaps_per_time_window(ws, spark, set_utc_t assert_check_and_split_results(checked, good, bad, expected, ["event_ts", "val"]) +def test_apply_checks_and_split_has_no_sequence_gaps(ws, spark): + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + schema = "invoice_no int, val int" + test_df = spark.createDataFrame( + [ + (1001, 1), # 1002 missing -> gap, boundary row is quarantined + (1003, 2), # consecutive with 1004 -> valid + (1004, 3), # valid + ], + schema, + ) + checks = [ + DQDatasetRule( + criticality="error", + check_func=check_funcs.has_no_sequence_gaps, + column="invoice_no", + ), + ] + + checked = dq_engine.apply_checks(test_df, checks) + good, bad = dq_engine.apply_checks_and_split(test_df, checks) + + expected_schema = schema + REPORTING_COLUMNS + expected = spark.createDataFrame( + [ + [ + 1001, + 1, + [ + build_quality_violation( + "invoice_no_has_no_sequence_gaps", + "Gap in sequence: no data between the value at 1001.0 and the next present value at 1003.0", + ["invoice_no"], + function="has_no_sequence_gaps", + ), + ], + None, + ], + [1003, 2, None, None], + [1004, 3, None, None], + ], + expected_schema, + ) + + assert_check_and_split_results(checked, good, bad, expected, ["invoice_no", "val"]) + + def test_apply_checks_passed(ws, spark): dq_engine = DQEngine(ws) test_df = spark.createDataFrame([[1, 3, 3]], SCHEMA) @@ -7242,6 +7289,13 @@ def test_apply_checks_all_checks_using_classes(ws, spark): column="col6", check_func_kwargs={"window_minutes": 1440}, ), + # has_no_sequence_gaps check + DQDatasetRule( + criticality="error", + check_func=check_funcs.has_no_sequence_gaps, + column="col2", + check_func_kwargs={"step": 1}, + ), # aggr_matches_dataset check — row count matches the reference dataset DQDatasetRule( criticality="error", diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index a8ce29920..ee72d9408 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -21,6 +21,7 @@ compare_datasets, is_data_fresh_per_time_window, has_no_gaps_per_time_window, + has_no_sequence_gaps, has_valid_schema, sql_query, aggr_matches_dataset, @@ -3920,6 +3921,319 @@ def test_has_no_gaps_per_time_window_trailing_gap_group_by(spark: SparkSession, assertDataFrameEqual(actual, expected, checkRowOrder=False) +def _sequence_gap_violation_message(bucket: str, next_bucket: str) -> str: + return f"Gap in sequence: no data between the value at {bucket} and the next present value at {next_bucket}" + + +def test_has_no_sequence_gaps(spark: SparkSession): + schema = "invoice_no int, val int" + data = [ + (1001, 1), + (1001, 2), # duplicate of the first value -> same bucket, no gap between them + (1003, 3), # 1002 is missing -> gap after 1001 + (1004, 4), # consecutive with 1003 -> no gap + (None, 5), # null passes with no violation + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no") + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("invoice_no", "val", condition) + + expected_schema = f"invoice_no int, val int, {condition_column} string" + expected = spark.createDataFrame( + [ + {"invoice_no": 1001, "val": 1, condition_column: _sequence_gap_violation_message("1001.0", "1003.0")}, + {"invoice_no": 1001, "val": 2, condition_column: _sequence_gap_violation_message("1001.0", "1003.0")}, + {"invoice_no": 1003, "val": 3, condition_column: None}, + {"invoice_no": 1004, "val": 4, condition_column: None}, + {"invoice_no": None, "val": 5, condition_column: None}, + ], + expected_schema, + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_no_gaps(spark: SparkSession): + schema = "invoice_no int, val int" + df = spark.createDataFrame([(1, 1), (2, 2), (3, 3)], schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no") + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("invoice_no", "val", condition) + + expected = spark.createDataFrame( + [ + {"invoice_no": 1, "val": 1, condition_column: None}, + {"invoice_no": 2, "val": 2, condition_column: None}, + {"invoice_no": 3, "val": 3, condition_column: None}, + ], + f"invoice_no int, val int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_multiple_gaps(spark: SparkSession): + schema = "invoice_no int, val int" + data = [ + (1, 1), # 2 missing -> gap to 3 + (3, 2), # 4..6 missing -> multi-value gap to 7 + (7, 3), # highest value -> nothing beyond it is reported + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no") + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("invoice_no", "val", condition) + + expected = spark.createDataFrame( + [ + {"invoice_no": 1, "val": 1, condition_column: _sequence_gap_violation_message("1.0", "3.0")}, + {"invoice_no": 3, "val": 2, condition_column: _sequence_gap_violation_message("3.0", "7.0")}, + {"invoice_no": 7, "val": 3, condition_column: None}, + ], + f"invoice_no int, val int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_custom_step(spark: SparkSession): + # Values are expected every 10, so 100 -> 110 is consecutive on the grid while 110 -> 130 skips + # the 120 bucket. Values inside a bucket (115) collapse onto that bucket's start. + schema = "reading int, val int" + data = [ + (100, 1), + (110, 2), + (115, 3), # same bucket as 110 + (130, 4), # 120 bucket missing -> gap after the 110 bucket + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column="reading", step=10) + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("reading", "val", condition) + + expected = spark.createDataFrame( + [ + {"reading": 100, "val": 1, condition_column: None}, + {"reading": 110, "val": 2, condition_column: _sequence_gap_violation_message("110.0", "130.0")}, + {"reading": 115, "val": 3, condition_column: _sequence_gap_violation_message("110.0", "130.0")}, + {"reading": 130, "val": 4, condition_column: None}, + ], + f"reading int, val int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_negative_values(spark: SparkSession): + # The grid is aligned to zero and extends below it, so gaps are detected across the sign boundary. + schema = "offset int, val int" + data = [ + (-3, 1), # -2 missing -> gap to -1 + (-1, 2), + (0, 3), + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column="offset") + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("offset", "val", condition) + + expected = spark.createDataFrame( + [ + {"offset": -3, "val": 1, condition_column: _sequence_gap_violation_message("-3.0", "-1.0")}, + {"offset": -1, "val": 2, condition_column: None}, + {"offset": 0, "val": 3, condition_column: None}, + ], + f"offset int, val int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_single_row(spark: SparkSession): + schema = "invoice_no int, val int" + df = spark.createDataFrame([(1001, 1)], schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no") + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("invoice_no", "val", condition) + + expected = spark.createDataFrame( + [{"invoice_no": 1001, "val": 1, condition_column: None}], + f"invoice_no int, val int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_empty_dataframe(spark: SparkSession): + schema = "invoice_no int, val int" + df = spark.createDataFrame([], schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no") + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("invoice_no", "val", condition) + + expected = spark.createDataFrame([], f"invoice_no int, val int, {condition_column} string") + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_all_null(spark: SparkSession): + schema = "invoice_no int, val int" + df = spark.createDataFrame([(None, 1), (None, 2)], schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no") + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("invoice_no", "val", condition) + + expected = spark.createDataFrame( + [ + {"invoice_no": None, "val": 1, condition_column: None}, + {"invoice_no": None, "val": 2, condition_column: None}, + ], + f"invoice_no int, val int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_group_by(spark: SparkSession): + schema = "customer string, invoice_no int, val int" + data = [ + ("A", 1, 1), # customer A: 2 missing -> gap on this boundary row + ("A", 3, 2), + ("B", 1, 3), # customer B: consecutive, no gaps + ("B", 2, 4), + ("B", 3, 5), + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no", group_by=["customer"]) + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("customer", "invoice_no", "val", condition) + + expected = spark.createDataFrame( + [ + { + "customer": "A", + "invoice_no": 1, + "val": 1, + condition_column: _sequence_gap_violation_message("1.0", "3.0"), + }, + {"customer": "A", "invoice_no": 3, "val": 2, condition_column: None}, + {"customer": "B", "invoice_no": 1, "val": 3, condition_column: None}, + {"customer": "B", "invoice_no": 2, "val": 4, condition_column: None}, + {"customer": "B", "invoice_no": 3, "val": 5, condition_column: None}, + ], + f"customer string, invoice_no int, val int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_group_by_bounded_independently(spark: SparkSession): + # Each group is bounded by its own highest present value, so a group whose sequence simply stops + # earlier than another group's is not a gap - only interior gaps are reported. + schema = "customer string, invoice_no int" + data = [ + ("A", 1), # customer A ends at 2 while customer B runs to 9: not a gap + ("A", 2), + ("B", 1), + ("B", 2), + ("B", 9), # interior gap within customer B -> flagged on the boundary row + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no", group_by=["customer"]) + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("customer", "invoice_no", condition) + + expected = spark.createDataFrame( + [ + {"customer": "A", "invoice_no": 1, condition_column: None}, + {"customer": "A", "invoice_no": 2, condition_column: None}, + {"customer": "B", "invoice_no": 1, condition_column: None}, + {"customer": "B", "invoice_no": 2, condition_column: _sequence_gap_violation_message("2.0", "9.0")}, + {"customer": "B", "invoice_no": 9, condition_column: None}, + ], + f"customer string, invoice_no int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_group_by_null_key(spark: SparkSession): + schema = "customer string, invoice_no int" + df = spark.createDataFrame([(None, 1), (None, 3)], schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no", group_by=["customer"]) + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("customer", "invoice_no", condition) + + expected = spark.createDataFrame( + [ + {"customer": None, "invoice_no": 1, condition_column: _sequence_gap_violation_message("1.0", "3.0")}, + {"customer": None, "invoice_no": 3, condition_column: None}, + ], + f"customer string, invoice_no int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_group_by_mixed_null_and_named_keys(spark: SparkSession): + # A batch mixing a NULL-key group with named groups: the null-safe join must keep each group's gaps + # isolated (NULL rows must not match a named group's buckets, and vice versa). + schema = "customer string, invoice_no int" + data = [ + (None, 1), # NULL group: 2 missing -> gap on this boundary row + (None, 3), + ("A", 10), # customer A: 11 missing -> gap on this boundary row + ("A", 12), + ("B", 1), # customer B: consecutive, no gaps + ("B", 2), + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column="invoice_no", group_by=["customer"]) + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("customer", "invoice_no", condition) + + expected = spark.createDataFrame( + [ + {"customer": None, "invoice_no": 1, condition_column: _sequence_gap_violation_message("1.0", "3.0")}, + {"customer": None, "invoice_no": 3, condition_column: None}, + {"customer": "A", "invoice_no": 10, condition_column: _sequence_gap_violation_message("10.0", "12.0")}, + {"customer": "A", "invoice_no": 12, condition_column: None}, + {"customer": "B", "invoice_no": 1, condition_column: None}, + {"customer": "B", "invoice_no": 2, condition_column: None}, + ], + f"customer string, invoice_no int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + +def test_has_no_sequence_gaps_group_by_column_expression(spark: SparkSession): + schema = "customer string, invoice_no int" + data = [ + ("A", 1), + ("A", 3), # 2 missing -> gap on the 1 boundary row + ("B", 1), + ("B", 2), + ] + df = spark.createDataFrame(data, schema) + + condition, apply_method = has_no_sequence_gaps(column=F.col("invoice_no"), group_by=[F.col("customer")]) + condition_column = get_column_name_or_alias(condition) + actual = apply_method(df).select("customer", "invoice_no", condition) + + expected = spark.createDataFrame( + [ + {"customer": "A", "invoice_no": 1, condition_column: _sequence_gap_violation_message("1.0", "3.0")}, + {"customer": "A", "invoice_no": 3, condition_column: None}, + {"customer": "B", "invoice_no": 1, condition_column: None}, + {"customer": "B", "invoice_no": 2, condition_column: None}, + ], + f"customer string, invoice_no int, {condition_column} string", + ) + assertDataFrameEqual(actual, expected, checkRowOrder=False) + + def test_has_valid_schema_invalid_schema_exceptions(): expected_schema = "INVALID_SCHEMA" with pytest.raises(InvalidParameterError, match=f"Invalid schema string '{expected_schema}'.*"): diff --git a/tests/perf/test_apply_checks.py b/tests/perf/test_apply_checks.py index 8239e9876..ca3450a07 100644 --- a/tests/perf/test_apply_checks.py +++ b/tests/perf/test_apply_checks.py @@ -1415,6 +1415,21 @@ def test_benchmark_has_no_gaps_per_time_window(benchmark, ws, generated_df): assert actual_count == EXPECTED_ROWS +def test_benchmark_has_no_sequence_gaps(benchmark, ws, generated_df): + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + checks = [ + DQDatasetRule( + criticality="error", + check_func=check_funcs.has_no_sequence_gaps, + column="col2", + check_func_kwargs={"step": 1}, + ), + ] + checked = dq_engine.apply_checks(generated_df, checks) + actual_count = benchmark(lambda: checked.count()) + assert actual_count == EXPECTED_ROWS + + @pytest.mark.parametrize( "generated_timestamp_df", [{"n_rows": DEFAULT_ROWS, "n_columns": 5}], diff --git a/tests/resources/all_dataset_checks.yaml b/tests/resources/all_dataset_checks.yaml index 16bb81e5a..32d1f1b52 100644 --- a/tests/resources/all_dataset_checks.yaml +++ b/tests/resources/all_dataset_checks.yaml @@ -238,6 +238,14 @@ column: col6 window_minutes: 1440 +# has_no_sequence_gaps check +- criticality: error + check: + function: has_no_sequence_gaps + arguments: + column: col2 + step: 1 + # has_valid_schema check - criticality: error check: diff --git a/tests/unit/test_check_func_signatures.py b/tests/unit/test_check_func_signatures.py index 13c6fbe7e..e2b8ec1da 100644 --- a/tests/unit/test_check_func_signatures.py +++ b/tests/unit/test_check_func_signatures.py @@ -136,6 +136,7 @@ "curr_timestamp", ), "has_no_gaps_per_time_window": ("column", "window_minutes", "group_by", "trailing_gap", "curr_timestamp"), + "has_no_sequence_gaps": ("column", "step", "group_by"), "has_valid_schema": ("expected_schema", "ref_df_name", "ref_table", "columns", "strict", "exclude_columns"), "is_valid_json": ("column",), "has_json_keys": ("column", "keys", "require_all"), diff --git a/tests/unit/test_dataset_checks.py b/tests/unit/test_dataset_checks.py index 3c44cb3aa..4a795a226 100644 --- a/tests/unit/test_dataset_checks.py +++ b/tests/unit/test_dataset_checks.py @@ -2,7 +2,12 @@ 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, + has_no_sequence_gaps, +) from databricks.labs.dqx.rule import DQDatasetRule from databricks.labs.dqx.errors import InvalidParameterError, UnsafeSqlQueryError, MissingParameterError @@ -252,6 +257,17 @@ def test_has_no_gaps_per_time_window_curr_timestamp_without_trailing_gap(): ) +@pytest.mark.parametrize("step", [0, -1, -0.5, None, True, False, "1"]) +def test_has_no_sequence_gaps_exceptions(step): + with pytest.raises(InvalidParameterError, match="step must be a positive number"): + has_no_sequence_gaps(column="invoice_no", step=step) + + +def test_has_no_sequence_gaps_invalid_group_by(): + with pytest.raises(InvalidParameterError, match="group_by must be a list"): + has_no_sequence_gaps(column="invoice_no", group_by="customer_id") + + @pytest.mark.parametrize( "expected_schema, ref_df_name, ref_table", [