diff --git a/demos/dqx_demo_library.py b/demos/dqx_demo_library.py index 19d3bbbe3..3d67ac193 100644 --- a/demos/dqx_demo_library.py +++ b/demos/dqx_demo_library.py @@ -943,11 +943,15 @@ def not_ends_with(column: str, suffix: str) -> Column: # COMMAND ---------- # MAGIC %md -# MAGIC #### Using `sql_query` check +# MAGIC #### Using `sql_query` check - Row-level validation +# MAGIC +# MAGIC The `sql_query` check supports two modes: +# MAGIC - **Row-level validation** (with `merge_columns`): Query results are joined back to mark specific rows +# MAGIC - **Dataset-level validation** (without `merge_columns`): Check result applies to all rows # COMMAND ---------- -# using DQX classes +# Row-level validation example: Check each sensor against its threshold from databricks.labs.dqx.rule import DQDatasetRule from databricks.labs.dqx.check_funcs import sql_query @@ -973,7 +977,7 @@ def not_ends_with(column: str, suffix: str) -> Column: check_func=sql_query, check_func_kwargs={ "query": query, - "merge_columns": ["sensor_id"], + "merge_columns": ["sensor_id"], # Results joined back by sensor_id "condition_column": "condition", # the check fails if this column evaluates to True "msg": "one of the sensor reading is greater than limit", "name": "sensor_reading_check", @@ -990,6 +994,41 @@ def not_ends_with(column: str, suffix: str) -> Column: # COMMAND ---------- +# MAGIC %md +# MAGIC #### Using `sql_query` check - Dataset-level validation +# MAGIC +# MAGIC When `merge_columns` is not provided, the check applies to all rows (all pass or all fail together). +# MAGIC This is useful for dataset-level aggregate validations. + +# COMMAND ---------- + +# Dataset-level validation example: Check total sensor count +dataset_query = """ + SELECT COUNT(DISTINCT sensor_id) < 1 AS condition + FROM {{ sensor }} +""" + +checks = [ + DQDatasetRule( + criticality="warn", + check_func=sql_query, + check_func_kwargs={ + "query": dataset_query, + # No merge_columns = dataset-level check (all rows get same result) + "condition_column": "condition", + "msg": "Dataset has no sensors", + "name": "dataset_has_sensors", + "input_placeholder": "sensor", + }, + ), +] + +ref_dfs = {"sensor_specs": sensor_specs_df} +valid_and_quarantine_df = dq_engine.apply_checks(sensor_df, checks, ref_dfs=ref_dfs) +display(valid_and_quarantine_df) + +# COMMAND ---------- + # using YAML declarative approach checks = yaml.safe_load( """ @@ -1028,6 +1067,30 @@ def not_ends_with(column: str, suffix: str) -> Column: # COMMAND ---------- +# YAML example for dataset-level validation (without merge_columns) +checks_dataset_level = yaml.safe_load( + """ + - criticality: warn + check: + function: sql_query + arguments: + # No merge_columns = dataset-level validation + condition_column: condition + msg: Dataset has no sensors + name: dataset_has_sensors + input_placeholder: sensor + query: | + SELECT COUNT(DISTINCT sensor_id) < 1 AS condition + FROM {{ sensor }} + """ +) + +ref_dfs = {"sensor_specs": sensor_specs_df} +valid_and_quarantine_df = dq_engine.apply_checks_by_metadata(sensor_df, checks_dataset_level, ref_dfs=ref_dfs) +display(valid_and_quarantine_df) + +# COMMAND ---------- + # MAGIC %md # MAGIC #### Defining custom python dataset-level check diff --git a/docs/dqx/docs/reference/benchmarks.mdx b/docs/dqx/docs/reference/benchmarks.mdx index 18c1bc423..653d48a90 100644 --- a/docs/dqx/docs/reference/benchmarks.mdx +++ b/docs/dqx/docs/reference/benchmarks.mdx @@ -23,6 +23,7 @@ sidebar_position: 13 | test_benchmark_compare_datasets | 3.598445 | 3.556993 | 3.430710 | 3.793938 | 0.158157 | 0.280218 | 3.466942 | 3.747160 | 5 | 0 | 2 | 0.28 | | test_benchmark_foreach_compare_datasets[n_rows_100000000_n_columns_5] | 25.879615 | 25.919933 | 25.536855 | 26.071184 | 0.217230 | 0.307223 | 25.748681 | 26.055904 | 5 | 0 | 1 | 0.04 | | test_benchmark_foreach_foreign_key[n_rows_100000000_n_columns_5] | 24.264873 | 22.893218 | 20.587308 | 29.037093 | 4.062789 | 7.705522 | 20.652819 | 28.358341 | 5 | 0 | 1 | 0.04 | +| test_benchmark_foreach_has_no_outliers[n_rows_100000000_n_columns_5] | 22.524313 | 22.347593 | 22.104944 | 22.924248 | 0.374170 | 0.646915 | 22.271984 | 22.918899 | 5 | 0 | 3 | 0.04 | | test_benchmark_foreach_has_valid_schema[n_rows_100000000_n_columns_5] | 1.068582 | 1.050490 | 0.979350 | 1.219259 | 0.092674 | 0.112164 | 1.003924 | 1.116088 | 5 | 0 | 1 | 0.94 | | test_benchmark_foreach_is_aggr_equal[n_rows_100000000_n_columns_5] | 1.239298 | 1.213153 | 1.192442 | 1.341836 | 0.060654 | 0.068928 | 1.200719 | 1.269646 | 5 | 0 | 1 | 0.81 | | test_benchmark_foreach_is_aggr_not_equal[n_rows_100000000_n_columns_5] | 1.264898 | 1.250273 | 1.218577 | 1.345211 | 0.051090 | 0.071957 | 1.225905 | 1.297862 | 5 | 0 | 1 | 0.79 | @@ -54,6 +55,7 @@ sidebar_position: 13 | test_benchmark_foreach_sql_query[n_rows_100000000_n_columns_5] | 4.578799 | 4.602143 | 4.442396 | 4.644892 | 0.083901 | 0.113694 | 4.530776 | 4.644470 | 5 | 0 | 1 | 0.22 | | test_benchmark_foreign_key | 31.784272 | 31.787610 | 31.414708 | 32.123221 | 0.269713 | 0.386951 | 31.597198 | 31.984149 | 5 | 0 | 2 | 0.03 | | test_benchmark_has_dimension | 0.215338 | 0.213285 | 0.210530 | 0.223131 | 0.005056 | 0.007086 | 0.211819 | 0.218905 | 5 | 0 | 1 | 4.64 | +| test_benchmark_has_no_outliers | 0.234952 | 0.228169 | 0.224165 | 0.257274 | 0.013649 | 0.017354 | 0.225936 | 0.243290 | 5 | 0 | 1 | 4.26 | | test_benchmark_has_valid_schema | 0.172078 | 0.172141 | 0.163793 | 0.181081 | 0.006715 | 0.009295 | 0.167010 | 0.176305 | 6 | 0 | 2 | 5.81 | | test_benchmark_has_x_coordinate_between | 0.217192 | 0.213656 | 0.209310 | 0.236233 | 0.011150 | 0.012638 | 0.209410 | 0.222048 | 5 | 0 | 1 | 4.60 | | test_benchmark_has_y_coordinate_between | 0.218497 | 0.219630 | 0.209352 | 0.234111 | 0.010103 | 0.013743 | 0.209584 | 0.223327 | 5 | 0 | 1 | 4.58 | diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 19851eecb..d5942f847 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1393,7 +1393,7 @@ You can also define your own custom dataset-level checks (see [Creating custom c | `is_aggr_equal` | Checks whether the aggregated values over group of rows or all rows are equal to the provided limit. | `column`: column to check (can be a string column name or a column expression), optional for 'count' aggregation; `limit`: limit as number, column name or sql expression (string literals must be single quoted, e.g. 'string_value'); `aggr_type`: aggregation function (default: "count"), supports 20 curated functions (count, sum, avg, stddev, percentile, etc.) plus any Databricks built-in aggregate; `group_by`: (optional) list of columns or column expressions to group the rows for aggregation (no grouping by default); `row_filter`: (optional) SQL expression to filter rows before aggregation; `aggr_params`: (optional) dict of parameters for aggregates requiring them | | `is_aggr_not_equal` | Checks whether the aggregated values over group of rows or all rows are not equal to the provided limit. | `column`: column to check (can be a string column name or a column expression), optional for 'count' aggregation; `limit`: limit as number, column name or sql expression (string literals must be single quoted, e.g. 'string_value'); `aggr_type`: aggregation function (default: "count"), supports 20 curated functions (count, sum, avg, stddev, percentile, etc.) plus any Databricks built-in aggregate; `group_by`: (optional) list of columns or column expressions to group the rows for aggregation (no grouping by default); `row_filter`: (optional) SQL expression to filter rows before aggregation; `aggr_params`: (optional) dict of parameters for aggregates requiring them | | `foreign_key` (aka is_in_list) | Checks whether input column or columns can be found in the reference DataFrame or Table (foreign key check). It supports foreign key check on single and composite keys. This check can be used to validate whether values in the input column(s) exist in a predefined list of allowed values (stored in the reference DataFrame or Table). It serves as a scalable alternative to `is_in_list` row-level checks, when working with large lists. | `columns`: columns to check (can be a list of string column names or column expressions); `ref_columns`: columns to check for existence in the reference DataFrame or Table (can be a list string column name or a column expression); `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; negate: if True the condition is negated (i.e. the check fails when the foreign key values exist in the reference DataFrame/Table), if False the check fails when the foreign key values do not exist in the reference | -| `sql_query` | Checks whether the condition column produced by a SQL query is satisfied. The check expects the query to return a boolean condition column indicating whether a record meets the requirement (True = fail, False = pass), and one or more merge columns so that results can be joined back to the input DataFrame to preserve all original records. Important considerations: if merge columns aren't unique, multiple query rows can attach to a single input row, potentially causing false positives. Performance tip: since the check must join back to the input DataFrame to retain original records, writing a custom dataset-level rule is usually more performant than `sql_query` check. | `query`: query string, must return all merge columns and condition column; `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`: 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; `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 | +| `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`: expected schema as a DDL string (e.g., "id INT, name STRING") or StructType object; `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 | @@ -1649,7 +1649,7 @@ Complex data types are supported as well. ref_df_name: ref_df_key negate: true -# sql_query check +# sql_query check with merge_columns (row-level validation) - criticality: error check: function: sql_query @@ -1665,6 +1665,18 @@ Complex data types are supported as well. name: sql_query_violation # optional negate: false # optional, default False +# sql_query check without merge_columns (dataset-level validation) +- criticality: error + check: + function: sql_query + arguments: + # sql query for dataset-level check (must return 1 record) + query: SELECT COUNT(*) = 0 AS condition FROM {{ input_view }} + input_placeholder: input_view # name to be used in the sql query as `{{ input_view }}` to refer to the input DataFrame + condition_column: condition # the check fails if this column evaluates to True + msg: dataset has no records # optional + name: dataset_is_empty # optional + # compare_datasets check - criticality: error check: @@ -2021,7 +2033,7 @@ checks = [ "negate": True }, - # sql_query check + # sql_query check with merge_columns (row-level validation) DQDatasetRule( criticality="error", check_func=sql_query, @@ -2034,7 +2046,21 @@ checks = [ "msg": "sql query check failed", # optional "name": "sql_query_violation", # optional "negate": False # optional, default False - }, + }, + ), + + # sql_query check without merge_columns (dataset-level validation) + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + # sql query for dataset-level check (must return 1 record) + "query": "SELECT COUNT(*) = 0 AS condition FROM {{ input_view }}", + "input_placeholder": "input_view", # name to be used in the sql query as `{{ input_placeholder }}` to refer to the input DataFrame + "condition_column": "condition", # the check fails if this column evaluates to True + "msg": "dataset has no records", # optional + "name": "dataset_is_empty", # optional + }, ), # compare_datasets check diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 9a968370d..7b5f88687 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -3,7 +3,7 @@ import warnings import ipaddress import uuid -from collections.abc import Callable +from collections.abc import Callable, Sequence from enum import Enum from itertools import zip_longest import operator as py_operator @@ -1265,7 +1265,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> @register_rule("dataset") def sql_query( query: str, - merge_columns: list[str], + merge_columns: list[str] | None = None, msg: str | None = None, name: str | None = None, negate: bool = False, @@ -1276,13 +1276,24 @@ def sql_query( """ Checks whether the condition column generated by SQL query is met. + Supports two modes: + - Row-level validation (merge_columns provided): Query results are joined back to specific rows + - Dataset-level validation (merge_columns omitted or None): All rows get the same check result + + Use dataset-level for aggregate validations like "total count > 100" or "avg(amount) < 1000". + Use row-level when you need to identify specific problematic rows. + Args: query: SQL query that must return as a minimum a condition column and - all merge columns. The resulting DataFrame is automatically joined back to the input DataFrame - using the merge_columns. Reference DataFrames when provided in the ref_dfs parameter are registered as temp view. + all merge columns (if provided). When merge_columns are provided, the resulting DataFrame is + automatically joined back to the input DataFrame. When merge_columns are not provided, the check + applies to all rows (either all pass or all fail), making it useful for dataset-level validation + with custom_metrics. Reference DataFrames when provided in the ref_dfs parameter are registered as temp view. + merge_columns: Optional (can be None or omitted). List of columns to join results back to input DataFrame. + - If provided: Row-level validation - different rows can have different results + - If None/omitted: Dataset-level validation - all rows get same result + When provided, columns must form a unique key to avoid duplicate records. condition_column: Column name indicating violation (boolean). Fail the check if True, pass it if False - merge_columns: List of columns for join back to the input DataFrame. - They must provide a unique key for the join, otherwise a duplicate records may be produced. msg: Optional custom message or Column expression. name: Optional name for the result. negate: If True, the condition is negated (i.e., the check fails when the condition is False). @@ -1295,13 +1306,23 @@ def sql_query( Tuple (condition column, apply function). Raises: - MissingParameterError: if *merge_columns* is None. - InvalidParameterError: if *merge_columns* is an empty list. UnsafeSqlQueryError: if the SQL query fails the safety check (e.g., contains disallowed operations). """ _validate_sql_query_params(query, merge_columns) - alias_name = name if name else "_".join(merge_columns) + f"_query_{condition_column}_violation" + # Normalize empty list to None (both mean "no merge columns" / dataset-level check) + if merge_columns is not None and not merge_columns: + merge_columns = None + + alias_name = ( + name + if name + else ( + "_".join(merge_columns) + f"_query_{condition_column}_violation" + if merge_columns + else f"query_{condition_column}_violation" + ) + ) unique_str = uuid.uuid4().hex # make sure any column added to the dataframe is unique unique_condition_column = f"{alias_name}_{condition_column}_{unique_str}" @@ -1338,6 +1359,12 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> "Resolved SQL query is not safe for execution. Please ensure it does not contain any unsafe operations." ) + # When merge_columns is None, the check applies to all rows (dataset-level check) + if merge_columns is None: + return _apply_dataset_level_sql_check( + df, spark, query_resolved, condition_column, unique_condition_column, row_filter + ) + # Resolve the SQL query against the input DataFrame and any reference DataFrames user_query_df = spark.sql(query_resolved).select( *merge_columns, F.col(condition_column).alias(unique_condition_column) @@ -2974,32 +3001,89 @@ def _is_valid_ipv6_cidr_block(cidr: str) -> bool: return False -def _validate_sql_query_params(query: str, merge_columns: list[str]) -> None: +def _apply_dataset_level_sql_check( + df: DataFrame, + spark: SparkSession, + query_resolved: str, + condition_column: str, + unique_condition_column: str, + row_filter: str | None, +) -> DataFrame: + """ + Apply a dataset-level SQL check where all rows get the same validation result. + + Args: + df: Input DataFrame to apply the check to. + spark: SparkSession for executing SQL. + query_resolved: The resolved SQL query (with placeholders replaced). + condition_column: Name of the condition column in the query result. + unique_condition_column: Unique name for the condition column in the output. + row_filter: Optional SQL expression for filtering which rows receive the check result. + + Returns: + DataFrame with the condition column added. + """ + # Query should only return the condition column + user_query_df = spark.sql(query_resolved).select(F.col(condition_column).alias(unique_condition_column)) + + # Capture up to two rows to detect accidental multi-row outputs + condition_rows = user_query_df.take(2) + if not condition_rows: + # No rows returned: treat as condition not met (False) + condition_value = False + elif len(condition_rows) > 1: + raise InvalidParameterError( + "Dataset-level sql_query without merge_columns must return exactly one row. " + "Provide merge_columns for row-level checks or aggregate the query to a single row." + ) + else: + condition_value = condition_rows[0][unique_condition_column] + + # Apply the condition consistently with row_filter behavior: + # - If row_filter is provided, only rows matching the filter get the condition value + # - Rows not matching the filter get None (consistent with row-level checks) + if row_filter: + filter_expr = F.expr(row_filter) + result_df = df.withColumn( + unique_condition_column, F.when(filter_expr, F.lit(condition_value)).otherwise(F.lit(None)) + ) + else: + # No filter: apply condition to all rows + result_df = df.withColumn(unique_condition_column, F.lit(condition_value)) + + return result_df + + +def _validate_sql_query_params(query: str, merge_columns: list[str] | None) -> None: """ Validate SQL query parameters to ensure correctness and safety. - This helper verifies that: - - The SQL query is provided and is a non-empty string. - - The 'merge_columns' parameter is provided and is a non-empty list of column names. - - The 'merge_columns' list contains valid column names. + This helper verifies that the SQL query is safe for execution. Args: query: The SQL query string to validate. - merge_columns: The list of column names to validate. + merge_columns: Optional list of column names (validated when provided). Raises: - MissingParameterError: If any required parameter is missing. - InvalidParameterError: If any parameter is invalid. UnsafeSqlQueryError: If the SQL query is unsafe. + InvalidParameterError: If merge_columns is provided but not a sequence of strings. """ - if merge_columns is None: - raise MissingParameterError("'merge_columns' is required and must be a non-empty list of column names.") - if not merge_columns: - raise InvalidParameterError("'merge_columns' must contain at least one column.") if not is_sql_query_safe(query): raise UnsafeSqlQueryError( "Provided SQL query is not safe for execution. Please ensure it does not contain any unsafe operations." ) + if merge_columns is None: + return + + if not isinstance(merge_columns, Sequence) or isinstance(merge_columns, str): + raise InvalidParameterError( + "'merge_columns' must be a sequence of column names (e.g., list or tuple) when provided." + ) + + invalid_columns = [col for col in merge_columns if not isinstance(col, str) or not col] + if invalid_columns: + raise InvalidParameterError("'merge_columns' entries must be non-empty strings.") + def _calculate_median_absolute_deviation(df: DataFrame, column: str, filter_condition: str | None) -> tuple[Any, Any]: """ diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 7a0d63405..38fd7b9cd 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -33,6 +33,27 @@ EXTRA_PARAMS = ExtraParams(run_time_overwrite=RUN_TIME.isoformat(), run_id_overwrite=RUN_ID) +def build_quality_violation( + name: str, + message: str, + columns: list[str] | None, + *, + function: str = "is_not_null_and_not_empty", +) -> dict[str, Any]: + """Helper for constructing expected violation entries with shared metadata.""" + + return { + "name": name, + "message": message, + "columns": columns, + "filter": None, + "function": function, + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + } + + @pytest.fixture def webbrowser_open(): with patch("webbrowser.open") as mock_open: @@ -301,16 +322,9 @@ def expected_quality_checking_output(spark) -> DataFrame: 3, None, [ - { - "name": "name_is_not_null_and_not_empty", - "message": "Column 'name' value is null or empty", - "columns": ["name"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } + build_quality_violation( + "name_is_not_null_and_not_empty", "Column 'name' value is null or empty", ["name"] + ) ], None, ], @@ -318,16 +332,9 @@ def expected_quality_checking_output(spark) -> DataFrame: None, "c", [ - { - "name": "id_is_not_null", - "message": "Column 'id' value is null", - "columns": ["id"], - "filter": None, - "function": "is_not_null", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - }, + build_quality_violation( + "id_is_not_null", "Column 'id' value is null", ["id"], function="is_not_null" + ) ], None, ], @@ -335,16 +342,9 @@ def expected_quality_checking_output(spark) -> DataFrame: 3, None, [ - { - "name": "name_is_not_null_and_not_empty", - "message": "Column 'name' value is null or empty", - "columns": ["name"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } + build_quality_violation( + "name_is_not_null_and_not_empty", "Column 'name' value is null or empty", ["name"] + ) ], None, ], diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 647336f5d..764e7120d 100644 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -26,7 +26,7 @@ from databricks.labs.dqx import check_funcs from tests.conftest import TEST_CATALOG -from tests.integration.conftest import REPORTING_COLUMNS, RUN_TIME, EXTRA_PARAMS, RUN_ID +from tests.integration.conftest import REPORTING_COLUMNS, RUN_TIME, EXTRA_PARAMS, RUN_ID, build_quality_violation SCHEMA = "a: int, b: int, c: int" @@ -3065,6 +3065,584 @@ def test_apply_checks_with_sql_query_and_ref_df(ws, spark): assert_df_equality(checked, expected, ignore_nullable=True) +def test_apply_checks_with_sql_query_without_merge_columns(ws, spark): + """Test sql_query check without merge_columns - dataset-level validation.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 3], [2, None, 3], [1, None, 4], [None, None, None]], SCHEMA) + + # Dataset-level check without merge_columns + # Query returns a single row with a condition column + # The check will fail because COUNT(*) > 2 is True (we have 4 rows) + query_fail = "SELECT COUNT(*) > 2 AS condition FROM {{input_view}}" + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": query_fail, + "condition_column": "condition", + "msg": "Dataset has more than 2 rows", + "name": "dataset_check_fail", + }, + ), + ] + + checked = dq_engine.apply_checks(test_df, checks) + + # All rows should have the same error result since it's a dataset-level check + expected = spark.createDataFrame( + [ + [ + 1, + 3, + 3, + [ + { + "name": "dataset_check_fail", + "message": "Dataset has more than 2 rows", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 2, + None, + 3, + [ + { + "name": "dataset_check_fail", + "message": "Dataset has more than 2 rows", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 1, + None, + 4, + [ + { + "name": "dataset_check_fail", + "message": "Dataset has more than 2 rows", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + None, + None, + None, + [ + { + "name": "dataset_check_fail", + "message": "Dataset has more than 2 rows", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_and_filter(ws, spark): + """Test sql_query check without merge_columns with row_filter - should only apply to filtered rows.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 3], [2, 10, 3], [1, 20, 4], [None, None, None]], SCHEMA) + + # Dataset-level check without merge_columns but WITH a filter + # The query checks if SUM(a) > 2, but only for rows where b >= 10 + # Only 2 rows match the filter: [2, 10, 3] and [1, 20, 4], and SUM(a) = 3 > 2, so condition is True + query_with_filter = "SELECT SUM(a) > 2 AS condition FROM {{input_view}}" + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + filter="b >= 10", # Only apply to rows where b >= 10 + check_func_kwargs={ + "query": query_with_filter, + "condition_column": "condition", + "msg": "Filtered sum check failed", + "name": "dataset_filtered_check", + }, + ), + ] + + checked = dq_engine.apply_checks(test_df, checks) + + # Only rows matching the filter (b >= 10) should have the error + # Rows not matching the filter should have None for both _errors and _warnings + expected = spark.createDataFrame( + [ + [ + 1, + 3, + 3, + None, # b=3 doesn't match filter (b >= 10), so no error + None, + ], + [ + 2, + 10, + 3, + [ + { + "name": "dataset_filtered_check", + "message": "Filtered sum check failed", + "columns": None, + "filter": "b >= 10", + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 1, + 20, + 4, + [ + { + "name": "dataset_filtered_check", + "message": "Filtered sum check failed", + "columns": None, + "filter": "b >= 10", + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + None, + None, + None, + None, # b=None doesn't match filter, so no error + None, + ], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_passes(ws, spark): + """Test sql_query without merge_columns where condition evaluates to False (no violations).""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 3], [2, None, 3]], SCHEMA) + + # Dataset-level check that passes: COUNT(*) > 100 is False (we only have 2 rows) + query_pass = "SELECT COUNT(*) > 100 AS condition FROM {{input_view}}" + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": query_pass, + "condition_column": "condition", + "msg": "Dataset has too many rows", + "name": "dataset_size_check", + }, + ), + ] + + checked = dq_engine.apply_checks(test_df, checks) + + # All rows should have None for errors since condition is False + expected = spark.createDataFrame( + [ + [1, 3, 3, None, None], + [2, None, 3, None, None], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_empty_result(ws, spark): + """Test sql_query without merge_columns where query returns no rows (treated as False).""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 3], [2, None, 3]], SCHEMA) + + # Query with WHERE 1=0 returns no rows + query_no_rows = "SELECT TRUE AS condition FROM {{input_view}} WHERE 1=0" + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": query_no_rows, + "condition_column": "condition", + "msg": "Empty query result", + "name": "empty_query_check", + }, + ), + ] + + checked = dq_engine.apply_checks(test_df, checks) + + # No rows from query means condition is False, so no violations + expected = spark.createDataFrame( + [ + [1, 3, 3, None, None], + [2, None, 3, None, None], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_negate(ws, spark): + """Test sql_query without merge_columns with negate=True.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 3], [2, None, 3]], SCHEMA) + + # Query that returns False, but with negate=True should fail + query_false = "SELECT COUNT(*) > 100 AS condition FROM {{input_view}}" + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": query_false, + "condition_column": "condition", + "msg": "Dataset does not have enough rows", + "name": "dataset_min_size_check", + "negate": True, # Fail when condition is False + }, + ), + ] + + checked = dq_engine.apply_checks(test_df, checks) + + # With negate=True, False becomes violation + expected = spark.createDataFrame( + [ + [ + 1, + 3, + 3, + [ + { + "name": "dataset_min_size_check", + "message": "Dataset does not have enough rows", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 2, + None, + 3, + [ + { + "name": "dataset_min_size_check", + "message": "Dataset does not have enough rows", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_warning(ws, spark): + """Test sql_query without merge_columns with warning criticality.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 3, 3], [2, None, 3], [3, 4, 5]], SCHEMA) + + # Dataset-level warning: COUNT(*) > 2 is True + query_warn = "SELECT COUNT(*) > 2 AS condition FROM {{input_view}}" + + checks = [ + DQDatasetRule( + criticality="warn", # Warning instead of error + check_func=sql_query, + check_func_kwargs={ + "query": query_warn, + "condition_column": "condition", + "msg": "Dataset has more than 2 rows (warning)", + "name": "dataset_warn_check", + }, + ), + ] + + checked = dq_engine.apply_checks(test_df, checks) + + # All rows should have warning (not error) + expected = spark.createDataFrame( + [ + [ + 1, + 3, + 3, + None, + [ + { + "name": "dataset_warn_check", + "message": "Dataset has more than 2 rows (warning)", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + ], + [ + 2, + None, + 3, + None, + [ + { + "name": "dataset_warn_check", + "message": "Dataset has more than 2 rows (warning)", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + ], + [ + 3, + 4, + 5, + None, + [ + { + "name": "dataset_warn_check", + "message": "Dataset has more than 2 rows (warning)", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + ], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_and_ref_df(ws, spark): + """Test sql_query without merge_columns (dataset-level) with reference DataFrame.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + + # Main dataset + test_df = spark.createDataFrame([[1, 10, 100], [2, 20, 200], [3, 30, 300]], SCHEMA) + + # Reference dataset with expected totals + ref_df = spark.createDataFrame([[600]], "expected_total: int") + + # Dataset-level check: violation occurs when totals do NOT match the reference table + query = """ + SELECT (SELECT SUM(c) FROM {{input_view}}) != (SELECT expected_total FROM {{expected_totals}}) AS condition + """ + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": query, + "condition_column": "condition", + "msg": "Total amount matches expected", + "name": "multi_dataset_check", + }, + ), + ] + + ref_dfs = {"expected_totals": ref_df} + checked = dq_engine.apply_checks(test_df, checks, ref_dfs=ref_dfs) + + # All rows should pass (SUM(100,200,300) = 600) + expected = spark.createDataFrame( + [ + [1, 10, 100, None, None], + [2, 20, 200, None, None], + [3, 30, 300, None, None], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_and_ref_df_fail(ws, spark): + """Test sql_query without merge_columns (dataset-level) with reference DataFrame that fails.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + + # Main dataset + test_df = spark.createDataFrame([[1, 10, 100], [2, 20, 200], [3, 30, 300]], SCHEMA) + + # Reference dataset with expected totals (wrong value) + ref_df = spark.createDataFrame([[999]], "expected_total: int") + + # Dataset-level check: violation occurs when totals do NOT match the reference table + query = """ + SELECT (SELECT SUM(c) FROM {{input_view}}) != (SELECT expected_total FROM {{expected_totals}}) AS condition + """ + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": query, + "condition_column": "condition", + "msg": "Total amount does not match expected", + "name": "multi_dataset_check", + }, + ), + ] + + ref_dfs = {"expected_totals": ref_df} + checked = dq_engine.apply_checks(test_df, checks, ref_dfs=ref_dfs) + + # All rows should fail (SUM(100,200,300) = 600 != 999) + expected = spark.createDataFrame( + [ + [ + 1, + 10, + 100, + [ + { + "name": "multi_dataset_check", + "message": "Total amount does not match expected", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 2, + 20, + 200, + [ + { + "name": "multi_dataset_check", + "message": "Total amount does not match expected", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + [ + 3, + 30, + 300, + [ + { + "name": "multi_dataset_check", + "message": "Total amount does not match expected", + "columns": None, + "filter": None, + "function": "sql_query", + "run_time": RUN_TIME, + "run_id": RUN_ID, + "user_metadata": {}, + }, + ], + None, + ], + ], + EXPECTED_SCHEMA, + ) + assert_df_equality(checked, expected, ignore_nullable=True) + + +def test_apply_checks_with_sql_query_without_merge_columns_multiple_rows_error(ws, spark): + """Ensure dataset-level sql_query errors when the query returns more than one row.""" + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + test_df = spark.createDataFrame([[1, 10, 100], [2, 20, 200]], SCHEMA) + + query_multi_rows = "SELECT a > 0 AS condition FROM {{input_view}}" + + checks = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": query_multi_rows, + "condition_column": "condition", + "msg": "Should never reach application", + "name": "dataset_multi_row_error", + }, + ), + ] + + with pytest.raises( + InvalidParameterError, match="Dataset-level sql_query without merge_columns must return exactly one row" + ): + dq_engine.apply_checks(test_df, checks) + + def test_apply_checks_with_sql_query_and_ref_table(ws, spark): dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) @@ -3881,18 +4459,7 @@ def test_apply_checks_by_metadata_with_custom_column_naming(ws, spark): 2, None, 4, - [ - { - "name": "b_is_null_or_empty", - "message": "Column 'b' value is null or empty", - "columns": ["b"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], + [build_quality_violation("b_is_null_or_empty", "Column 'b' value is null or empty", ["b"])], None, ], [ @@ -3900,47 +4467,14 @@ def test_apply_checks_by_metadata_with_custom_column_naming(ws, spark): 4, None, None, - [ - { - "name": "a_is_null_or_empty", - "message": "Column 'a' value is null or empty", - "columns": ["a"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], + [build_quality_violation("a_is_null_or_empty", "Column 'a' value is null or empty", ["a"])], ], [ None, None, None, - [ - { - "name": "b_is_null_or_empty", - "message": "Column 'b' value is null or empty", - "columns": ["b"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], - [ - { - "name": "a_is_null_or_empty", - "message": "Column 'a' value is null or empty", - "columns": ["a"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], + [build_quality_violation("b_is_null_or_empty", "Column 'b' value is null or empty", ["b"])], + [build_quality_violation("a_is_null_or_empty", "Column 'a' value is null or empty", ["a"])], ], ], EXPECTED_SCHEMA_WITH_CUSTOM_NAMES, @@ -3975,18 +4509,7 @@ def test_apply_checks_by_metadata_with_custom_column_naming_fallback_to_default( 2, None, 4, - [ - { - "name": "b_is_null_or_empty", - "message": "Column 'b' value is null or empty", - "columns": ["b"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], + [build_quality_violation("b_is_null_or_empty", "Column 'b' value is null or empty", ["b"])], None, ], [ @@ -3994,47 +4517,14 @@ def test_apply_checks_by_metadata_with_custom_column_naming_fallback_to_default( 4, None, None, - [ - { - "name": "a_is_null_or_empty", - "message": "Column 'a' value is null or empty", - "columns": ["a"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], + [build_quality_violation("a_is_null_or_empty", "Column 'a' value is null or empty", ["a"])], ], [ None, None, None, - [ - { - "name": "b_is_null_or_empty", - "message": "Column 'b' value is null or empty", - "columns": ["b"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], - [ - { - "name": "a_is_null_or_empty", - "message": "Column 'a' value is null or empty", - "columns": ["a"], - "filter": None, - "function": "is_not_null_and_not_empty", - "run_time": RUN_TIME, - "run_id": RUN_ID, - "user_metadata": {}, - } - ], + [build_quality_violation("b_is_null_or_empty", "Column 'b' value is null or empty", ["b"])], + [build_quality_violation("a_is_null_or_empty", "Column 'a' value is null or empty", ["a"])], ], ], EXPECTED_SCHEMA, @@ -7294,7 +7784,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): [ { "name": "count_greater_than_limit", - "message": "Count 3 in column '*' is greater than limit: 0", + "message": "Count value 3 in column '*' is greater than limit: 0", "columns": ["*"], "filter": None, "function": "is_aggr_not_greater_than", @@ -7314,7 +7804,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): }, { "name": "count_less_than_limit", - "message": "Count 3 in column '*' is less than limit: 10", + "message": "Count value 3 in column '*' is less than limit: 10", "columns": ["*"], "filter": None, "function": "is_aggr_not_less_than", @@ -7324,7 +7814,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): }, { "name": "count_equal_to_limit", - "message": "Count 3 in column '*' is equal to limit: 3", + "message": "Count value 3 in column '*' is equal to limit: 3", "columns": ["*"], "filter": None, "function": "is_aggr_not_equal", @@ -7403,7 +7893,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): [ { "name": "count_greater_than_limit", - "message": "Count 3 in column '*' is greater than limit: 0", + "message": "Count value 3 in column '*' is greater than limit: 0", "columns": ["*"], "filter": None, "function": "is_aggr_not_greater_than", @@ -7423,7 +7913,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): }, { "name": "count_less_than_limit", - "message": "Count 3 in column '*' is less than limit: 10", + "message": "Count value 3 in column '*' is less than limit: 10", "columns": ["*"], "filter": None, "function": "is_aggr_not_less_than", @@ -7433,7 +7923,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): }, { "name": "count_equal_to_limit", - "message": "Count 3 in column '*' is equal to limit: 3", + "message": "Count value 3 in column '*' is equal to limit: 3", "columns": ["*"], "filter": None, "function": "is_aggr_not_equal", @@ -7542,7 +8032,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): [ { "name": "count_greater_than_limit", - "message": "Count 3 in column '*' is greater than limit: 0", + "message": "Count value 3 in column '*' is greater than limit: 0", "columns": ["*"], "filter": None, "function": "is_aggr_not_greater_than", @@ -7572,7 +8062,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): }, { "name": "count_less_than_limit", - "message": "Count 3 in column '*' is less than limit: 10", + "message": "Count value 3 in column '*' is less than limit: 10", "columns": ["*"], "filter": None, "function": "is_aggr_not_less_than", @@ -7582,7 +8072,7 @@ def test_apply_aggr_checks_by_metadata(ws, spark): }, { "name": "count_equal_to_limit", - "message": "Count 3 in column '*' is equal to limit: 3", + "message": "Count value 3 in column '*' is equal to limit: 3", "columns": ["*"], "filter": None, "function": "is_aggr_not_equal", diff --git a/tests/perf/.benchmarks/baseline.json b/tests/perf/.benchmarks/baseline.json index 66e28fffd..61d4c963f 100644 --- a/tests/perf/.benchmarks/baseline.json +++ b/tests/perf/.benchmarks/baseline.json @@ -344,6 +344,46 @@ "iterations": 1 } }, + { + "group": "test_benchmark_foreach_is_not_greater_than_100000000_rows_5_columns", + "name": "test_benchmark_foreach_has_no_outliers[n_rows_100000000_n_columns_5]", + "fullname": "tests/perf/test_apply_checks.py::test_benchmark_foreach_has_no_outliers[n_rows_100000000_n_columns_5]", + "params": { + "generated_integer_df": { + "n_rows": 100000000, + "n_columns": 5 + } + }, + "param": "n_rows_100000000_n_columns_5", + "extra_info": {}, + "options": { + "disable_gc": false, + "timer": "perf_counter", + "min_rounds": 5, + "max_time": 1.0, + "min_time": 0.000005, + "warmup": false + }, + "stats": { + "min": 22.104943891999937, + "max": 22.924248213999817, + "mean": 22.52431289160013, + "stddev": 0.37417012250544546, + "rounds": 5, + "median": 22.3475933360005, + "iqr": 0.6469148584997129, + "q1": 22.27198371500026, + "q3": 22.918898573499973, + "iqr_outliers": 0, + "stddev_outliers": 3, + "outliers": "3;0", + "ld15iqr": 22.104943891999937, + "hd15iqr": 22.924248213999817, + "ops": 0.04439647081855823, + "total": 112.62156445800065, + "iterations": 1 + } + }, { "group": "test_benchmark_foreach_has_valid_schema_100000000_rows_5_columns", "name": "test_benchmark_foreach_has_valid_schema[n_rows_100000000_n_columns_5]", @@ -1576,6 +1616,41 @@ "iterations": 1 } }, + { + "group": null, + "name": "test_benchmark_has_no_outliers", + "fullname": "tests/perf/test_apply_checks.py::test_benchmark_has_no_outliers", + "params": null, + "param": null, + "extra_info": {}, + "options": { + "disable_gc": false, + "timer": "perf_counter", + "min_rounds": 5, + "max_time": 1.0, + "min_time": 0.000005, + "warmup": false + }, + "stats": { + "min": 0.2241649769994183, + "max": 0.2572740440000416, + "mean": 0.23495235979971768, + "stddev": 0.01364878347106346, + "rounds": 5, + "median": 0.22816866199991637, + "iqr": 0.017353771250554928, + "q1": 0.22593578549935955, + "q3": 0.24328955674991448, + "iqr_outliers": 0, + "stddev_outliers": 1, + "outliers": "1;0", + "ld15iqr": 0.2241649769994183, + "hd15iqr": 0.2572740440000416, + "ops": 4.256181980263736, + "total": 1.1747617989985883, + "iterations": 1 + } + }, { "group": null, "name": "test_benchmark_has_valid_schema", diff --git a/tests/resources/all_dataset_checks.yaml b/tests/resources/all_dataset_checks.yaml index 39402900b..390ee9923 100644 --- a/tests/resources/all_dataset_checks.yaml +++ b/tests/resources/all_dataset_checks.yaml @@ -163,7 +163,7 @@ - ref_col2 ref_df_name: ref_df_key -# sql_query check +# sql_query check with merge_columns (row-level validation) - criticality: error check: function: sql_query @@ -179,6 +179,17 @@ name: sql_query_violation # optional negate: false # optional, default False +# sql_query check without merge_columns (dataset-level validation) +- criticality: warn + check: + function: sql_query + arguments: + # sql query for dataset-level check - no merge_columns needed + query: SELECT COUNT(*) = 0 AS condition FROM {{ input_view }} + condition_column: condition # the check fails if this column evaluates to True + msg: dataset has no records # optional + name: dataset_is_empty # optional + # apply check to multiple columns - criticality: error check: diff --git a/tests/unit/test_dataset_checks.py b/tests/unit/test_dataset_checks.py index 58180f159..871d37ded 100644 --- a/tests/unit/test_dataset_checks.py +++ b/tests/unit/test_dataset_checks.py @@ -145,13 +145,26 @@ def test_compare_datasets_invalid_tolerance_exceptions(abs_tolerance, rel_tolera ) -def test_sql_query_missing_merge_columns(): - with pytest.raises(InvalidParameterError, match="'merge_columns' must contain at least one column"): - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={"query": "SELECT 1", "merge_columns": [], "condition_column": "condition"}, - ) +def test_sql_query_empty_merge_columns(): + """Test that empty list for merge_columns is treated the same as None (dataset-level check).""" + # Should not raise an error - empty list is normalized to None + rule = DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={"query": "SELECT FALSE AS condition", "merge_columns": [], "condition_column": "condition"}, + ) + assert rule.check_func is sql_query + + +def test_sql_query_without_merge_columns(): + """Test that merge_columns is optional and can be None.""" + # Should not raise an error + rule = DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={"query": "SELECT FALSE AS condition", "condition_column": "condition"}, + ) + assert rule.check_func is sql_query def test_sql_query_unsafe(): @@ -164,6 +177,36 @@ def test_sql_query_unsafe(): ) +def test_sql_query_merge_columns_as_string_raises(): + """Ensure merge_columns must be provided as a sequence, not a single string.""" + with pytest.raises( + InvalidParameterError, match="'merge_columns' must be a sequence of column names \\(e.g., list or tuple\\)" + ): + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT FALSE AS condition", + "merge_columns": "id", + "condition_column": "condition", + }, + ) + + +def test_sql_query_merge_columns_invalid_entries_raise(): + """Ensure merge_columns entries must be non-empty strings.""" + with pytest.raises(InvalidParameterError, match="'merge_columns' entries must be non-empty strings."): + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT FALSE AS condition", + "merge_columns": ["id", ""], + "condition_column": "condition", + }, + ) + + @pytest.mark.parametrize( "lookback_windows, min_records_per_window, window_minutes, expected_message", [