From 5bb84801d4de27597f95c11d660db95a1042681d Mon Sep 17 00:00:00 2001 From: Varun Bhandary Date: Mon, 24 Nov 2025 18:11:41 +0000 Subject: [PATCH 01/18] feat: Make merge_columns optional in sql_query check Now supports two modes: - Row-level validation (with merge_columns): results joined back to specific rows - Dataset-level validation (without merge_columns): all rows get same result This makes aggregate validations potentialy much faster and works with custom_metrics. Empty list is treated as None. backward compatible. Includes tests, docs, and demo updates. --- demos/dqx_demo_library.py | 69 +++- docs/dqx/docs/reference/quality_checks.mdx | 2 +- src/databricks/labs/dqx/check_funcs.py | 67 ++-- test_sql_query_optional_merge_columns.py | 373 ++++++++++++++++++ tests/integration/test_apply_checks.py | 415 +++++++++++++++++++++ tests/resources/all_dataset_checks.yaml | 13 +- tests/unit/test_dataset_checks.py | 27 +- 7 files changed, 933 insertions(+), 33 deletions(-) create mode 100644 test_sql_query_optional_merge_columns.py diff --git a/demos/dqx_demo_library.py b/demos/dqx_demo_library.py index d414317cb..6b91982ec 100644 --- a/demos/dqx_demo_library.py +++ b/demos/dqx_demo_library.py @@ -722,11 +722,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 @@ -752,7 +756,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", @@ -769,6 +773,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( """ @@ -807,6 +846,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/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 4587eab77..3341b900d 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1392,7 +1392,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; `aggr_type`: aggregation function to use, such as "count" (default), "sum", "avg", "min", and "max"; `group_by`: (optional) list of columns or column expressions to group the rows for aggregation (no grouping by default) | | `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; `aggr_type`: aggregation function to use, such as "count" (default), "sum", "avg", "min", and "max"; `group_by`: (optional) list of columns or column expressions to group the rows for aggregation (no grouping by default) | | `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 | diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 5ab24df1f..7eba6ce5c 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1151,7 +1151,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, @@ -1164,11 +1164,14 @@ def sql_query( 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. 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. + merge_columns: Optional list of columns for join back to the input DataFrame. + When provided, they must provide a unique key for the join, otherwise duplicate records may be produced. + When not provided (None or empty list), the check result applies to all rows in the dataset. 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). @@ -1181,13 +1184,19 @@ 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}" @@ -1224,6 +1233,31 @@ 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: + # Query should only return the condition column + user_query_df = spark.sql(query_resolved).select(F.col(condition_column).alias(unique_condition_column)) + + # Get the first (and should be only) row's condition value + # If the query returns no rows, treat it as condition not met (False) + condition_result = user_query_df.first() + condition_value = condition_result[unique_condition_column] if condition_result else False + + # 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 + # 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) @@ -2680,27 +2714,18 @@ def _is_valid_ipv6_cidr_block(cidr: str) -> bool: return False -def _validate_sql_query_params(query: str, merge_columns: list[str]) -> None: +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 (not validated here as empty list is normalized to None). Raises: - MissingParameterError: If any required parameter is missing. - InvalidParameterError: If any parameter is invalid. UnsafeSqlQueryError: If the SQL query is unsafe. """ - 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." diff --git a/test_sql_query_optional_merge_columns.py b/test_sql_query_optional_merge_columns.py new file mode 100644 index 000000000..3cb554896 --- /dev/null +++ b/test_sql_query_optional_merge_columns.py @@ -0,0 +1,373 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # Test: sql_query with Optional merge_columns +# MAGIC +# MAGIC This notebook tests the new feature where `merge_columns` is now optional in the `sql_query` check function. +# MAGIC +# MAGIC **Features to test:** +# MAGIC 1. Dataset-level checks (merge_columns=None) - all rows get same result +# MAGIC 2. Dataset-level checks with row_filter - only filtered rows marked +# MAGIC 3. Empty list treated as None +# MAGIC 4. Row-level checks (merge_columns provided) - still work as before +# MAGIC 5. Performance comparison + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. Install the wheel file +# MAGIC +# MAGIC Upload the wheel file `databricks_labs_dqx-0.10.0-py3-none-any.whl` to DBFS or Workspace, then install it: + +# COMMAND ---------- + +# Install from workspace/DBFS (update path as needed) +# %pip install /Workspace/Users/your.email@domain.com/databricks_labs_dqx-0.10.0-py3-none-any.whl --force-reinstall +# Or if uploaded to DBFS: +# %pip install /dbfs/path/to/databricks_labs_dqx-0.10.0-py3-none-any.whl --force-reinstall + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Setup Test Data + +# COMMAND ---------- + +from pyspark.sql import SparkSession +from databricks.sdk import WorkspaceClient + +spark = SparkSession.builder.getOrCreate() +ws = WorkspaceClient() + +# Create test DataFrame +test_df = spark.createDataFrame([ + [1, 10, 100], + [2, 20, 200], + [3, 30, 300], + [4, 5, 50], + [5, 15, 150], +], ["id", "value", "amount"]) + +display(test_df) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. Test Dataset-Level Check (merge_columns=None) +# MAGIC +# MAGIC All rows should get the same check result (pass or fail together) + +# COMMAND ---------- + +from databricks.labs.dqx.engine import DQEngine +from databricks.labs.dqx.rule import DQDatasetRule +from databricks.labs.dqx.check_funcs import sql_query + +dq_engine = DQEngine(workspace_client=ws) + +# Dataset-level check: total count > 3 +# This should PASS (we have 5 rows), so no violations +checks_pass = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT COUNT(*) < 10 AS condition FROM {{input_view}}", + # No merge_columns = dataset-level check + "condition_column": "condition", + "msg": "Dataset has too many rows", + "name": "dataset_size_check_pass", + }, + ), +] + +result_pass = dq_engine.apply_checks(test_df, checks_pass) +print("āœ… Test 1: Dataset-level check that PASSES") +display(result_pass) +# Expected: All rows should have None for _errors and _warnings + +# COMMAND ---------- + +# Dataset-level check that FAILS +# This should FAIL (we have 5 rows > 3), so ALL rows get marked +checks_fail = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT COUNT(*) > 3 AS condition FROM {{input_view}}", + # No merge_columns = dataset-level check + "condition_column": "condition", + "msg": "Dataset has more than 3 rows", + "name": "dataset_size_check_fail", + }, + ), +] + +result_fail = dq_engine.apply_checks(test_df, checks_fail) +print("āŒ Test 2: Dataset-level check that FAILS") +display(result_fail) +# Expected: ALL rows should have the error + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 4. Test Dataset-Level Check with row_filter +# MAGIC +# MAGIC Only filtered rows should get the check result, others should have None + +# COMMAND ---------- + +# Dataset-level check with filter: check only rows where value >= 20 +# For filtered rows (value >= 20), check if SUM(amount) > 500 +# Filtered rows: id=2,3,5 with amounts 200,300,150 = 650 > 500, so condition is TRUE +checks_filter = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + filter="value >= 20", # Only apply to these rows + check_func_kwargs={ + "query": "SELECT SUM(amount) > 500 AS condition FROM {{input_view}}", + # No merge_columns = dataset-level check + "condition_column": "condition", + "msg": "High value rows have too much amount", + "name": "filtered_dataset_check", + }, + ), +] + +result_filter = dq_engine.apply_checks(test_df, checks_filter) +print("šŸ” Test 3: Dataset-level check with row_filter") +display(result_filter) +# Expected: Only rows where value >= 20 (id=2,3,5) should have errors +# Rows where value < 20 (id=1,4) should have None + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 5. Test Empty List Treated as None + +# COMMAND ---------- + +# Empty list should behave the same as None +checks_empty_list = [ + DQDatasetRule( + criticality="warn", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT COUNT(*) > 3 AS condition FROM {{input_view}}", + "merge_columns": [], # Empty list = same as None + "condition_column": "condition", + "msg": "Empty list test", + "name": "empty_list_check", + }, + ), +] + +result_empty = dq_engine.apply_checks(test_df, checks_empty_list) +print("šŸ“‹ Test 4: Empty list for merge_columns") +display(result_empty) +# Expected: Should work same as merge_columns=None (all rows get same result) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 6. Test Row-Level Check (Backward Compatibility) +# MAGIC +# MAGIC With merge_columns provided, should work as before + +# COMMAND ---------- + +# Row-level check: mark rows where value > 20 +checks_row_level = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT id, value > 20 AS condition FROM {{input_view}}", + "merge_columns": ["id"], # Join back by id + "condition_column": "condition", + "msg": "Value exceeds threshold", + "name": "row_level_check", + }, + ), +] + +result_row = dq_engine.apply_checks(test_df, checks_row_level) +print("šŸ“ Test 5: Row-level check (backward compatibility)") +display(result_row) +# Expected: Only rows where value > 20 (id=2,3) should have errors + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 7. Test with Metrics Observer +# MAGIC +# MAGIC Dataset-level checks work great with custom metrics! + +# COMMAND ---------- + +from databricks.labs.dqx.metrics_observer import DQMetricsObserver + +# Create observer with custom metrics +custom_metrics = [ + "avg(amount) as avg_amount", + "sum(amount) as total_amount", + "max(value) as max_value", +] + +observer = DQMetricsObserver(name="test_metrics", custom_metrics=custom_metrics) +dq_engine_with_metrics = DQEngine(workspace_client=ws, observer=observer) + +checks_with_metrics = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT COUNT(*) > 3 AS condition FROM {{input_view}}", + # No merge_columns = perfect for aggregate validations + "condition_column": "condition", + "msg": "Dataset validation with metrics", + "name": "dataset_with_metrics", + }, + ), +] + +result_metrics, observation = dq_engine_with_metrics.apply_checks(test_df, checks_with_metrics) + +# Trigger action to get metrics +result_metrics.count() + +print("šŸ“Š Test 6: Dataset-level check with custom metrics") +print("\nMetrics observed:") +print(observation.get) +display(result_metrics) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 8. Test with Negation + +# COMMAND ---------- + +# Dataset-level check with negate=True +# Query returns False (COUNT < 2), but negate=True means it should fail +checks_negate = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT COUNT(*) < 2 AS condition FROM {{input_view}}", + # No merge_columns + "condition_column": "condition", + "msg": "Dataset should have at least 2 rows (negated check)", + "name": "negated_check", + "negate": True, # Flip the logic + }, + ), +] + +result_negate = dq_engine.apply_checks(test_df, checks_negate) +print("šŸ”„ Test 7: Dataset-level check with negate=True") +display(result_negate) +# Expected: Should fail because COUNT >= 2 but negate flips it + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 9. Performance Comparison +# MAGIC +# MAGIC Compare dataset-level vs row-level performance + +# COMMAND ---------- + +import time + +# Create larger dataset for meaningful comparison +large_df = spark.range(0, 1000000).selectExpr( + "id", + "id % 100 as group_id", + "rand() * 1000 as value" +) + +print(f"Dataset size: {large_df.count()} rows") + +# COMMAND ---------- + +# Test 1: Dataset-level check (should be fast) +start = time.time() + +checks_dataset = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT AVG(value) > 500 AS condition FROM {{input_view}}", + # No merge_columns = fast! + "condition_column": "condition", + "name": "perf_dataset_check", + }, + ), +] + +result_dataset = dq_engine.apply_checks(large_df, checks_dataset) +result_dataset.count() # Trigger action + +dataset_time = time.time() - start +print(f"āš” Dataset-level check time: {dataset_time:.2f} seconds") + +# COMMAND ---------- + +# Test 2: Row-level check (will be slower due to join) +start = time.time() + +checks_row = [ + DQDatasetRule( + criticality="error", + check_func=sql_query, + check_func_kwargs={ + "query": "SELECT group_id, AVG(value) > 500 AS condition FROM {{input_view}} GROUP BY group_id", + "merge_columns": ["group_id"], # Requires join + "condition_column": "condition", + "name": "perf_row_check", + }, + ), +] + +result_row = dq_engine.apply_checks(large_df, checks_row) +result_row.count() # Trigger action + +row_time = time.time() - start +print(f"šŸ¢ Row-level check time: {row_time:.2f} seconds") + +# COMMAND ---------- + +print(f"\nšŸ“Š Performance Summary:") +print(f"Dataset-level: {dataset_time:.2f}s") +print(f"Row-level: {row_time:.2f}s") +print(f"Speedup: {row_time/dataset_time:.1f}x faster!") +print(f"\nāœ… Dataset-level checks are ~{row_time/dataset_time:.1f}x faster than row-level!") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## āœ… All Tests Complete! +# MAGIC +# MAGIC ### Summary of New Features: +# MAGIC 1. āœ… `merge_columns` is now optional (can be None or omitted) +# MAGIC 2. āœ… Empty list `[]` is treated the same as `None` +# MAGIC 3. āœ… Dataset-level checks apply result to all rows (or filtered rows) +# MAGIC 4. āœ… Works with `row_filter` - only filtered rows get marked +# MAGIC 5. āœ… Backward compatible - existing row-level checks still work +# MAGIC 6. āœ… Works great with metrics observer for aggregate validations +# MAGIC 7. āœ… Significantly faster than row-level checks (5-20x) +# MAGIC +# MAGIC ### Use Cases: +# MAGIC - **Dataset-level (no merge_columns)**: Aggregate validations, dataset metrics, all-or-nothing checks +# MAGIC - **Row-level (with merge_columns)**: Identify specific problematic rows, per-group validations + +# COMMAND ---------- + diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index f0c274bd6..a9106fd9a 100644 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -3064,6 +3064,421 @@ 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_and_ref_table(ws, spark): dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) diff --git a/tests/resources/all_dataset_checks.yaml b/tests/resources/all_dataset_checks.yaml index 4d5e523ef..8711e0b2c 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 records # optional + name: dataset_not_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..3fac67b3e 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 == 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 == sql_query def test_sql_query_unsafe(): From 43812c86f8ce34f6cd503f0db9bf426fe2ee006e Mon Sep 17 00:00:00 2001 From: Varun Bhandary Date: Mon, 24 Nov 2025 20:41:25 +0000 Subject: [PATCH 02/18] fmt & tests: sql check complexity and reuse dataset level check. Added multi-dataset integration test. --- src/databricks/labs/dqx/check_funcs.py | 101 ++++-- test_sql_query_optional_merge_columns.py | 373 ----------------------- tests/integration/test_apply_checks.py | 173 +++++++++-- tests/unit/test_dataset_checks.py | 4 +- 4 files changed, 226 insertions(+), 425 deletions(-) delete mode 100644 test_sql_query_optional_merge_columns.py diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 7eba6ce5c..1e84304df 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1162,16 +1162,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 (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: Optional list of columns for join back to the input DataFrame. - When provided, they must provide a unique key for the join, otherwise duplicate records may be produced. - When not provided (None or empty list), the check result applies to all rows in the dataset. 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). @@ -1187,15 +1195,19 @@ def sql_query( UnsafeSqlQueryError: if the SQL query fails the safety check (e.g., contains disallowed operations). """ _validate_sql_query_params(query, merge_columns) - + # 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" + 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 @@ -1235,28 +1247,9 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> # When merge_columns is None, the check applies to all rows (dataset-level check) if merge_columns is None: - # Query should only return the condition column - user_query_df = spark.sql(query_resolved).select(F.col(condition_column).alias(unique_condition_column)) - - # Get the first (and should be only) row's condition value - # If the query returns no rows, treat it as condition not met (False) - condition_result = user_query_df.first() - condition_value = condition_result[unique_condition_column] if condition_result else False - - # 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 + 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( @@ -2714,14 +2707,58 @@ def _is_valid_ipv6_cidr_block(cidr: str) -> bool: return False -def _validate_sql_query_params(query: str, merge_columns: list[str] | None) -> 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)) + + # Get the first (and should be only) row's condition value + # If the query returns no rows, treat it as condition not met (False) + condition_result = user_query_df.first() + condition_value = condition_result[unique_condition_column] if condition_result else False + + # 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 safe for execution. Args: query: The SQL query string to validate. - merge_columns: Optional list of column names (not validated here as empty list is normalized to None). Raises: UnsafeSqlQueryError: If the SQL query is unsafe. diff --git a/test_sql_query_optional_merge_columns.py b/test_sql_query_optional_merge_columns.py deleted file mode 100644 index 3cb554896..000000000 --- a/test_sql_query_optional_merge_columns.py +++ /dev/null @@ -1,373 +0,0 @@ -# Databricks notebook source -# MAGIC %md -# MAGIC # Test: sql_query with Optional merge_columns -# MAGIC -# MAGIC This notebook tests the new feature where `merge_columns` is now optional in the `sql_query` check function. -# MAGIC -# MAGIC **Features to test:** -# MAGIC 1. Dataset-level checks (merge_columns=None) - all rows get same result -# MAGIC 2. Dataset-level checks with row_filter - only filtered rows marked -# MAGIC 3. Empty list treated as None -# MAGIC 4. Row-level checks (merge_columns provided) - still work as before -# MAGIC 5. Performance comparison - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 1. Install the wheel file -# MAGIC -# MAGIC Upload the wheel file `databricks_labs_dqx-0.10.0-py3-none-any.whl` to DBFS or Workspace, then install it: - -# COMMAND ---------- - -# Install from workspace/DBFS (update path as needed) -# %pip install /Workspace/Users/your.email@domain.com/databricks_labs_dqx-0.10.0-py3-none-any.whl --force-reinstall -# Or if uploaded to DBFS: -# %pip install /dbfs/path/to/databricks_labs_dqx-0.10.0-py3-none-any.whl --force-reinstall - -# COMMAND ---------- - -dbutils.library.restartPython() - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 2. Setup Test Data - -# COMMAND ---------- - -from pyspark.sql import SparkSession -from databricks.sdk import WorkspaceClient - -spark = SparkSession.builder.getOrCreate() -ws = WorkspaceClient() - -# Create test DataFrame -test_df = spark.createDataFrame([ - [1, 10, 100], - [2, 20, 200], - [3, 30, 300], - [4, 5, 50], - [5, 15, 150], -], ["id", "value", "amount"]) - -display(test_df) - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 3. Test Dataset-Level Check (merge_columns=None) -# MAGIC -# MAGIC All rows should get the same check result (pass or fail together) - -# COMMAND ---------- - -from databricks.labs.dqx.engine import DQEngine -from databricks.labs.dqx.rule import DQDatasetRule -from databricks.labs.dqx.check_funcs import sql_query - -dq_engine = DQEngine(workspace_client=ws) - -# Dataset-level check: total count > 3 -# This should PASS (we have 5 rows), so no violations -checks_pass = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT COUNT(*) < 10 AS condition FROM {{input_view}}", - # No merge_columns = dataset-level check - "condition_column": "condition", - "msg": "Dataset has too many rows", - "name": "dataset_size_check_pass", - }, - ), -] - -result_pass = dq_engine.apply_checks(test_df, checks_pass) -print("āœ… Test 1: Dataset-level check that PASSES") -display(result_pass) -# Expected: All rows should have None for _errors and _warnings - -# COMMAND ---------- - -# Dataset-level check that FAILS -# This should FAIL (we have 5 rows > 3), so ALL rows get marked -checks_fail = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT COUNT(*) > 3 AS condition FROM {{input_view}}", - # No merge_columns = dataset-level check - "condition_column": "condition", - "msg": "Dataset has more than 3 rows", - "name": "dataset_size_check_fail", - }, - ), -] - -result_fail = dq_engine.apply_checks(test_df, checks_fail) -print("āŒ Test 2: Dataset-level check that FAILS") -display(result_fail) -# Expected: ALL rows should have the error - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 4. Test Dataset-Level Check with row_filter -# MAGIC -# MAGIC Only filtered rows should get the check result, others should have None - -# COMMAND ---------- - -# Dataset-level check with filter: check only rows where value >= 20 -# For filtered rows (value >= 20), check if SUM(amount) > 500 -# Filtered rows: id=2,3,5 with amounts 200,300,150 = 650 > 500, so condition is TRUE -checks_filter = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - filter="value >= 20", # Only apply to these rows - check_func_kwargs={ - "query": "SELECT SUM(amount) > 500 AS condition FROM {{input_view}}", - # No merge_columns = dataset-level check - "condition_column": "condition", - "msg": "High value rows have too much amount", - "name": "filtered_dataset_check", - }, - ), -] - -result_filter = dq_engine.apply_checks(test_df, checks_filter) -print("šŸ” Test 3: Dataset-level check with row_filter") -display(result_filter) -# Expected: Only rows where value >= 20 (id=2,3,5) should have errors -# Rows where value < 20 (id=1,4) should have None - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 5. Test Empty List Treated as None - -# COMMAND ---------- - -# Empty list should behave the same as None -checks_empty_list = [ - DQDatasetRule( - criticality="warn", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT COUNT(*) > 3 AS condition FROM {{input_view}}", - "merge_columns": [], # Empty list = same as None - "condition_column": "condition", - "msg": "Empty list test", - "name": "empty_list_check", - }, - ), -] - -result_empty = dq_engine.apply_checks(test_df, checks_empty_list) -print("šŸ“‹ Test 4: Empty list for merge_columns") -display(result_empty) -# Expected: Should work same as merge_columns=None (all rows get same result) - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 6. Test Row-Level Check (Backward Compatibility) -# MAGIC -# MAGIC With merge_columns provided, should work as before - -# COMMAND ---------- - -# Row-level check: mark rows where value > 20 -checks_row_level = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT id, value > 20 AS condition FROM {{input_view}}", - "merge_columns": ["id"], # Join back by id - "condition_column": "condition", - "msg": "Value exceeds threshold", - "name": "row_level_check", - }, - ), -] - -result_row = dq_engine.apply_checks(test_df, checks_row_level) -print("šŸ“ Test 5: Row-level check (backward compatibility)") -display(result_row) -# Expected: Only rows where value > 20 (id=2,3) should have errors - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 7. Test with Metrics Observer -# MAGIC -# MAGIC Dataset-level checks work great with custom metrics! - -# COMMAND ---------- - -from databricks.labs.dqx.metrics_observer import DQMetricsObserver - -# Create observer with custom metrics -custom_metrics = [ - "avg(amount) as avg_amount", - "sum(amount) as total_amount", - "max(value) as max_value", -] - -observer = DQMetricsObserver(name="test_metrics", custom_metrics=custom_metrics) -dq_engine_with_metrics = DQEngine(workspace_client=ws, observer=observer) - -checks_with_metrics = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT COUNT(*) > 3 AS condition FROM {{input_view}}", - # No merge_columns = perfect for aggregate validations - "condition_column": "condition", - "msg": "Dataset validation with metrics", - "name": "dataset_with_metrics", - }, - ), -] - -result_metrics, observation = dq_engine_with_metrics.apply_checks(test_df, checks_with_metrics) - -# Trigger action to get metrics -result_metrics.count() - -print("šŸ“Š Test 6: Dataset-level check with custom metrics") -print("\nMetrics observed:") -print(observation.get) -display(result_metrics) - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 8. Test with Negation - -# COMMAND ---------- - -# Dataset-level check with negate=True -# Query returns False (COUNT < 2), but negate=True means it should fail -checks_negate = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT COUNT(*) < 2 AS condition FROM {{input_view}}", - # No merge_columns - "condition_column": "condition", - "msg": "Dataset should have at least 2 rows (negated check)", - "name": "negated_check", - "negate": True, # Flip the logic - }, - ), -] - -result_negate = dq_engine.apply_checks(test_df, checks_negate) -print("šŸ”„ Test 7: Dataset-level check with negate=True") -display(result_negate) -# Expected: Should fail because COUNT >= 2 but negate flips it - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## 9. Performance Comparison -# MAGIC -# MAGIC Compare dataset-level vs row-level performance - -# COMMAND ---------- - -import time - -# Create larger dataset for meaningful comparison -large_df = spark.range(0, 1000000).selectExpr( - "id", - "id % 100 as group_id", - "rand() * 1000 as value" -) - -print(f"Dataset size: {large_df.count()} rows") - -# COMMAND ---------- - -# Test 1: Dataset-level check (should be fast) -start = time.time() - -checks_dataset = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT AVG(value) > 500 AS condition FROM {{input_view}}", - # No merge_columns = fast! - "condition_column": "condition", - "name": "perf_dataset_check", - }, - ), -] - -result_dataset = dq_engine.apply_checks(large_df, checks_dataset) -result_dataset.count() # Trigger action - -dataset_time = time.time() - start -print(f"āš” Dataset-level check time: {dataset_time:.2f} seconds") - -# COMMAND ---------- - -# Test 2: Row-level check (will be slower due to join) -start = time.time() - -checks_row = [ - DQDatasetRule( - criticality="error", - check_func=sql_query, - check_func_kwargs={ - "query": "SELECT group_id, AVG(value) > 500 AS condition FROM {{input_view}} GROUP BY group_id", - "merge_columns": ["group_id"], # Requires join - "condition_column": "condition", - "name": "perf_row_check", - }, - ), -] - -result_row = dq_engine.apply_checks(large_df, checks_row) -result_row.count() # Trigger action - -row_time = time.time() - start -print(f"šŸ¢ Row-level check time: {row_time:.2f} seconds") - -# COMMAND ---------- - -print(f"\nšŸ“Š Performance Summary:") -print(f"Dataset-level: {dataset_time:.2f}s") -print(f"Row-level: {row_time:.2f}s") -print(f"Speedup: {row_time/dataset_time:.1f}x faster!") -print(f"\nāœ… Dataset-level checks are ~{row_time/dataset_time:.1f}x faster than row-level!") - -# COMMAND ---------- - -# MAGIC %md -# MAGIC ## āœ… All Tests Complete! -# MAGIC -# MAGIC ### Summary of New Features: -# MAGIC 1. āœ… `merge_columns` is now optional (can be None or omitted) -# MAGIC 2. āœ… Empty list `[]` is treated the same as `None` -# MAGIC 3. āœ… Dataset-level checks apply result to all rows (or filtered rows) -# MAGIC 4. āœ… Works with `row_filter` - only filtered rows get marked -# MAGIC 5. āœ… Backward compatible - existing row-level checks still work -# MAGIC 6. āœ… Works great with metrics observer for aggregate validations -# MAGIC 7. āœ… Significantly faster than row-level checks (5-20x) -# MAGIC -# MAGIC ### Use Cases: -# MAGIC - **Dataset-level (no merge_columns)**: Aggregate validations, dataset metrics, all-or-nothing checks -# MAGIC - **Row-level (with merge_columns)**: Identify specific problematic rows, per-group validations - -# COMMAND ---------- - diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index a9106fd9a..67a4e72aa 100644 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -3073,7 +3073,7 @@ def test_apply_checks_with_sql_query_without_merge_columns(ws, spark): # 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", @@ -3086,9 +3086,9 @@ def test_apply_checks_with_sql_query_without_merge_columns(ws, spark): }, ), ] - + 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( [ @@ -3179,7 +3179,7 @@ def test_apply_checks_with_sql_query_without_merge_columns_and_filter(ws, spark) # 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", @@ -3193,9 +3193,9 @@ def test_apply_checks_with_sql_query_without_merge_columns_and_filter(ws, spark) }, ), ] - + 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( @@ -3263,7 +3263,7 @@ def test_apply_checks_with_sql_query_without_merge_columns_passes(ws, spark): # 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", @@ -3276,9 +3276,9 @@ def test_apply_checks_with_sql_query_without_merge_columns_passes(ws, spark): }, ), ] - + checked = dq_engine.apply_checks(test_df, checks) - + # All rows should have None for errors since condition is False expected = spark.createDataFrame( [ @@ -3297,7 +3297,7 @@ def test_apply_checks_with_sql_query_without_merge_columns_empty_result(ws, spar # 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", @@ -3310,9 +3310,9 @@ def test_apply_checks_with_sql_query_without_merge_columns_empty_result(ws, spar }, ), ] - + checked = dq_engine.apply_checks(test_df, checks) - + # No rows from query means condition is False, so no violations expected = spark.createDataFrame( [ @@ -3331,7 +3331,7 @@ def test_apply_checks_with_sql_query_without_merge_columns_negate(ws, spark): # Query that returns False, but with negate=True should fail query_false = "SELECT COUNT(*) > 100 AS condition FROM {{input_view}}" - + checks = [ DQDatasetRule( criticality="error", @@ -3345,9 +3345,9 @@ def test_apply_checks_with_sql_query_without_merge_columns_negate(ws, spark): }, ), ] - + checked = dq_engine.apply_checks(test_df, checks) - + # With negate=True, False becomes violation expected = spark.createDataFrame( [ @@ -3400,7 +3400,7 @@ def test_apply_checks_with_sql_query_without_merge_columns_warning(ws, spark): # 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 @@ -3413,9 +3413,9 @@ def test_apply_checks_with_sql_query_without_merge_columns_warning(ws, spark): }, ), ] - + checked = dq_engine.apply_checks(test_df, checks) - + # All rows should have warning (not error) expected = spark.createDataFrame( [ @@ -3479,6 +3479,143 @@ def test_apply_checks_with_sql_query_without_merge_columns_warning(ws, spark): 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: verify total amount matches expected value in 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: verify total amount matches expected value in 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_and_ref_table(ws, spark): dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) diff --git a/tests/unit/test_dataset_checks.py b/tests/unit/test_dataset_checks.py index 3fac67b3e..d7a47cf9c 100644 --- a/tests/unit/test_dataset_checks.py +++ b/tests/unit/test_dataset_checks.py @@ -153,7 +153,7 @@ def test_sql_query_empty_merge_columns(): check_func=sql_query, check_func_kwargs={"query": "SELECT FALSE AS condition", "merge_columns": [], "condition_column": "condition"}, ) - assert rule.check_func == sql_query + assert rule.check_func is sql_query def test_sql_query_without_merge_columns(): @@ -164,7 +164,7 @@ def test_sql_query_without_merge_columns(): check_func=sql_query, check_func_kwargs={"query": "SELECT FALSE AS condition", "condition_column": "condition"}, ) - assert rule.check_func == sql_query + assert rule.check_func is sql_query def test_sql_query_unsafe(): From 24abf0de748cf8c0c010635ce45a7691ffb20a55 Mon Sep 17 00:00:00 2001 From: Varun Bhandary Date: Tue, 25 Nov 2025 19:11:03 +0000 Subject: [PATCH 03/18] improve: tighten sql_query validation + share test helper --- src/databricks/labs/dqx/check_funcs.py | 34 +++++-- tests/integration/conftest.py | 60 +++++------ tests/integration/test_apply_checks.py | 132 +++++++------------------ tests/unit/test_dataset_checks.py | 30 ++++++ 4 files changed, 123 insertions(+), 133 deletions(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 1e84304df..dd8a98697 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 @@ -2732,10 +2732,18 @@ def _apply_dataset_level_sql_check( # Query should only return the condition column user_query_df = spark.sql(query_resolved).select(F.col(condition_column).alias(unique_condition_column)) - # Get the first (and should be only) row's condition value - # If the query returns no rows, treat it as condition not met (False) - condition_result = user_query_df.first() - condition_value = condition_result[unique_condition_column] if condition_result else False + # 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 @@ -2752,18 +2760,32 @@ def _apply_dataset_level_sql_check( return result_df -def _validate_sql_query_params(query: str, _merge_columns: list[str] | None) -> None: +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 safe for execution. Args: query: The SQL query string to validate. + merge_columns: Optional list of column names (validated when provided). Raises: UnsafeSqlQueryError: If the SQL query is unsafe. + InvalidParameterError: If merge_columns is provided but not a sequence of strings. """ 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 isinstance(merge_columns, str) or not isinstance(merge_columns, Sequence): + 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.") diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 6400415a5..a6db3ea2b 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], + *, + 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: @@ -298,16 +319,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, ], @@ -315,16 +329,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, ], @@ -332,16 +339,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 67a4e72aa..8b470cc12 100644 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -23,7 +23,7 @@ from databricks.labs.dqx.schema import dq_result_schema from databricks.labs.dqx import check_funcs import databricks.labs.dqx.geo.check_funcs as geo_check_funcs -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 from tests.conftest import TEST_CATALOG @@ -3616,6 +3616,32 @@ def test_apply_checks_with_sql_query_without_merge_columns_and_ref_df_fail(ws, s 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) @@ -4432,18 +4458,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, ], [ @@ -4451,47 +4466,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, @@ -4526,18 +4508,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, ], [ @@ -4545,47 +4516,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, diff --git a/tests/unit/test_dataset_checks.py b/tests/unit/test_dataset_checks.py index d7a47cf9c..871d37ded 100644 --- a/tests/unit/test_dataset_checks.py +++ b/tests/unit/test_dataset_checks.py @@ -177,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", [ From 55383eabf77fd3a763f00875e96e3db482026957 Mon Sep 17 00:00:00 2001 From: Varun Bhandary Date: Tue, 25 Nov 2025 19:34:09 +0000 Subject: [PATCH 04/18] fix: test assertion when totals do not match --- tests/integration/test_apply_checks.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 8b470cc12..4a76b1bd2 100644 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -3489,9 +3489,9 @@ def test_apply_checks_with_sql_query_without_merge_columns_and_ref_df(ws, spark) # Reference dataset with expected totals ref_df = spark.createDataFrame([[600]], "expected_total: int") - # Dataset-level check: verify total amount matches expected value in reference table + # 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 + SELECT (SELECT SUM(c) FROM {{input_view}}) != (SELECT expected_total FROM {{expected_totals}}) AS condition """ checks = [ @@ -3532,9 +3532,9 @@ def test_apply_checks_with_sql_query_without_merge_columns_and_ref_df_fail(ws, s # Reference dataset with expected totals (wrong value) ref_df = spark.createDataFrame([[999]], "expected_total: int") - # Dataset-level check: verify total amount matches expected value in reference table + # 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 + SELECT (SELECT SUM(c) FROM {{input_view}}) != (SELECT expected_total FROM {{expected_totals}}) AS condition """ checks = [ From 2f43afadbbbf4f74c59e1624f71d8c74b138550b Mon Sep 17 00:00:00 2001 From: Greg Hansen <163584195+ghanse@users.noreply.github.com> Date: Mon, 24 Nov 2025 22:29:21 +0100 Subject: [PATCH 05/18] Support Custom Folder Installation for CLI Commands (#942) ## Changes This PR adds `--install-folder` arguments for the DQX CLI commands. This allows users to run the CLI commands when DQX is installed in a custom folder. When users install DQX in a custom folder, they should add arguments to run DQX CLI commands: ``` databricks labs dqx open-dashboards --install-folder "/Workspace/..." ``` ### Linked issues Resolves #671 ### Tests - [x] manually tested - [ ] added unit tests - [x] added integration tests - [ ] added end-to-end tests - [ ] added performance tests --- docs/dqx/docs/reference/cli.mdx | 28 +++- labs.yml | 24 ++- pyproject.toml | 2 +- src/databricks/labs/dqx/cli.py | 38 +++-- .../labs/dqx/contexts/global_context.py | 7 +- .../labs/dqx/contexts/workspace_context.py | 6 +- src/databricks/labs/dqx/installer/install.py | 2 +- tests/integration/conftest.py | 3 + tests/integration/test_cli.py | 155 ++++++++++++++++++ 9 files changed, 246 insertions(+), 19 deletions(-) diff --git a/docs/dqx/docs/reference/cli.mdx b/docs/dqx/docs/reference/cli.mdx index 65a6cdede..6faa46d7c 100644 --- a/docs/dqx/docs/reference/cli.mdx +++ b/docs/dqx/docs/reference/cli.mdx @@ -26,7 +26,9 @@ databricks labs uninstall dqx By default, DQX is installed under the user's home folder (for example `/Users//.dqx`). Use the `DQX_FORCE_INSTALL` env var to force a global or user install. Provide a custom installation -folder during installation to override the default location. See the installation guide for details. +folder during installation to override the default location. See the installation guide for details. +Set the `--install-folder` argument When running commands for a DQX installation with a custom +installation folder. ## Configuration helpers @@ -35,8 +37,14 @@ folder during installation to override the default location. See the installatio # Open the installation config in the workspace databricks labs dqx open-remote-config +# Open the installation config in the workspace with a custom installation folder +databricks labs dqx open-remote-config --install-folder "/Workspace/my_dqx_folder" + # Open dashboards folder for the installation databricks labs dqx open-dashboards + +# Open dashboards folder for an installation with a custom installation folder +databricks labs dqx open-dashboards --install-folder "/Workspace/my_dqx_folder" ``` ## Workflows and logs @@ -45,8 +53,14 @@ databricks labs dqx open-dashboards # List installed workflows with their latest run state databricks labs dqx workflows +# List installed workflows with their latest run state with a custom installation folder +databricks labs dqx workflows --install-folder "/Workspace/my_dqx_folder" + # Show logs for the latest run of a workflow (profiler, quality-checker, e2e) databricks labs dqx logs --workflow quality-checker + +# Show logs for the latest run of a workflow (profiler, quality-checker, e2e) with a custom installation folder +databricks labs dqx logs --workflow quality-checker --install-folder "/Workspace/my_dqx_folder" ``` ## Running workflows @@ -55,6 +69,9 @@ databricks labs dqx logs --workflow quality-checker # Run profiler for all run configs in the configuration file: profile data and generate quality check candidates databricks labs dqx profile --timeout-minutes 20 +# Run profiler for all run configs in the configuration file with a custom installation folder: profile data and generate quality check candidates +databricks labs dqx profile --timeout-minutes 20 --install-folder "/Workspace/my_dqx_folder" + # Run profiler for a single run config in the configuration file: profile data and generate quality check candidates databricks labs dqx profile --run-config default --timeout-minutes 20 @@ -72,6 +89,9 @@ databricks labs dqx profile --run-config "default" --patterns "main.product001.* # Run quality checker for all run configs in the configuration file: apply quality checks and write results and optionally quarantine databricks labs dqx apply-checks --timeout-minutes 20 +# Run quality checker for all run configs in the configuration file with a custom installation folder: apply quality checks and write results and optionally quarantine +databricks labs dqx apply-checks --timeout-minutes 20 --install-folder "/Workspace/my_dqx_folder" + # Run quality checker for a single run config in the configuration file: apply quality checks and write results and optionally quarantine databricks labs dqx apply-checks --run-config default --timeout-minutes 20 @@ -81,6 +101,9 @@ databricks labs dqx apply-checks --run-config "default" --patterns "main.product # Run e2e (end-to-end) workflows for all run configs in the configuration file: profile data > generate quality checks > apply checks > write results and optionally quarantine databricks labs dqx e2e --timeout-minutes 20 +# Run e2e (end-to-end) workflows for all run configs in the configuration file with a custom installation folder: profile data > generate quality checks > apply checks > write results and optionally quarantine +databricks labs dqx e2e --timeout-minutes 20 --install-folder "/Workspace/my_dqx_folder" + # Run e2e (end-to-end) workflows for a single run config in the configuration file: profile data > generate quality checks > apply checks > write results and optionally quarantine databricks labs dqx e2e --run-config default --timeout-minutes 20 @@ -94,6 +117,9 @@ databricks labs dqx e2e --run-config "default" --patterns "main.product001.*;mai # Validate checks stored in the installation (defined in `checks_location`) for a single run config in the configuration file: databricks labs dqx validate-checks +# Validate checks stored in a custom installation folder (defined in `checks_location`) for a single run config in the configuration file: +databricks labs dqx validate-checks --install-folder "/Workspace/my_dqx_folder" + # Validate checks stored in the installation (defined in `checks_location`) for all run configs in the configuration file: databricks labs dqx validate-checks --run-config default ``` diff --git a/labs.yml b/labs.yml index 857e0c159..db69b0f13 100644 --- a/labs.yml +++ b/labs.yml @@ -11,14 +11,23 @@ min_python: 3.10 commands: - name: open-remote-config description: Opens remote configuration in the browser. + flags: + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. - name: open-dashboards description: Opens remote dashboards directory. + flags: + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. - name: installations description: Show installations by different users on the same workspace. table_template: |- Path\tVersion\Input {{range .}}{{.path}}\t{{.version}}\t{{.input}} {{end}} + flags: + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. - name: validate-checks description: Validate checks flags: @@ -27,6 +36,8 @@ commands: - name: validate-custom-check-functions description: Whether to validate custom check functions. default: true + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. table_template: |- Run Config\tErrors {{range .}}{{.run_config}}\t{{.error}} @@ -42,6 +53,8 @@ commands: description: (Optional) Semicolon separated list of location patterns to exclude from profiling. Useful if wanting to exclude existing output tables. - name: timeout-minutes description: The timeout in minutes for the cli to wait for the workflow to complete. + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. - name: apply-checks description: Apply data quality checks to the input data, and save the results to the output table and optionally quarantine table (based on the run config). flags: @@ -59,6 +72,8 @@ commands: default: "_dq_quarantine" - name: timeout-minutes description: The timeout in minutes for the cli to wait for the workflow to complete. + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. - name: e2e description: Run end to end workflow to profile input data and apply quality checks, and save the results to the output table and optionally quarantine table (based on the run config). flags: @@ -76,14 +91,21 @@ commands: default: "_dq_quarantine" - name: timeout-minutes description: The timeout in minutes for the cli to wait for the workflow to complete. + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. - name: workflows description: Show deployed workflows and their latest run state. table_template: |- Workflow\tWorkflow ID\tState\tStarted {{range .}}{{.workflow}}\t{{.workflow_id}}\t{{.state}}\t{{.started}} {{end}} + flags: + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. - name: logs description: Show logs from the latest job run flags: - name: workflow - description: Name of the workflow to show logs for, e.g. profiler. \ No newline at end of file + description: Name of the workflow to show logs for, e.g. profiler. + - name: install-folder + description: (Optional) Custom installation folder. If not provided, the default installation folder is used. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index cbf530962..8e79ca8f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "databricks-labs-blueprint>=0.9.1,<0.10", "databricks-sdk~=0.71", "databricks-labs-lsql>=0.5,<=0.16", - "sqlalchemy>=1.4,<3.0", + "sqlalchemy>=2.0,<3.0", ] [project.optional-dependencies] diff --git a/src/databricks/labs/dqx/cli.py b/src/databricks/labs/dqx/cli.py index bf8421d23..ca1ed9574 100644 --- a/src/databricks/labs/dqx/cli.py +++ b/src/databricks/labs/dqx/cli.py @@ -23,29 +23,31 @@ @dqx.command -def open_remote_config(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None): +def open_remote_config(w: WorkspaceClient, *, install_folder: str = "", ctx: WorkspaceContext | None = None): """ Opens remote configuration in the browser. Args: w: The WorkspaceClient instance to use for accessing the workspace. + install_folder: Optional custom installation folder path. ctx: The WorkspaceContext instance to use for accessing the workspace. """ - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) workspace_link = ctx.installation.workspace_link(WorkspaceConfig.__file__) webbrowser.open(workspace_link) @dqx.command -def open_dashboards(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None): +def open_dashboards(w: WorkspaceClient, *, install_folder: str = "", ctx: WorkspaceContext | None = None): """ Opens remote dashboard directory in the browser. Args: w: The WorkspaceClient instance to use for accessing the workspace. + install_folder: Optional custom installation folder path. ctx: The WorkspaceContext instance to use for accessing the workspace. """ - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) workspace_link = ctx.installation.workspace_link("") webbrowser.open(f"{workspace_link}dashboards/") @@ -85,6 +87,7 @@ def validate_checks( *, run_config: str = "", validate_custom_check_functions: bool = True, + install_folder: str = "", ctx: WorkspaceContext | None = None, ) -> list[dict]: """ @@ -94,9 +97,10 @@ def validate_checks( w: The WorkspaceClient instance to use for accessing the workspace. run_config: The name of the run configuration to use. If not provided, run it for all run configs. validate_custom_check_functions: Whether to validate custom check functions (default is True). + install_folder: Optional custom installation folder path. ctx: The WorkspaceContext instance to use for accessing the workspace. """ - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) config = ctx.installation.load(WorkspaceConfig) errors_list = [] @@ -154,6 +158,7 @@ def profile( patterns: str = "", exclude_patterns: str = "", timeout_minutes: int = 30, + install_folder: str = "", ctx: WorkspaceContext | None = None, ) -> None: """ @@ -168,10 +173,11 @@ def profile( exclude_patterns: Semicolon-separated list of location patterns to exclude. Useful to skip existing output and quarantine tables based on suffixes. timeout_minutes: The timeout for the workflow run in minutes (default is 30). + install_folder: Optional custom installation folder path. ctx: The WorkspaceContext instance to use for accessing the workspace. """ timeout = timedelta(minutes=timeout_minutes) - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) ctx.deployed_workflows.run_workflow( workflow="profiler", run_config_name=run_config, @@ -191,6 +197,7 @@ def apply_checks( output_table_suffix: str = "_dq_output", quarantine_table_suffix: str = "_dq_quarantine", timeout_minutes: int = 30, + install_folder: str = "", ctx: WorkspaceContext | None = None, ) -> None: """ @@ -207,10 +214,11 @@ def apply_checks( output_table_suffix: Suffix to append to the output table names (default is "_dq_output"). quarantine_table_suffix: Suffix to append to the quarantine table names (default is "_dq_quarantine"). timeout_minutes: The timeout for the workflow run in minutes (default is 30). + install_folder: Optional custom installation folder path. ctx: The WorkspaceContext instance to use for accessing the workspace. """ timeout = timedelta(minutes=timeout_minutes) - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) ctx.deployed_workflows.run_workflow( workflow="quality-checker", run_config_name=run_config, @@ -232,6 +240,7 @@ def e2e( output_table_suffix: str = "_dq_output", quarantine_table_suffix: str = "_dq_quarantine", timeout_minutes: int = 60, + install_folder: str = "", ctx: WorkspaceContext | None = None, ) -> None: """ @@ -251,10 +260,11 @@ def e2e( output_table_suffix: Suffix to append to the output table names (default is "_dq_output"). quarantine_table_suffix: Suffix to append to the quarantine table names (default is "_dq_quarantine"). timeout_minutes: The timeout for the workflow run in minutes (default is 60). + install_folder: Optional custom installation folder path. ctx: The WorkspaceContext instance to use for accessing the workspace. """ timeout = timedelta(minutes=timeout_minutes) - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) ctx.deployed_workflows.run_workflow( workflow="e2e", run_config_name=run_config, @@ -267,15 +277,16 @@ def e2e( @dqx.command -def workflows(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None): +def workflows(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None, install_folder: str = ""): """ Show deployed workflows and their state Args: w: The WorkspaceClient instance to use for accessing the workspace. ctx: The WorkspaceContext instance to use for accessing the workspace. + install_folder: Optional custom installation folder path. """ - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) logger.info("Fetching deployed jobs...") latest_job_status = ctx.deployed_workflows.latest_job_status() print(json.dumps(latest_job_status)) @@ -283,16 +294,19 @@ def workflows(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None): @dqx.command -def logs(w: WorkspaceClient, *, workflow: str | None = None, ctx: WorkspaceContext | None = None): +def logs( + w: WorkspaceClient, *, workflow: str | None = None, install_folder: str = "", ctx: WorkspaceContext | None = None +): """ Show logs of the latest job run. Args: w: The WorkspaceClient instance to use for accessing the workspace. workflow: The name of the workflow to show logs for. + install_folder: Optional custom installation folder path. ctx: The WorkspaceContext instance to use for accessing the workspace """ - ctx = ctx or WorkspaceContext(w) + ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None) ctx.deployed_workflows.relay_logs(workflow) diff --git a/src/databricks/labs/dqx/contexts/global_context.py b/src/databricks/labs/dqx/contexts/global_context.py index d0509e07f..b1fb9f06e 100644 --- a/src/databricks/labs/dqx/contexts/global_context.py +++ b/src/databricks/labs/dqx/contexts/global_context.py @@ -15,10 +15,11 @@ class GlobalContext(abc.ABC): GlobalContext class that provides a global context, including workspace client, """ - def __init__(self, named_parameters: dict[str, str] | None = None): + def __init__(self, named_parameters: dict[str, str] | None = None, install_folder: str | None = None): if not named_parameters: named_parameters = {} self._named_parameters = named_parameters + self._install_folder = install_folder def replace(self, **kwargs): """ @@ -48,6 +49,10 @@ def product_info(self) -> ProductInfo: @cached_property def installation(self) -> Installation: + if self._install_folder: + return Installation( + self.workspace_client, self.product_info.product_name(), install_folder=self._install_folder + ) return Installation.current(self.workspace_client, self.product_info.product_name()) @cached_property diff --git a/src/databricks/labs/dqx/contexts/workspace_context.py b/src/databricks/labs/dqx/contexts/workspace_context.py index bb97ffcfe..68c251beb 100644 --- a/src/databricks/labs/dqx/contexts/workspace_context.py +++ b/src/databricks/labs/dqx/contexts/workspace_context.py @@ -10,8 +10,10 @@ class WorkspaceContext(CliContext): WorkspaceContext class that extends CliContext to provide workspace-specific functionality. """ - def __init__(self, ws: WorkspaceClient, named_parameters: dict[str, str] | None = None): - super().__init__(named_parameters) + def __init__( + self, ws: WorkspaceClient, named_parameters: dict[str, str] | None = None, install_folder: str | None = None + ): + super().__init__(named_parameters, install_folder) self._ws = ws @cached_property diff --git a/src/databricks/labs/dqx/installer/install.py b/src/databricks/labs/dqx/installer/install.py index b76f88d43..4a05ac9e0 100644 --- a/src/databricks/labs/dqx/installer/install.py +++ b/src/databricks/labs/dqx/installer/install.py @@ -50,7 +50,7 @@ class WorkspaceInstaller(WorkspaceContext, InstallationMixin): """ def __init__(self, ws: WorkspaceClient, environ: dict[str, str] | None = None, install_folder: str | None = None): - super().__init__(ws) + super().__init__(ws, install_folder=install_folder) if not environ: environ = dict(os.environ.items()) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index a6db3ea2b..e06813da4 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -204,6 +204,9 @@ def setup_workflows_with_custom_folder( Set up the workflows with installation in the custom install folder. """ + if os.getenv("DATABRICKS_SERVERLESS_COMPUTE_ID"): + pytest.skip() + def create(_spark, **kwargs): installation_ctx_custom_install_folder.installation_service.run() diff --git a/tests/integration/test_cli.py b/tests/integration/test_cli.py index 688ba2a84..2601e1124 100644 --- a/tests/integration/test_cli.py +++ b/tests/integration/test_cli.py @@ -306,3 +306,158 @@ def test_logs(ws, installation_ctx, caplog): logs(installation_ctx.workspace_client, ctx=installation_ctx.workspace_installer) assert "No jobs to relay logs for" in caplog.text + + +def test_open_remote_config_with_custom_folder(ws, installation_ctx_custom_install_folder, webbrowser_open): + installation_ctx_custom_install_folder.installation.save(installation_ctx_custom_install_folder.config) + custom_folder = installation_ctx_custom_install_folder.installation.install_folder() + open_remote_config(w=installation_ctx_custom_install_folder.workspace_client, install_folder=custom_folder) + webbrowser_open.assert_called_once_with( + installation_ctx_custom_install_folder.installation.workspace_link(WorkspaceConfig.__file__) + ) + + +def test_open_dashboards_with_custom_folder(ws, installation_ctx_custom_install_folder, webbrowser_open): + installation_ctx_custom_install_folder.installation.save(installation_ctx_custom_install_folder.config) + custom_folder = installation_ctx_custom_install_folder.installation.install_folder() + open_dashboards(w=installation_ctx_custom_install_folder.workspace_client, install_folder=custom_folder) + webbrowser_open.assert_called_once_with( + installation_ctx_custom_install_folder.installation.workspace_link("") + "dashboards/" + ) + + +def test_validate_checks_with_custom_folder(ws, make_workspace_file, installation_ctx_custom_install_folder): + installation_ctx_custom_install_folder.installation.save(installation_ctx_custom_install_folder.config) + custom_folder = installation_ctx_custom_install_folder.installation.install_folder() + + checks = [{"criticality": "warn", "check": {"function": "is_not_null", "arguments": {"column": "a"}}}] + run_config = installation_ctx_custom_install_folder.config.get_run_config() + checks_location = f"{custom_folder}/{run_config.checks_location}" + make_workspace_file(path=checks_location, content=yaml.dump(checks)) + + errors_list = validate_checks( + installation_ctx_custom_install_folder.workspace_client, + run_config=run_config.name, + install_folder=custom_folder, + ) + + assert not errors_list + + +def test_validate_checks_with_custom_folder_invalid_checks( + ws, make_workspace_file, installation_ctx_custom_install_folder +): + installation_ctx_custom_install_folder.installation.save(installation_ctx_custom_install_folder.config) + custom_folder = installation_ctx_custom_install_folder.installation.install_folder() + + checks = [ + {"criticality": "warn", "check": {"function": "invalid_func", "arguments": {"column": "a"}}}, + {"criticality": "warn", "check_missing": {"function": "is_not_null", "arguments": {"column": "b"}}}, + ] + run_config = installation_ctx_custom_install_folder.config.get_run_config() + checks_location = f"{custom_folder}/{run_config.checks_location}" + make_workspace_file(path=checks_location, content=yaml.dump(checks)) + + errors = validate_checks(installation_ctx_custom_install_folder.workspace_client, install_folder=custom_folder) + + expected_errors = [ + "function 'invalid_func' is not defined", + "'check' field is missing", + ] + assert len(errors) == len(expected_errors) + for e in expected_errors: + assert any(e in error["error"] for error in errors) + + +def test_profile_with_custom_folder(ws, spark, setup_workflows_with_custom_folder, caplog): + installation_ctx, run_config = setup_workflows_with_custom_folder() + custom_folder = installation_ctx.installation.install_folder() + + profile(installation_ctx.workspace_client, run_config=run_config.name, install_folder=custom_folder) + + checks = DQEngine(ws, spark).load_checks( + config=InstallationChecksStorageConfig( + run_config_name=run_config.name, + assume_user=True, + product_name=installation_ctx.installation.product(), + install_folder=custom_folder, + ), + ) + assert checks, "Checks were not loaded correctly" + + status = ws.workspace.get_status(f"{custom_folder}/{run_config.profiler_config.summary_stats_file}") + assert status, f"Profile summary stats file {run_config.profiler_config.summary_stats_file} does not exist." + + with caplog.at_level(logging.INFO): + logs(installation_ctx.workspace_client, install_folder=custom_folder) + + assert "Completed profiler workflow run" in caplog.text + + +def test_apply_checks_with_custom_folder( + ws, spark, setup_workflows_with_custom_folder, caplog, expected_quality_checking_output +): + installation_ctx, run_config = setup_workflows_with_custom_folder(checks=True) + custom_folder = installation_ctx.installation.install_folder() + + apply_checks(installation_ctx.workspace_client, run_config=run_config.name, install_folder=custom_folder) + + checked_df = spark.table(run_config.output_config.location) + assert_df_equality(checked_df, expected_quality_checking_output, ignore_nullable=True) + + with caplog.at_level(logging.INFO): + logs(installation_ctx.workspace_client, install_folder=custom_folder) + + assert "Completed quality-checker workflow run" in caplog.text + + +def test_e2e_with_custom_folder(ws, spark, setup_workflows_with_custom_folder, caplog): + installation_ctx, run_config = setup_workflows_with_custom_folder() + custom_folder = installation_ctx.installation.install_folder() + + e2e(installation_ctx.workspace_client, run_config=run_config.name, install_folder=custom_folder) + + checked_df = spark.table(run_config.output_config.location) + input_df = spark.table(run_config.input_config.location) + + assert checked_df.count() == input_df.count(), "Output table is empty" + + with caplog.at_level(logging.INFO): + logs(installation_ctx.workspace_client, install_folder=custom_folder) + + assert "Completed e2e workflow run" in caplog.text + + +def test_workflows_with_custom_folder(ws, installation_ctx_custom_install_folder): + installation_ctx_custom_install_folder.installation_service.run() + custom_folder = installation_ctx_custom_install_folder.installation.install_folder() + + installed_workflows = workflows( + installation_ctx_custom_install_folder.workspace_client, + ctx=installation_ctx_custom_install_folder.workspace_installer, + install_folder=custom_folder, + ) + + expected_workflows_state = [{'workflow': 'profiler', 'state': 'UNKNOWN', 'started': ''}] + for state in expected_workflows_state: + assert contains_expected_workflows(installed_workflows, state) + + +def test_logs_with_custom_folder(ws, installation_ctx_custom_install_folder, caplog): + installation_ctx_custom_install_folder.installation_service.run() + custom_folder = installation_ctx_custom_install_folder.installation.install_folder() + + with caplog.at_level(logging.INFO): + logs(installation_ctx_custom_install_folder.workspace_client, install_folder=custom_folder) + + assert "No jobs to relay logs for" in caplog.text + + +def test_validate_checks_with_custom_folder_missing_file(ws, installation_ctx_custom_install_folder): + installation_ctx_custom_install_folder.installation.save(installation_ctx_custom_install_folder.config) + custom_folder = installation_ctx_custom_install_folder.installation.install_folder() + + file = f"{custom_folder}/{installation_ctx_custom_install_folder.config.get_run_config().checks_location}" + + with pytest.raises(NotFound, match=f"Checks file {file} missing"): + validate_checks(installation_ctx_custom_install_folder.workspace_client, install_folder=custom_folder) From 1694a721a8b995015bab86d5b2324833279882ed Mon Sep 17 00:00:00 2001 From: Varun Bhandary Date: Tue, 25 Nov 2025 21:21:21 +0000 Subject: [PATCH 06/18] fix: updated test assertion failure scenario totals mismatch / dataset empty --- tests/resources/all_dataset_checks.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/resources/all_dataset_checks.yaml b/tests/resources/all_dataset_checks.yaml index 8711e0b2c..db8f511df 100644 --- a/tests/resources/all_dataset_checks.yaml +++ b/tests/resources/all_dataset_checks.yaml @@ -185,10 +185,10 @@ function: sql_query arguments: # sql query for dataset-level check - no merge_columns needed - query: SELECT COUNT(*) > 0 AS condition FROM {{ input_view }} + query: SELECT COUNT(*) = 0 AS condition FROM {{ input_view }} condition_column: condition # the check fails if this column evaluates to True - msg: dataset has records # optional - name: dataset_not_empty # optional + msg: dataset has no records # optional + name: dataset_is_empty # optional # apply check to multiple columns - criticality: error From 5f4976c8a2613d18758adbd71398e031cc60d900 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 20:47:30 +0100 Subject: [PATCH 07/18] Apply suggestion from @mwojtyczka --- docs/dqx/docs/reference/quality_checks.mdx | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 59dc54fc5..5d2b4650e 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1663,6 +1663,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 + 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: From 03a1b1b0f025d02eda62e92ba7475e4c979edf4b Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 20:48:03 +0100 Subject: [PATCH 08/18] Apply suggestion from @mwojtyczka --- docs/dqx/docs/reference/quality_checks.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 5d2b4650e..c79b3a57f 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1647,7 +1647,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 From 49326a360c8819248aec8a982acfd53a64c1e966 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 20:49:35 +0100 Subject: [PATCH 09/18] Apply suggestion from @mwojtyczka --- docs/dqx/docs/reference/quality_checks.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index c79b3a57f..17a796c85 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -2024,7 +2024,7 @@ checks = [ "negate": True }, - # sql_query check + # sql_query check with merge_columns (row-level validation) DQDatasetRule( criticality="error", check_func=sql_query, From 7d98e2a39cd8fa57281a7ca079abc0c0178d7f56 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 20:51:32 +0100 Subject: [PATCH 10/18] Apply suggestion from @mwojtyczka --- docs/dqx/docs/reference/quality_checks.mdx | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 17a796c85..386452d1f 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -2040,6 +2040,20 @@ checks = [ }, ), + # 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 + "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 DQDatasetRule( criticality="error", From 623c2df038a708d180243ab8438cc0e05a60702c Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 20:52:36 +0100 Subject: [PATCH 11/18] Apply suggestion from @mwojtyczka --- docs/dqx/docs/reference/quality_checks.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 386452d1f..2e13d8fc1 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -1668,7 +1668,7 @@ Complex data types are supported as well. check: function: sql_query arguments: - # sql query for dataset-level check + # 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 From b2062d4334e126404cd2600a0d35ead8347b50b7 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 20:53:01 +0100 Subject: [PATCH 12/18] Apply suggestion from @mwojtyczka --- docs/dqx/docs/reference/quality_checks.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 2e13d8fc1..481718bf7 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -2045,7 +2045,7 @@ checks = [ criticality="error", check_func=sql_query, check_func_kwargs={ - # sql query for dataset-level check + # 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 From b943d5df52b1aa7cb8b4d10930c7267f455f88e4 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 21:01:06 +0100 Subject: [PATCH 13/18] fmt --- docs/dqx/docs/reference/quality_checks.mdx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 7c4932899..d5942f847 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -2046,7 +2046,7 @@ 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) @@ -2060,7 +2060,7 @@ checks = [ "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 From f707a23c1b43e3d0ce90db40890291631a3b5df6 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 21:03:46 +0100 Subject: [PATCH 14/18] Update src/databricks/labs/dqx/check_funcs.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/databricks/labs/dqx/check_funcs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 39b20406e..4e6c26db2 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -3075,7 +3075,7 @@ def _validate_sql_query_params(query: str, merge_columns: list[str] | None) -> N if merge_columns is None: return - if isinstance(merge_columns, str) or not isinstance(merge_columns, Sequence): + 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." ) From 24cdd3fbce1f1452f9599b8d725d5f4fcd08e2d6 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 21:03:57 +0100 Subject: [PATCH 15/18] Update src/databricks/labs/dqx/check_funcs.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/databricks/labs/dqx/check_funcs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 4e6c26db2..7b5f88687 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -1289,7 +1289,7 @@ def sql_query( 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. + 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. From ded5feef58e0f110a9416a335727b9650c3c358f Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 21:05:57 +0100 Subject: [PATCH 16/18] fmt --- tests/integration/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index e06813da4..38fd7b9cd 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -36,7 +36,7 @@ def build_quality_violation( name: str, message: str, - columns: list[str], + columns: list[str] | None, *, function: str = "is_not_null_and_not_empty", ) -> dict[str, Any]: From e50db33bd75a0784986ec0eba36e50b0f52a8cda Mon Sep 17 00:00:00 2001 From: mwojtyczka Date: Sun, 7 Dec 2025 21:23:30 +0000 Subject: [PATCH 17/18] Add pytest-benchmark performance baseline --- docs/dqx/docs/reference/benchmarks.mdx | 2 + tests/perf/.benchmarks/baseline.json | 75 ++++++++++++++++++++++++++ 2 files changed, 77 insertions(+) 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/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", From cddee6be095f709d54709574b8d6b401541124fc Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Sun, 7 Dec 2025 22:30:27 +0100 Subject: [PATCH 18/18] fixed tests --- tests/integration/test_apply_checks.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index 99a40d019..764e7120d 100644 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -7784,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", @@ -7804,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", @@ -7814,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", @@ -7893,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", @@ -7913,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", @@ -7923,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", @@ -8032,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", @@ -8062,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", @@ -8072,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",