From 3c5aae69568e0c882f9b7b6613e24797efa71b09 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 19 Jan 2026 15:27:01 -0500 Subject: [PATCH 01/15] Update handling of metadata columns during schema validation --- docs/dqx/docs/reference/quality_checks.mdx | 31 +++++++++++++++++++++- src/databricks/labs/dqx/check_funcs.py | 12 ++++++++- src/databricks/labs/dqx/engine.py | 13 +++++++++ tests/integration/test_dataset_checks.py | 24 +++++++++++++++++ 4 files changed, 78 insertions(+), 2 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 0093efb1d..a4f4218ec 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1654,7 +1654,7 @@ You can also define your own custom dataset-level checks (see [Creating custom c | `sql_query` | Checks whether the condition column produced by a SQL query is satisfied. The check supports two modes: **Row-level validation** (when `merge_columns` is provided) - query results are joined back to the input DataFrame to mark specific rows; **Dataset-level validation** (when `merge_columns` is None or empty) - the check result applies to all rows (or filtered rows if `row_filter` is used), making it ideal for aggregate validations with custom metrics. The query must return a boolean condition column (True = fail, False = pass). For row-level checks: if merge columns aren't unique, multiple query rows can attach to a single input row, potentially causing false positives. Performance tip: for complex queries, writing a custom dataset-level rule is usually more performant than `sql_query` check. | `query`: query string, must return condition column (and merge columns if provided); `input_placeholder`: name to be used in the sql query as `{{ input_placeholder }}` to refer to the input DataFrame, optional reference DataFrames are referred by the name provided in the dictionary of reference DataFrames (e.g. `{{ ref_df_key }}`, dictionary of DataFrames can be passed when applying checks); `merge_columns`: (optional) list of columns used for merging with the input DataFrame which must exist in the input DataFrame and be present in output of the sql query; when not provided (None or empty list), the check result applies to all rows in the dataset (dataset-level validation); `condition_column`: name of the column indicating a violation (False = pass, True = fail); `msg`: (optional) message to output; `name`: (optional) name of the resulting check (it can be overwritten by `name` specified at the check level); `negate`: if the condition should be negated | | `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) | | `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_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. | `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 | +| `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. | `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; `ignore_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); | **Compare datasets check** @@ -2001,6 +2001,16 @@ Complex data types are supported as well. - id - name +# has_valid_schema check with specific ignored columns +- criticality: warn + check: + function: has_valid_schema + arguments: + expected_schema: "id INT, name STRING, age INT, contact_info STRUCT" + ignore_columns: + - last_update_date + - last_updated_by + # has_valid_schema check using reference table - criticality: error check: @@ -2456,6 +2466,25 @@ checks = [ }, ), + # has_valid_schema check with specific ignored columns, expected schema defined using StructType + DQDatasetRule( + criticality="warn", + check_func=check_funcs.has_valid_schema, + check_func_kwargs={ + "expected_schema": StructType([ + StructField("id", IntegerType(), True), + StructField("name", StringType(), True), + StructField("age", IntegerType(), True), + StructField("contact_info", StructType([ + StructField("email", StringType(), True), + StructField("phone", StringType(), True), + StructField("address", StringType(), True), + ]), True) + ]), + "columns": ["last_update_date", "last_updated_by"], + }, + ), + # has_valid_schema check using reference table DQDatasetRule( criticality="error", diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index f88dda841..387459186 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1931,6 +1931,7 @@ def has_valid_schema( ref_table: str | None = None, columns: list[str | Column] | None = None, strict: bool = False, + ignore_columns: list[str] | None = None, ) -> tuple[Column, Callable]: """ Build a schema compatibility check condition and closure for dataset-level validation. @@ -1948,6 +1949,8 @@ def has_valid_schema( strict: Whether to perform strict schema validation (default: False). - False: Validates that all expected columns exist with compatible types (allows extra columns) - True: Validates exact schema match (same columns, same order, same types) + ignore_columns: Optional list of column names in the checked DataFrame schema to + ignore for validation. Returns: A tuple of: @@ -1976,6 +1979,12 @@ def has_valid_schema( if columns: column_names = [get_column_name_or_alias(col) if not isinstance(col, str) else col for col in columns] + ignore_column_names: list[str] | None = None + if ignore_columns: + ignore_column_names = [ + get_column_name_or_alias(col) if not isinstance(col, str) else col for col in ignore_columns + ] + expected_schema = _get_schema(expected_schema or types.StructType(), column_names) unique_str = uuid.uuid4().hex # make sure any column added to the dataframe is unique @@ -2003,7 +2012,8 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> else: _expected_schema = expected_schema - actual_schema = df.select(*columns).schema if columns else df.schema + base_df = df.select(*columns) if columns else df + actual_schema = base_df.drop(*(ignore_column_names or [])).schema if strict: errors = _get_strict_schema_comparison(actual_schema, _expected_schema) diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 2a9d57513..5b1a337bf 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -14,6 +14,7 @@ from pyspark.sql.streaming import StreamingQuery from databricks.labs.dqx.base import DQEngineBase, DQEngineCoreBase +from databricks.labs.dqx import check_funcs from databricks.labs.dqx.checks_resolver import resolve_custom_check_functions_from_path from databricks.labs.dqx.checks_serializer import deserialize_checks from databricks.labs.dqx.config_serializer import ConfigSerializer @@ -345,6 +346,17 @@ def _all_are_dq_rules(checks: list[DQRule]) -> bool: """Check if all elements in the checks list are instances of DQRule.""" return all(isinstance(check, DQRule) for check in checks) + def _get_checks_with_ignored_result_columns(self, checks: list[DQRule]) -> list[DQRule]: + """Get checks with ignored result columns for schema validation checks.""" + for check in checks: + if check.check_func is not check_funcs.has_valid_schema: + continue + existing_ignored = check.check_func_kwargs.get("ignore_columns") or [] + check.check_func_kwargs["ignore_columns"] = list( + set(existing_ignored + list(self._result_column_names.values())) + ) + return checks + def _append_empty_checks(self, df: DataFrame) -> DataFrame: """Append empty checks at the end of DataFrame. @@ -385,6 +397,7 @@ def _create_results_array( empty_result = F.lit(None).cast(dq_result_schema).alias(dest_col) return df.select("*", empty_result) + checks = self._get_checks_with_ignored_result_columns(checks) check_conditions = [] current_df = df diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index cb2374120..4da07d6ae 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -2774,3 +2774,27 @@ def test_has_valid_schema_with_ref_df_name(spark: SparkSession): "a string, b string, has_invalid_schema string", ) assert_df_equality(actual_condition_df, expected_condition_df, ignore_nullable=True) + + +def test_has_valid_schema_with_ignore_columns(spark: SparkSession): + test_df = spark.createDataFrame( + [ + ["str1", 1, 100.0, "extra"], + ["str2", 2, 200.0, "data"], + ], + "a string, b int, c double, d string", + ) + + expected_schema = "a string, b int, c string" + condition, apply_method = has_valid_schema(expected_schema, ignore_columns=["d"], strict=True) + actual_apply_df = apply_method(test_df, spark, {}) + actual_condition_df = actual_apply_df.select("a", "b", "c", "d", condition) + + expected_condition_df = spark.createDataFrame( + [ + ["str1", 1, 100.0, "extra", None], + ["str2", 2, 200.0, "data", None], + ], + "a string, b int, c double, d string, has_invalid_schema string", + ) + assert_df_equality(actual_condition_df, expected_condition_df, ignore_nullable=True) From 0e8d61bd31e63da100a994d03b84ee52cfbba5ce Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 20 Jan 2026 11:00:38 -0500 Subject: [PATCH 02/15] Update implementation, docs, and tests --- docs/dqx/docs/reference/quality_checks.mdx | 4 +- src/databricks/labs/dqx/check_funcs.py | 2 + src/databricks/labs/dqx/engine.py | 13 +- tests/integration/test_apply_checks.py | 177 +++++++++++++++++++++ 4 files changed, 187 insertions(+), 9 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index a4f4218ec..b5da2d16f 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1654,7 +1654,7 @@ You can also define your own custom dataset-level checks (see [Creating custom c | `sql_query` | Checks whether the condition column produced by a SQL query is satisfied. The check supports two modes: **Row-level validation** (when `merge_columns` is provided) - query results are joined back to the input DataFrame to mark specific rows; **Dataset-level validation** (when `merge_columns` is None or empty) - the check result applies to all rows (or filtered rows if `row_filter` is used), making it ideal for aggregate validations with custom metrics. The query must return a boolean condition column (True = fail, False = pass). For row-level checks: if merge columns aren't unique, multiple query rows can attach to a single input row, potentially causing false positives. Performance tip: for complex queries, writing a custom dataset-level rule is usually more performant than `sql_query` check. | `query`: query string, must return condition column (and merge columns if provided); `input_placeholder`: name to be used in the sql query as `{{ input_placeholder }}` to refer to the input DataFrame, optional reference DataFrames are referred by the name provided in the dictionary of reference DataFrames (e.g. `{{ ref_df_key }}`, dictionary of DataFrames can be passed when applying checks); `merge_columns`: (optional) list of columns used for merging with the input DataFrame which must exist in the input DataFrame and be present in output of the sql query; when not provided (None or empty list), the check result applies to all rows in the dataset (dataset-level validation); `condition_column`: name of the column indicating a violation (False = pass, True = fail); `msg`: (optional) message to output; `name`: (optional) name of the resulting check (it can be overwritten by `name` specified at the check level); `negate`: if the condition should be negated | | `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) | | `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_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. | `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; `ignore_columns`: (optional) list of columns to ignore during validation (if not provided, all columns are considered); | +| `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 `ignore_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; `ignore_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); | **Compare datasets check** @@ -2481,7 +2481,7 @@ checks = [ StructField("address", StringType(), True), ]), True) ]), - "columns": ["last_update_date", "last_updated_by"], + "ignore_columns": ["last_update_date", "last_updated_by"], }, ), diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 387459186..3ddce9aa3 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1941,6 +1941,8 @@ def has_valid_schema( In strict mode, validates that the schema matches exactly (same columns, same order, same types) for the columns specified in columns or for all columns if columns is not specified. + All columns in the `ignore_columns` list will be ignored even if the column is present in the `columns` list. + Args: expected_schema: Expected schema as a DDL string (e.g., "id INT, name STRING") or StructType object. ref_df_name: Name of the reference DataFrame (used when passing DataFrames directly). diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 5b1a337bf..2d58abca6 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -346,15 +346,14 @@ def _all_are_dq_rules(checks: list[DQRule]) -> bool: """Check if all elements in the checks list are instances of DQRule.""" return all(isinstance(check, DQRule) for check in checks) - def _get_checks_with_ignored_result_columns(self, checks: list[DQRule]) -> list[DQRule]: - """Get checks with ignored result columns for schema validation checks.""" + @staticmethod + def _preselect_schema_validation_columns(df: DataFrame, checks: list[DQRule]) -> list[DQRule]: + """Determine columns for schema validation checks.""" for check in checks: if check.check_func is not check_funcs.has_valid_schema: continue - existing_ignored = check.check_func_kwargs.get("ignore_columns") or [] - check.check_func_kwargs["ignore_columns"] = list( - set(existing_ignored + list(self._result_column_names.values())) - ) + if not check.check_func_kwargs.get("columns"): + check.check_func_kwargs["columns"] = df.columns return checks def _append_empty_checks(self, df: DataFrame) -> DataFrame: @@ -397,7 +396,7 @@ def _create_results_array( empty_result = F.lit(None).cast(dq_result_schema).alias(dest_col) return df.select("*", empty_result) - checks = self._get_checks_with_ignored_result_columns(checks) + checks = DQEngineCore._preselect_schema_validation_columns(df, checks) check_conditions = [] current_df = df diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 054615bab..a4a83a290 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -8777,6 +8777,183 @@ def test_apply_checks_with_is_data_fresh_per_time_window(ws, spark, set_utc_time assert_df_equality(checked.sort("id"), expected, ignore_nullable=True) +def test_apply_checks_with_has_valid_schema_ignores_result_columns(ws, spark): + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + + schema = "id int, v1 int, v2 string" + test_df = spark.createDataFrame( + [ + [1, 10, "x"], + [2, None, "y"], + ], + schema, + ) + + checks = [ + DQRowRule( + name="v1_is_not_null", + criticality="error", + check_func=check_funcs.is_not_null, + column="v1", + ), + DQDatasetRule( + name="has_valid_schema", + criticality="warn", + check_func=check_funcs.has_valid_schema, + check_func_kwargs={ + "expected_schema": "id int, v1 int, v2 string", + "strict": True, + }, + ) + ] + + checked = dq_engine.apply_checks(test_df, checks) + + expected_schema = schema + REPORTING_COLUMNS + expected = spark.createDataFrame( + [ + [ + 1, + 10, + "x", + None, + None, + ], + [ + 2, + None, + "y", + [ + { + "name": "v1_is_null", + "message": "Column 'v1' value is null", + "columns": ["v1"], + "filter": None, + "function": "is_not_null", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None + ], + ], + expected_schema, + ) + + assert_df_equality(checked.sort("id"), expected.sort("id"), ignore_nullable=True) + + +def test_apply_checks_with_has_valid_schema_ignores_generated_columns(ws, spark, set_utc_timezone): + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + + schema = "id int, col1 timestamp" + test_df = spark.createDataFrame( + [ + [1, datetime(2025, 1, 1)], + [1, datetime(2025, 1, 1)], + [2, datetime(2025, 1, 2)], + [3, None], + ], + schema, + ) + + checks = [ + DQDatasetRule( + name="aggr_count_not_equal_to_limit", + criticality="error", + check_func=check_funcs.is_aggr_equal, + column="id", + check_func_kwargs={"aggr_type": "count", "limit": 1}, + ), + DQDatasetRule( + name="has_valid_schema", + criticality="warn", + check_func=check_funcs.has_valid_schema, + check_func_kwargs={"expected_schema": "id int, col1 timestamp", "strict": True}, + ), + ] + + checked = dq_engine.apply_checks_by_metadata(test_df, checks) + + expected_schema = schema + REPORTING_COLUMNS + expected = spark.createDataFrame( + [ + [ + 1, + datetime(2025, 1, 1), + [ + { + "name": "aggr_count_not_equal_to_limit", + "message": "Count value 4 in column 'id' is not equal to limit: 1", + "columns": ["id"], + "filter": None, + "function": "is_aggr_equal", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 1, + datetime(2025, 1, 1), + [ + { + "name": "aggr_count_not_equal_to_limit", + "message": "Count value 4 in column 'id' is not equal to limit: 1", + "columns": ["id"], + "filter": None, + "function": "is_aggr_equal", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 2, + datetime(2025, 1, 2), + [ + { + "name": "aggr_count_not_equal_to_limit", + "message": "Count value 4 in column 'id' is not equal to limit: 1", + "columns": ["id"], + "filter": None, + "function": "is_aggr_equal", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 3, + None, + [ + { + "name": "aggr_count_not_equal_to_limit", + "message": "Count value 4 in column 'id' is not equal to limit: 1", + "columns": ["id"], + "filter": None, + "function": "is_aggr_equal", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + ], + expected_schema, + ) + + assert_df_equality(checked.sort("id"), expected.sort("id"), ignore_nullable=True) + + def test_apply_checks_and_save_in_tables_for_patterns_missing_output_suffix(ws, spark): dq_engine = DQEngine(ws) From 49624f387fdbc236ecf7d4fba8a320041798f60d Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 20 Jan 2026 11:20:35 -0500 Subject: [PATCH 03/15] Fix formatting --- tests/integration/test_apply_checks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index a4a83a290..5970830d3 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -8804,7 +8804,7 @@ def test_apply_checks_with_has_valid_schema_ignores_result_columns(ws, spark): "expected_schema": "id int, v1 int, v2 string", "strict": True, }, - ) + ), ] checked = dq_engine.apply_checks(test_df, checks) @@ -8835,7 +8835,7 @@ def test_apply_checks_with_has_valid_schema_ignores_result_columns(ws, spark): "user_metadata": {}, }, ], - None + None, ], ], expected_schema, From 64aa561a8f5d1e3793d03f292c6dc2c0f6d9e04b Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 20 Jan 2026 14:47:32 -0500 Subject: [PATCH 04/15] Fix implementation and tests --- src/databricks/labs/dqx/check_funcs.py | 6 ++++-- src/databricks/labs/dqx/engine.py | 12 ++++++++---- tests/integration/test_apply_checks.py | 3 +-- tests/integration/test_dataset_checks.py | 2 +- 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 3ddce9aa3..81da6f7d4 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -2014,8 +2014,10 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> else: _expected_schema = expected_schema - base_df = df.select(*columns) if columns else df - actual_schema = base_df.drop(*(ignore_column_names or [])).schema + selected_column_names = column_names if column_names else df.columns + if ignore_column_names: + selected_column_names = [col for col in selected_column_names if col not in ignore_column_names] + actual_schema = df.select(*selected_column_names).schema if strict: errors = _get_strict_schema_comparison(actual_schema, _expected_schema) diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 2d58abca6..8311a0594 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -346,14 +346,18 @@ def _all_are_dq_rules(checks: list[DQRule]) -> bool: """Check if all elements in the checks list are instances of DQRule.""" return all(isinstance(check, DQRule) for check in checks) - @staticmethod - def _preselect_schema_validation_columns(df: DataFrame, checks: list[DQRule]) -> list[DQRule]: + def _preselect_schema_validation_columns(self, df: DataFrame, checks: list[DQRule]) -> list[DQRule]: """Determine columns for schema validation checks.""" for check in checks: if check.check_func is not check_funcs.has_valid_schema: continue + + # Default columns to all columns of the current DataFrame if not explicitly set if not check.check_func_kwargs.get("columns"): - check.check_func_kwargs["columns"] = df.columns + check.check_func_kwargs["columns"] = [ + col for col in df.columns if col not in set(self._result_column_names.values()) + ] + return checks def _append_empty_checks(self, df: DataFrame) -> DataFrame: @@ -396,7 +400,7 @@ def _create_results_array( empty_result = F.lit(None).cast(dq_result_schema).alias(dest_col) return df.select("*", empty_result) - checks = DQEngineCore._preselect_schema_validation_columns(df, checks) + checks = self._preselect_schema_validation_columns(df, checks) check_conditions = [] current_df = df diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 5970830d3..05a87c14d 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -8863,8 +8863,7 @@ def test_apply_checks_with_has_valid_schema_ignores_generated_columns(ws, spark, name="aggr_count_not_equal_to_limit", criticality="error", check_func=check_funcs.is_aggr_equal, - column="id", - check_func_kwargs={"aggr_type": "count", "limit": 1}, + check_func_kwargs={"aggr_type": "count", "column": "id", "limit": 1}, ), DQDatasetRule( name="has_valid_schema", diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index 4da07d6ae..80de951e1 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -2785,7 +2785,7 @@ def test_has_valid_schema_with_ignore_columns(spark: SparkSession): "a string, b int, c double, d string", ) - expected_schema = "a string, b int, c string" + expected_schema = "a string, b int, c double" condition, apply_method = has_valid_schema(expected_schema, ignore_columns=["d"], strict=True) actual_apply_df = apply_method(test_df, spark, {}) actual_condition_df = actual_apply_df.select("a", "b", "c", "d", condition) From 36b7dcaa1310276ee3f77758e1a9f58fd1493166 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 20 Jan 2026 17:20:54 -0500 Subject: [PATCH 05/15] Remove cached 'check' attribute during preselection --- src/databricks/labs/dqx/engine.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 8311a0594..bf292b5d1 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -357,6 +357,11 @@ def _preselect_schema_validation_columns(self, df: DataFrame, checks: list[DQRul check.check_func_kwargs["columns"] = [ col for col in df.columns if col not in set(self._result_column_names.values()) ] + # Clear cached rule check so the updated columns are used when evaluating. + try: + object.__delattr__(check, "check") + except AttributeError: + pass return checks From 8083221171f25536b33dd5ab3bea5e7ce4519039 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 21 Jan 2026 09:47:04 -0500 Subject: [PATCH 06/15] Fix tests --- tests/integration/test_apply_checks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 05a87c14d..87a67d485 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -8825,7 +8825,7 @@ def test_apply_checks_with_has_valid_schema_ignores_result_columns(ws, spark): "y", [ { - "name": "v1_is_null", + "name": "v1_is_not_null", "message": "Column 'v1' value is null", "columns": ["v1"], "filter": None, @@ -8873,7 +8873,7 @@ def test_apply_checks_with_has_valid_schema_ignores_generated_columns(ws, spark, ), ] - checked = dq_engine.apply_checks_by_metadata(test_df, checks) + checked = dq_engine.apply_checks(test_df, checks) expected_schema = schema + REPORTING_COLUMNS expected = spark.createDataFrame( From 02291eabad3cd4b942b0a2cad25ecc31b6e46594 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 21 Jan 2026 17:18:12 -0500 Subject: [PATCH 07/15] Refactor --- src/databricks/labs/dqx/engine.py | 49 ++++++++++++++++++------------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index bf292b5d1..33d8b1894 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -3,6 +3,7 @@ import logging from concurrent import futures from collections.abc import Callable +from dataclasses import replace from datetime import datetime from functools import cached_property from typing import Any @@ -346,24 +347,28 @@ def _all_are_dq_rules(checks: list[DQRule]) -> bool: """Check if all elements in the checks list are instances of DQRule.""" return all(isinstance(check, DQRule) for check in checks) - def _preselect_schema_validation_columns(self, df: DataFrame, checks: list[DQRule]) -> list[DQRule]: - """Determine columns for schema validation checks.""" - for check in checks: - if check.check_func is not check_funcs.has_valid_schema: - continue - - # Default columns to all columns of the current DataFrame if not explicitly set - if not check.check_func_kwargs.get("columns"): - check.check_func_kwargs["columns"] = [ - col for col in df.columns if col not in set(self._result_column_names.values()) - ] - # Clear cached rule check so the updated columns are used when evaluating. - try: - object.__delattr__(check, "check") - except AttributeError: - pass - - return checks + def _preselect_schema_validation_columns(self, df: DataFrame, check: DQRule) -> DQRule: + """ + Determines the selected columns for schema validation checks. This is required when using `has_valid_schema` to + ignore columns added during quality checking. This includes: + * DQX result columns (e.g. '_warnings' and '_errors') + * Internal columns used for aggregate checks + + Args: + df: Input DataFrame + check: DQRule to be modified + """ + + # Default columns to all columns of the current DataFrame if not explicitly set + if check.check_func_kwargs.get("columns"): + return check + + if check.check_func_args and len(check.check_func_args) >= 3: + return check + + rule_kwargs = check.check_func_kwargs.copy() + rule_kwargs["columns"] = [col for col in df.columns if col not in set(self._result_column_names.values())] + return replace(check, check_func_kwargs=rule_kwargs) def _append_empty_checks(self, df: DataFrame) -> DataFrame: """Append empty checks at the end of DataFrame. @@ -405,13 +410,17 @@ def _create_results_array( empty_result = F.lit(None).cast(dq_result_schema).alias(dest_col) return df.select("*", empty_result) - checks = self._preselect_schema_validation_columns(df, checks) check_conditions = [] current_df = df for check in checks: + normalized_check = ( + self._preselect_schema_validation_columns(df, check) + if check.check_func is check_funcs.has_valid_schema + else check + ) manager = DQRuleManager( - check=check, + check=normalized_check, df=current_df, spark=self.spark, run_id=self.run_id, From d2f88ca4e4fe9c22371091738d502c7550653622 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 22 Jan 2026 09:16:42 +0100 Subject: [PATCH 08/15] Apply suggestion from @mwojtyczka refactor --- src/databricks/labs/dqx/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 33d8b1894..0f0070e2c 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -352,7 +352,7 @@ def _preselect_schema_validation_columns(self, df: DataFrame, check: DQRule) -> Determines the selected columns for schema validation checks. This is required when using `has_valid_schema` to ignore columns added during quality checking. This includes: * DQX result columns (e.g. '_warnings' and '_errors') - * Internal columns used for aggregate checks + * Internal columns used for dataset-level checks Args: df: Input DataFrame From 1f92e4816476500bd16f5e7feece739fdbf8601b Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 23 Jan 2026 10:24:32 -0500 Subject: [PATCH 09/15] Introduce registry for schema preselection --- src/databricks/labs/dqx/check_funcs.py | 3 ++- src/databricks/labs/dqx/engine.py | 25 +++++++++++++------------ src/databricks/labs/dqx/rule.py | 9 +++++++++ 3 files changed, 24 insertions(+), 13 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 893720367..450b9fb26 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -13,7 +13,7 @@ from pyspark.sql import types from pyspark.sql import Column, DataFrame, SparkSession from pyspark.sql.window import Window -from databricks.labs.dqx.rule import register_rule +from databricks.labs.dqx.rule import register_rule, register_for_schema_preselection from databricks.labs.dqx.utils import ( get_column_name_or_alias, is_sql_query_safe, @@ -1925,6 +1925,7 @@ def apply(df: DataFrame) -> DataFrame: @register_rule("dataset") +@register_for_schema_preselection("columns") def has_valid_schema( expected_schema: str | types.StructType | None = None, ref_df_name: str | None = None, diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 0f0070e2c..1e9a6e014 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -1,6 +1,7 @@ import copy -import os +import inspect import logging +import os from concurrent import futures from collections.abc import Callable from dataclasses import replace @@ -15,7 +16,6 @@ from pyspark.sql.streaming import StreamingQuery from databricks.labs.dqx.base import DQEngineBase, DQEngineCoreBase -from databricks.labs.dqx import check_funcs from databricks.labs.dqx.checks_resolver import resolve_custom_check_functions_from_path from databricks.labs.dqx.checks_serializer import deserialize_checks from databricks.labs.dqx.config_serializer import ConfigSerializer @@ -39,6 +39,7 @@ ColumnArguments, DefaultColumnNames, DQRule, + SCHEMA_PRESELECTION_REGISTRY, ) from databricks.labs.dqx.checks_validator import ChecksValidator, ChecksValidationStatus from databricks.labs.dqx.schema import dq_result_schema @@ -358,16 +359,20 @@ def _preselect_schema_validation_columns(self, df: DataFrame, check: DQRule) -> df: Input DataFrame check: DQRule to be modified """ - - # Default columns to all columns of the current DataFrame if not explicitly set - if check.check_func_kwargs.get("columns"): + if check.check_func.__name__ not in SCHEMA_PRESELECTION_REGISTRY: return check - if check.check_func_args and len(check.check_func_args) >= 3: + schema_argument = SCHEMA_PRESELECTION_REGISTRY[check.check_func.__name__] + if check.check_func_kwargs.get(schema_argument): return check + if check.check_func_args: + check_func_signature = inspect.signature(check.check_func) + if check_func_signature.parameters.get(schema_argument): + return check + rule_kwargs = check.check_func_kwargs.copy() - rule_kwargs["columns"] = [col for col in df.columns if col not in set(self._result_column_names.values())] + rule_kwargs[schema_argument] = [col for col in df.columns if col not in set(self._result_column_names.values())] return replace(check, check_func_kwargs=rule_kwargs) def _append_empty_checks(self, df: DataFrame) -> DataFrame: @@ -414,11 +419,7 @@ def _create_results_array( current_df = df for check in checks: - normalized_check = ( - self._preselect_schema_validation_columns(df, check) - if check.check_func is check_funcs.has_valid_schema - else check - ) + normalized_check = self._preselect_schema_validation_columns(df, check) manager = DQRuleManager( check=normalized_check, df=current_df, diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 96275bc53..9f1f78099 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -16,6 +16,7 @@ CHECK_FUNC_REGISTRY: dict[str, str] = {} +SCHEMA_PRESELECTION_REGISTRY: dict[str, str] = {} def register_rule(rule_type: str) -> Callable: @@ -26,6 +27,14 @@ def wrapper(func: Callable) -> Callable: return wrapper +def register_for_schema_preselection(schema_argument: str) -> Callable: + def wrapper(func: Callable) -> Callable: + SCHEMA_PRESELECTION_REGISTRY[func.__name__] = schema_argument + return func + + return wrapper + + class Criticality(Enum): """Enum class to represent criticality of the check.""" From f6e724649b11830691128d54be79fd023f39522e Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 23 Jan 2026 12:24:20 -0500 Subject: [PATCH 10/15] Add test coverage for positional and keyword "columns" arguments --- tests/integration/test_apply_checks.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 87a67d485..b1afc170f 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -8865,6 +8865,18 @@ def test_apply_checks_with_has_valid_schema_ignores_generated_columns(ws, spark, check_func=check_funcs.is_aggr_equal, check_func_kwargs={"aggr_type": "count", "column": "id", "limit": 1}, ), + DQDatasetRule( + name="has_valid_schema", + criticality="warn", + check_func=check_funcs.has_valid_schema, + check_func_kwargs={"expected_schema": "id int", "columns": ["id"], "strict": True} + ), + DQDatasetRule( + name="has_valid_schema", + criticality="warn", + check_func=check_funcs.has_valid_schema, + check_func_args=["id int", None, None, ["id"], True] + ), DQDatasetRule( name="has_valid_schema", criticality="warn", From c3f558724c09e791548d5e52c050575e434afd9d Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 23 Jan 2026 12:24:48 -0500 Subject: [PATCH 11/15] Fmt --- tests/integration/test_apply_checks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index b1afc170f..268344144 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -8869,13 +8869,13 @@ def test_apply_checks_with_has_valid_schema_ignores_generated_columns(ws, spark, name="has_valid_schema", criticality="warn", check_func=check_funcs.has_valid_schema, - check_func_kwargs={"expected_schema": "id int", "columns": ["id"], "strict": True} + check_func_kwargs={"expected_schema": "id int", "columns": ["id"], "strict": True}, ), DQDatasetRule( name="has_valid_schema", criticality="warn", check_func=check_funcs.has_valid_schema, - check_func_args=["id int", None, None, ["id"], True] + check_func_args=["id int", None, None, ["id"], True], ), DQDatasetRule( name="has_valid_schema", From c91c18fd2f0b40350c759f435b8cbdc4125eb7ce Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 23 Jan 2026 12:27:25 -0500 Subject: [PATCH 12/15] Remove downstreams workflow --- .github/workflows/downstreams.yml | 54 ------------------------------- 1 file changed, 54 deletions(-) delete mode 100644 .github/workflows/downstreams.yml diff --git a/.github/workflows/downstreams.yml b/.github/workflows/downstreams.yml deleted file mode 100644 index 87bf06f41..000000000 --- a/.github/workflows/downstreams.yml +++ /dev/null @@ -1,54 +0,0 @@ -name: downstreams - -on: - pull_request: - types: [opened, synchronize] - merge_group: - types: [checks_requested] - push: - # Always run on push to main. The build cache can only be reused - # if it was saved by a run from the repository's default branch. - # The run result will be identical to that from the merge queue - # because the commit is identical, yet we need to perform it to - # seed the build cache. - branches: - - main - -permissions: - id-token: write - contents: read - pull-requests: write - -jobs: - compatibility: - strategy: - fail-fast: false - matrix: - downstream: - #- name: ucx - - name: remorph - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v4.2.2 - with: - fetch-depth: 0 - - - name: Install Python - uses: actions/setup-python@v5 - with: - cache: 'pip' - cache-dependency-path: '**/pyproject.toml' - python-version: '3.12' - - - name: Install toolchain - run: | - pip install hatch==1.15.0 - - - name: Check downstream compatibility - uses: databrickslabs/sandbox/downstreams@downstreams/v0.0.1 - with: - repo: ${{ matrix.downstream.name }} - org: databrickslabs - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} From 51159688485bbce0409e7b4a7168ebf07bce831d Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 23 Jan 2026 19:12:47 -0500 Subject: [PATCH 13/15] Update method signatures and ignore column handling --- src/databricks/labs/dqx/check_funcs.py | 3 ++- src/databricks/labs/dqx/engine.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 450b9fb26..55631adfc 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -2017,7 +2017,8 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> selected_column_names = column_names if column_names else df.columns if ignore_column_names: - selected_column_names = [col for col in selected_column_names if col not in ignore_column_names] + ignore_set = set(ignore_column_names) + selected_column_names = [col for col in selected_column_names if col not in ignore_set] actual_schema = df.select(*selected_column_names).schema if strict: diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 1e9a6e014..1e8336ab5 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -348,7 +348,7 @@ def _all_are_dq_rules(checks: list[DQRule]) -> bool: """Check if all elements in the checks list are instances of DQRule.""" return all(isinstance(check, DQRule) for check in checks) - def _preselect_schema_validation_columns(self, df: DataFrame, check: DQRule) -> DQRule: + def _preselect_original_columns(self, df: DataFrame, check: DQRule) -> DQRule: """ Determines the selected columns for schema validation checks. This is required when using `has_valid_schema` to ignore columns added during quality checking. This includes: @@ -419,7 +419,7 @@ def _create_results_array( current_df = df for check in checks: - normalized_check = self._preselect_schema_validation_columns(df, check) + normalized_check = self._preselect_original_columns(df, check) manager = DQRuleManager( check=normalized_check, df=current_df, From ab4df8bb7092fa3b704a47bce88f4c3b37031409 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sat, 24 Jan 2026 11:04:04 +0100 Subject: [PATCH 14/15] refactor, improved test coverage --- src/databricks/labs/dqx/check_funcs.py | 8 +++---- src/databricks/labs/dqx/engine.py | 28 ++++++++++++++---------- src/databricks/labs/dqx/rule.py | 6 ++--- tests/integration/test_dataset_checks.py | 24 ++++++++++++++++++++ 4 files changed, 48 insertions(+), 18 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 55631adfc..6cca8aa48 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -13,7 +13,7 @@ from pyspark.sql import types from pyspark.sql import Column, DataFrame, SparkSession from pyspark.sql.window import Window -from databricks.labs.dqx.rule import register_rule, register_for_schema_preselection +from databricks.labs.dqx.rule import register_rule, register_for_original_columns_preselection from databricks.labs.dqx.utils import ( get_column_name_or_alias, is_sql_query_safe, @@ -1925,14 +1925,14 @@ def apply(df: DataFrame) -> DataFrame: @register_rule("dataset") -@register_for_schema_preselection("columns") +@register_for_original_columns_preselection() def has_valid_schema( expected_schema: str | types.StructType | None = None, ref_df_name: str | None = None, ref_table: str | None = None, columns: list[str | Column] | None = None, strict: bool = False, - ignore_columns: list[str] | None = None, + ignore_columns: list[str | Column] | None = None, ) -> tuple[Column, Callable]: """ Build a schema compatibility check condition and closure for dataset-level validation. @@ -1952,7 +1952,7 @@ def has_valid_schema( strict: Whether to perform strict schema validation (default: False). - False: Validates that all expected columns exist with compatible types (allows extra columns) - True: Validates exact schema match (same columns, same order, same types) - ignore_columns: Optional list of column names in the checked DataFrame schema to + ignore_columns: Optional list of columns in the checked DataFrame schema to ignore for validation. Returns: diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py index 1e8336ab5..cc97f8e9a 100644 --- a/src/databricks/labs/dqx/engine.py +++ b/src/databricks/labs/dqx/engine.py @@ -39,7 +39,7 @@ ColumnArguments, DefaultColumnNames, DQRule, - SCHEMA_PRESELECTION_REGISTRY, + CHECK_FUNC_REGISTRY_ORIGINAL_COLUMNS_PRESELECTION, ) from databricks.labs.dqx.checks_validator import ChecksValidator, ChecksValidationStatus from databricks.labs.dqx.schema import dq_result_schema @@ -350,29 +350,34 @@ def _all_are_dq_rules(checks: list[DQRule]) -> bool: def _preselect_original_columns(self, df: DataFrame, check: DQRule) -> DQRule: """ - Determines the selected columns for schema validation checks. This is required when using `has_valid_schema` to - ignore columns added during quality checking. This includes: - * DQX result columns (e.g. '_warnings' and '_errors') - * Internal columns used for dataset-level checks + Certain data quality checks (such as has_valid_schema) require access to the DataFrame's original schema—before + any DQX metadata columns, e.g. + * DQX result columns (e.g. '_warnings' and '_errors') + * Internal columns added by dataset-level checks + To enable this, check functions that need the original schema must be registered with + the register_for_original_columns_preselection decorator. Args: df: Input DataFrame - check: DQRule to be modified + check: Updated DQRule """ - if check.check_func.__name__ not in SCHEMA_PRESELECTION_REGISTRY: + # check func does not require original columns + if check.check_func.__name__ not in CHECK_FUNC_REGISTRY_ORIGINAL_COLUMNS_PRESELECTION: return check - schema_argument = SCHEMA_PRESELECTION_REGISTRY[check.check_func.__name__] - if check.check_func_kwargs.get(schema_argument): + # columns already provided in the check func kwargs + if check.check_func_kwargs.get("columns"): return check + # columns already provided in the check func args if check.check_func_args: check_func_signature = inspect.signature(check.check_func) - if check_func_signature.parameters.get(schema_argument): + if check_func_signature.parameters.get("columns"): return check + # preselect original columns rule_kwargs = check.check_func_kwargs.copy() - rule_kwargs[schema_argument] = [col for col in df.columns if col not in set(self._result_column_names.values())] + rule_kwargs["columns"] = [col for col in df.columns if col not in set(self._result_column_names.values())] return replace(check, check_func_kwargs=rule_kwargs) def _append_empty_checks(self, df: DataFrame) -> DataFrame: @@ -419,6 +424,7 @@ def _create_results_array( current_df = df for check in checks: + # each check pass may add new columns to the df and certain checks require original columns normalized_check = self._preselect_original_columns(df, check) manager = DQRuleManager( check=normalized_check, diff --git a/src/databricks/labs/dqx/rule.py b/src/databricks/labs/dqx/rule.py index 9f1f78099..0dd48b103 100644 --- a/src/databricks/labs/dqx/rule.py +++ b/src/databricks/labs/dqx/rule.py @@ -16,7 +16,7 @@ CHECK_FUNC_REGISTRY: dict[str, str] = {} -SCHEMA_PRESELECTION_REGISTRY: dict[str, str] = {} +CHECK_FUNC_REGISTRY_ORIGINAL_COLUMNS_PRESELECTION: set[str] = set() def register_rule(rule_type: str) -> Callable: @@ -27,9 +27,9 @@ def wrapper(func: Callable) -> Callable: return wrapper -def register_for_schema_preselection(schema_argument: str) -> Callable: +def register_for_original_columns_preselection() -> Callable: def wrapper(func: Callable) -> Callable: - SCHEMA_PRESELECTION_REGISTRY[func.__name__] = schema_argument + CHECK_FUNC_REGISTRY_ORIGINAL_COLUMNS_PRESELECTION.add(func.__name__) return func return wrapper diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index 80de951e1..3feb4f6c0 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -2798,3 +2798,27 @@ def test_has_valid_schema_with_ignore_columns(spark: SparkSession): "a string, b int, c double, d string, has_invalid_schema string", ) assert_df_equality(actual_condition_df, expected_condition_df, ignore_nullable=True) + + +def test_has_valid_schema_with_ignore_columns_as_expression(spark: SparkSession): + test_df = spark.createDataFrame( + [ + ["str1", 1, 100.0, "extra"], + ["str2", 2, 200.0, "data"], + ], + "a string, b int, c double, d string", + ) + + expected_schema = "a string, b int, c double" + condition, apply_method = has_valid_schema(expected_schema, ignore_columns=[F.col("d")], strict=True) + actual_apply_df = apply_method(test_df, spark, {}) + actual_condition_df = actual_apply_df.select("a", "b", "c", "d", condition) + + expected_condition_df = spark.createDataFrame( + [ + ["str1", 1, 100.0, "extra", None], + ["str2", 2, 200.0, "data", None], + ], + "a string, b int, c double, d string, has_invalid_schema string", + ) + assert_df_equality(actual_condition_df, expected_condition_df, ignore_nullable=True) From b6728d84f9c37878d11852efbb84773264abff66 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sat, 24 Jan 2026 19:39:44 +0100 Subject: [PATCH 15/15] renamed ignore_columns to exclude_columns --- docs/dqx/docs/reference/quality_checks.mdx | 6 +++--- src/databricks/labs/dqx/check_funcs.py | 18 +++++++++--------- tests/integration/test_dataset_checks.py | 8 ++++---- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index dc21897e9..e78dc4d9d 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1654,7 +1654,7 @@ You can also define your own custom dataset-level checks (see [Creating custom c | `sql_query` | Checks whether the condition column produced by a SQL query is satisfied. The check supports two modes: **Row-level validation** (when `merge_columns` is provided) - query results are joined back to the input DataFrame to mark specific rows; **Dataset-level validation** (when `merge_columns` is None or empty) - the check result applies to all rows (or filtered rows if `row_filter` is used), making it ideal for aggregate validations with custom metrics. The query must return a boolean condition column (True = fail, False = pass). For row-level checks: if merge columns aren't unique, multiple query rows can attach to a single input row, potentially causing false positives. Performance tip: for complex queries, writing a custom dataset-level rule is usually more performant than `sql_query` check. | `query`: query string, must return condition column (and merge columns if provided); `input_placeholder`: name to be used in the sql query as `{{ input_placeholder }}` to refer to the input DataFrame, optional reference DataFrames are referred by the name provided in the dictionary of reference DataFrames (e.g. `{{ ref_df_key }}`, dictionary of DataFrames can be passed when applying checks); `merge_columns`: (optional) list of columns used for merging with the input DataFrame which must exist in the input DataFrame and be present in output of the sql query; when not provided (None or empty list), the check result applies to all rows in the dataset (dataset-level validation); `condition_column`: name of the column indicating a violation (False = pass, True = fail); `msg`: (optional) message to output; `name`: (optional) name of the resulting check (it can be overwritten by `name` specified at the check level); `negate`: if the condition should be negated | | `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) | | `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_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 `ignore_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; `ignore_columns`: (optional) list of columns to ignore during validation (if not provided, all columns are considered); | +| `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); | **Compare datasets check** @@ -2007,7 +2007,7 @@ Complex data types are supported as well. function: has_valid_schema arguments: expected_schema: "id INT, name STRING, age INT, contact_info STRUCT" - ignore_columns: + exclude_columns: - last_update_date - last_updated_by @@ -2481,7 +2481,7 @@ checks = [ StructField("address", StringType(), True), ]), True) ]), - "ignore_columns": ["last_update_date", "last_updated_by"], + "exclude_columns": ["last_update_date", "last_updated_by"], }, ), diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 6cca8aa48..78eca1d91 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1932,7 +1932,7 @@ def has_valid_schema( ref_table: str | None = None, columns: list[str | Column] | None = None, strict: bool = False, - ignore_columns: list[str | Column] | None = None, + exclude_columns: list[str | Column] | None = None, ) -> tuple[Column, Callable]: """ Build a schema compatibility check condition and closure for dataset-level validation. @@ -1942,7 +1942,7 @@ def has_valid_schema( In strict mode, validates that the schema matches exactly (same columns, same order, same types) for the columns specified in columns or for all columns if columns is not specified. - All columns in the `ignore_columns` list will be ignored even if the column is present in the `columns` list. + All columns in the `exclude_columns` list will be ignored even if the column is present in the `columns` list. Args: expected_schema: Expected schema as a DDL string (e.g., "id INT, name STRING") or StructType object. @@ -1952,7 +1952,7 @@ def has_valid_schema( strict: Whether to perform strict schema validation (default: False). - False: Validates that all expected columns exist with compatible types (allows extra columns) - True: Validates exact schema match (same columns, same order, same types) - ignore_columns: Optional list of columns in the checked DataFrame schema to + exclude_columns: Optional list of columns in the checked DataFrame schema to ignore for validation. Returns: @@ -1982,10 +1982,10 @@ def has_valid_schema( if columns: column_names = [get_column_name_or_alias(col) if not isinstance(col, str) else col for col in columns] - ignore_column_names: list[str] | None = None - if ignore_columns: - ignore_column_names = [ - get_column_name_or_alias(col) if not isinstance(col, str) else col for col in ignore_columns + exclude_column_names: list[str] | None = None + if exclude_columns: + exclude_column_names = [ + get_column_name_or_alias(col) if not isinstance(col, str) else col for col in exclude_columns ] expected_schema = _get_schema(expected_schema or types.StructType(), column_names) @@ -2016,8 +2016,8 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> _expected_schema = expected_schema selected_column_names = column_names if column_names else df.columns - if ignore_column_names: - ignore_set = set(ignore_column_names) + if exclude_column_names: + ignore_set = set(exclude_column_names) selected_column_names = [col for col in selected_column_names if col not in ignore_set] actual_schema = df.select(*selected_column_names).schema diff --git a/tests/integration/test_dataset_checks.py b/tests/integration/test_dataset_checks.py index 3feb4f6c0..d3f3d1de3 100644 --- a/tests/integration/test_dataset_checks.py +++ b/tests/integration/test_dataset_checks.py @@ -2776,7 +2776,7 @@ def test_has_valid_schema_with_ref_df_name(spark: SparkSession): assert_df_equality(actual_condition_df, expected_condition_df, ignore_nullable=True) -def test_has_valid_schema_with_ignore_columns(spark: SparkSession): +def test_has_valid_schema_with_exclude_columns(spark: SparkSession): test_df = spark.createDataFrame( [ ["str1", 1, 100.0, "extra"], @@ -2786,7 +2786,7 @@ def test_has_valid_schema_with_ignore_columns(spark: SparkSession): ) expected_schema = "a string, b int, c double" - condition, apply_method = has_valid_schema(expected_schema, ignore_columns=["d"], strict=True) + condition, apply_method = has_valid_schema(expected_schema, exclude_columns=["d"], strict=True) actual_apply_df = apply_method(test_df, spark, {}) actual_condition_df = actual_apply_df.select("a", "b", "c", "d", condition) @@ -2800,7 +2800,7 @@ def test_has_valid_schema_with_ignore_columns(spark: SparkSession): assert_df_equality(actual_condition_df, expected_condition_df, ignore_nullable=True) -def test_has_valid_schema_with_ignore_columns_as_expression(spark: SparkSession): +def test_has_valid_schema_with_exclude_columns_as_expression(spark: SparkSession): test_df = spark.createDataFrame( [ ["str1", 1, 100.0, "extra"], @@ -2810,7 +2810,7 @@ def test_has_valid_schema_with_ignore_columns_as_expression(spark: SparkSession) ) expected_schema = "a string, b int, c double" - condition, apply_method = has_valid_schema(expected_schema, ignore_columns=[F.col("d")], strict=True) + condition, apply_method = has_valid_schema(expected_schema, exclude_columns=[F.col("d")], strict=True) actual_apply_df = apply_method(test_df, spark, {}) actual_condition_df = actual_apply_df.select("a", "b", "c", "d", condition)