diff --git a/demos/dqx_demo_ai_assisted_checks_generation.py b/demos/dqx_demo_ai_assisted_checks_generation.py index 41f7e7399..0c7e79d40 100644 --- a/demos/dqx_demo_ai_assisted_checks_generation.py +++ b/demos/dqx_demo_ai_assisted_checks_generation.py @@ -41,7 +41,7 @@ import os, yaml from databricks.labs.dqx.profiler.generator import DQGenerator -from databricks.labs.dqx.config import LLMModelConfig +from databricks.labs.dqx.config import LLMModelConfig, InputConfig from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient @@ -95,6 +95,6 @@ # COMMAND ---------- -checks = generator.generate_dq_rules_ai_assisted(user_input=user_requirement, table_name=table_name) +checks = generator.generate_dq_rules_ai_assisted(user_input=user_requirement, input_config=InputConfig(location=table_name)) print("======== Generated checks =========") print(checks) diff --git a/docs/dqx/docs/demos.mdx b/docs/dqx/docs/demos.mdx index 9de02f061..d12ac89aa 100644 --- a/docs/dqx/docs/demos.mdx +++ b/docs/dqx/docs/demos.mdx @@ -12,6 +12,7 @@ Import the following notebooks in the Databricks workspace to try DQX out: * [DQX Demo Notebook](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_demo_library.py) - demonstrates how to use DQX as a library. * [DQX Demo Notebook for Data Quality Summary Metrics](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_demo_summary_metrics.py) - demonstrates how to generate summary-level data quality metrics when validating data with DQX. * [DQX Demo Notebook for Profiling and Applying Checks at Scale on Multiple Tables](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_multi_table_demo.py) - demonstrates how to use DQX as a library at scale to apply checks on multiple tables. +* [DQX Demo Notebook for AI-assisted checks generation](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_demo_ai_assisted_checks_generation.py) - demonstrates how to generate DQX rules/checks with LLM using natural language. * [DQX Demo Notebook for Spark Structured Streaming (Native End-to-End Approach)](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_streaming_demo_native.py) - demonstrates how to use DQX as a library with Spark Structured Streaming, using the built-in end-to-end method to handle both reading and writing. * [DQX Demo Notebook for Spark Structured Streaming (DIY Approach)](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_streaming_demo_diy.py) - demonstrates how to use DQX as a library with Spark Structured Streaming, while handling reading and writing on your own outside DQX using Spark API. * [DQX Demo Notebook for Lakeflow Pipelines (formerly DLT)](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_dlt_demo.py) - demonstrates how to use DQX as a library with Lakeflow Pipelines. diff --git a/docs/dqx/docs/guide/ai_assisted_quality_checks_generation.mdx b/docs/dqx/docs/guide/ai_assisted_quality_checks_generation.mdx index 1514913bc..0a6f4cdb4 100644 --- a/docs/dqx/docs/guide/ai_assisted_quality_checks_generation.mdx +++ b/docs/dqx/docs/guide/ai_assisted_quality_checks_generation.mdx @@ -14,7 +14,7 @@ DQX provides the capability to generate data quality rule candidates using AI/LL The AI-assisted quality checks generation is available via the `DQGenerator` class and supports: - Analyzing natural language descriptions of data quality requirements. -- Optionally inspecting table schemas to understand data structure. +- Optionally inspecting input data schemas to understand data structure. - Generating appropriate data quality rules that match the requirements using built-in and custom check functions. - Validating the generated rules to ensure they are syntactically correct. @@ -81,21 +81,22 @@ Here's a simple example of generating quality rules from a natural language desc -### Using with Table Schema +### Using with Input Data Schema -For better results, you can provide a fully qualified table name of the input data. -The LLM will analyze the table schema to generate more accurate rules: +For better results, you can provide a fully qualified table name or path of the input data. +The LLM will analyze the schema of the input data to generate more accurate rules: ```python from databricks.labs.dqx.profiler.generator import DQGenerator + from databricks.labs.dqx.config import InputConfig from databricks.sdk import WorkspaceClient ws = WorkspaceClient() generator = DQGenerator(workspace_client=ws) - # Generate rules with table schema awareness + # Generate rules with input data schema awareness user_input = """ All customer records must have complete contact information. Email addresses must be valid. @@ -105,7 +106,7 @@ The LLM will analyze the table schema to generate more accurate rules: checks = generator.generate_dq_rules_ai_assisted( user_input=user_input, - table_name="catalog1.schema1.customers" + input_config=InputConfig(location="catalog1.schema1.customers") ) print(checks) @@ -209,6 +210,7 @@ Here's a complete workflow showing how to generate, review, and save AI-assisted ```python from databricks.labs.dqx.profiler.generator import DQGenerator from databricks.labs.dqx.engine import DQEngine + from databricks.labs.dqx.config import InputConfig from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig from databricks.sdk import WorkspaceClient import yaml @@ -233,7 +235,8 @@ Here's a complete workflow showing how to generate, review, and save AI-assisted # Generate quality rules using AI checks = generator.generate_dq_rules_ai_assisted( user_input=business_requirements, - table_name="production.users.registrations" + # input location could be a table or path + input_config=InputConfig(location="production.users.registrations") ) # Review the generated checks @@ -265,6 +268,7 @@ You can combine AI-assisted rules with profiler-generated rules for comprehensiv ```python from databricks.labs.dqx.profiler.profiler import DQProfiler from databricks.labs.dqx.profiler.generator import DQGenerator + from databricks.labs.dqx.config import InputConfig from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient @@ -291,7 +295,7 @@ You can combine AI-assisted rules with profiler-generated rules for comprehensiv ai_checks = generator.generate_dq_rules_ai_assisted( user_input=business_requirements, - table_name="catalog1.schema1.sales_data" + input_config=InputConfig(location="catalog1.schema1.sales_data") ) # Step 3: Combine both sets of rules @@ -329,7 +333,7 @@ GDPR Compliance Requirements: checks = generator.generate_dq_rules_ai_assisted( user_input=compliance_requirements, - table_name="gdpr.users.personal_data" + input_config=InputConfig(location="gdpr.users.personal_data") ) ``` @@ -349,7 +353,7 @@ Financial Transaction Rules: checks = generator.generate_dq_rules_ai_assisted( user_input=financial_requirements, - table_name="finance.transactions.daily" + input_config=InputConfig(location="finance.transactions.daily") ) ``` @@ -369,7 +373,7 @@ IoT Sensor Data Quality Rules: checks = generator.generate_dq_rules_ai_assisted( user_input=iot_requirements, - table_name="iot.sensors.readings" + input_config=InputConfig(location="iot.sensors.readings") ) ``` @@ -410,7 +414,7 @@ Use a slash-separated format for specifying the scope and key, such as `api_key: ## Best Practices 1. **Be Specific**: Provide clear and specific business requirements in your input. The more detailed your description, the better the generated rules. However, be aware that the model endpoint that you use have token limitations [see Databricks Model APIs limits and quotas](https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/limits). -2. **Include Table Name**: When possible, provide the fully qualified table name to help the LLM understand the actual data structure. +2. **Include Input Config**: When possible, provide input config specifying location of input data as a fully qualified table name or path to help the LLM understand the actual data structure. 3. **Review Generated Rules**: Always review the generated rules before applying them to production data. The AI may not perfectly understand all nuances of your requirements so treat the generated rules as candidates and review them. 4. **Combine Approaches**: Use AI-assisted generation for business logic rules and profiler-based generation for statistical (technical) rules. 5. **Iterate**: If the generated rules don't match your expectations, refine your input description and regenerate, or update manually. @@ -432,7 +436,7 @@ pip install 'databricks-labs-dqx[llm]' **Solution**: - Make your input description more specific and detailed. -- Provide the table name so the LLM can analyze the actual schema. +- Provide the input data location so the LLM can analyze the actual schema. - Break complex requirements into simpler, more focused descriptions. - Include examples in your description. @@ -452,7 +456,7 @@ pip install 'databricks-labs-dqx[llm]' **Solution**: - Check the validation error messages for specific issues. -- Verify that the column names mentioned in your requirements exist in the table. +- Verify that the column names mentioned in your requirements exist in the input data. - Ensure the rule types requested are supported by DQX (see [Quality Checks Reference](/docs/reference/quality_checks)). - Simplify your requirements and regenerate. diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index 395b6c80b..4587eab77 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_view }}`, 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 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 | | `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 | @@ -1748,7 +1748,7 @@ from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient ref_df = spark.table("catalog1.schema1.ref_table") -ref_dfs = {"ref_view": ref_df} # mapping of reference data name to DataFrame, multiple references can be provided +ref_dfs = {"ref_df_key": ref_df} # mapping of reference data name to DataFrame, multiple references can be provided dq_engine = DQEngine(WorkspaceClient()) valid_and_quarantine_df = dq_engine.apply_checks_by_metadata(df, checks, ref_dfs=ref_dfs) @@ -1756,18 +1756,18 @@ good_df, quarantine_df = dq_engine.apply_checks_by_metadata_and_split(df, check ``` The reference DataFrames are used in selected Dataset-level checks: -* `foreign_key`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_view"`. +* `foreign_key`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_df_key"`. The value of `ref_df_name` must match the key in the `ref_dfs` dictionary. -* `compare_datasets`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_view"`. +* `compare_datasets`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_df_key"`. The value of `ref_df_name` must match the key in the `ref_dfs` dictionary. * `sql_query`: the reference DataFrames are registered as temporary views and can be used in the sql query. - For example, if you have a reference DataFrame named `ref_view`, you can use it in the SQL query as `{{ ref_view }}`: + For example, if you have a reference DataFrame named `ref_df_key`, you can use it in the SQL query as `{{ ref_df_key }}`: ```sql - SELECT col1, col2, SUM(col3) = 0 AS condition FROM {{ input_view }} input JOIN {{ ref_view }} ref ON input.col1 = ref.ref_col1 GROUP BY col1, col2 + SELECT col1, col2, SUM(col3) = 0 AS condition FROM {{ input_view }} input JOIN {{ ref_df_key }} ref ON input.col1 = ref.ref_col1 GROUP BY col1, col2 ``` You can also use reference table directly in the `sql` query without supplying it as a DataFrame: @@ -1783,7 +1783,7 @@ Each reference dataset should be specified as a fully qualified table name (cata Example: ``` reference_tables: - ref_view: # <- name of the reference data + ref_df_key: # <- name of the reference data input_config: # <- input data configuration for the reference table format: delta location: catalog1.schema1.ref_table @@ -2143,7 +2143,7 @@ from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient ref_df = spark.table("catalog1.schema1.ref_table") -ref_dfs = {"ref_view": ref_df} +ref_dfs = {"ref_df_key": ref_df} dq_engine = DQEngine(WorkspaceClient()) valid_and_quarantine_df = dq_engine.apply_checks(df, checks, ref_dfs=ref_dfs) @@ -2151,18 +2151,18 @@ good_df, quarantine_df = dq_engine.apply_checks_and_split(df, checks, ref_dfs=r ``` The reference DataFrames are used in selected Dataset-level checks: -* `foreign_key`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_view"`. +* `foreign_key`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_df_key"`. The value of `ref_df_name` must match the key in the `ref_dfs` dictionary. -* `compare_datasets`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_view"`. +* `compare_datasets`: required for this check if `ref_df_name` argument is specified and not `ref_table`, e.g. `ref_df_name="ref_df_key"`. The value of `ref_df_name` must match the key in the `ref_dfs` dictionary. * `sql_query`: the reference DataFrames are registered as temporary views and can be used in the sql query. - For example, if you have a reference DataFrame named `ref_view`, you can use it in the SQL query as `{{ ref_view }}`: + For example, if you have a reference DataFrame named `ref_df_key`, you can use it in the SQL query as `{{ ref_df_key }}`: ```sql - SELECT col1, col2, SUM(col3) = 0 AS condition FROM {{ input_view }} input JOIN {{ ref_view }} ref ON input.col1 = ref.ref_col1 GROUP BY col1, col2 + SELECT col1, col2, SUM(col3) = 0 AS condition FROM {{ input_view }} input JOIN {{ ref_df_key }} ref ON input.col1 = ref.ref_col1 GROUP BY col1, col2 ``` You can also use reference table directly in the `sql` query without supplying it as a DataFrame: diff --git a/src/databricks/labs/dqx/llm/llm_core.py b/src/databricks/labs/dqx/llm/llm_core.py index 2f83d26e4..b9e147e95 100644 --- a/src/databricks/labs/dqx/llm/llm_core.py +++ b/src/databricks/labs/dqx/llm/llm_core.py @@ -93,11 +93,12 @@ class DspyRuleSignature(dspy.Signature): available_functions: str = dspy.InputField(desc="JSON string of available DQX check functions") quality_rules: str = dspy.OutputField( desc=( - "Return a valid JSON array of data quality rules. Use double quotes only. " + "Return a valid JSON array of data quality rules. Use double quotes only and capitalize SQL keywords. " "Criticality can be error or warn. " + "Filter may be used to apply the rule to the relevant records only. " "Check function name and doc to select the appropriate check function. " - "Format: [{\"criticality\":\"error\",\"check\":{\"function\":\"name\",\"arguments\":{\"column\":\"col\"}}}] " - "Example: [{\"criticality\":\"error\",\"check\":{\"function\":\"is_not_null\",\"arguments\":{\"column\":\"customer_id\"}}}]" + "Format: [{\"criticality\":\"error\",\"check\":{\"function\":\"name\",\"arguments\":{\"column\":\"col\"}},\"filter\":\"expression\"}] " + "Example: [{\"criticality\":\"error\",\"check\":{\"function\":\"is_not_null\",\"arguments\":{\"column\":\"customer_id\"}},\"filter\":\"customer_name is not null\"}]" ) ) reasoning: str = dspy.OutputField(desc="Explanation of why these rules were chosen") diff --git a/src/databricks/labs/dqx/llm/llm_utils.py b/src/databricks/labs/dqx/llm/llm_utils.py index 8dbd7cb48..98aaa9981 100644 --- a/src/databricks/labs/dqx/llm/llm_utils.py +++ b/src/databricks/labs/dqx/llm/llm_utils.py @@ -7,8 +7,11 @@ import json import yaml import dspy # type: ignore +from pyspark.sql import SparkSession from databricks.labs.dqx.checks_resolver import resolve_check_function from databricks.labs.dqx.rule import CHECK_FUNC_REGISTRY +from databricks.labs.dqx.config import InputConfig +from databricks.labs.dqx.io import read_input_data logger = logging.getLogger(__name__) @@ -100,6 +103,23 @@ def create_optimizer_training_set(custom_check_functions: dict[str, Callable] | return examples +def get_column_metadata(spark: SparkSession, input_config: InputConfig) -> str: + """ + Get the column metadata for a given table. + + Args: + input_config (InputConfig): Input configuration for the table. + spark (SparkSession): The Spark session used to access the table. + + Returns: + str: A JSON string containing the column metadata with columns wrapped in a "columns" key. + """ + df = read_input_data(spark, input_config) + columns = [{"name": field.name, "type": field.dataType.simpleString()} for field in df.schema.fields] + schema_info = {"columns": columns} + return json.dumps(schema_info) + + def _load_training_examples() -> list[dict[str, Any]]: """A function to load the training examples from the llm/resources/training_examples.yml file. diff --git a/src/databricks/labs/dqx/llm/resources/training_examples.yml b/src/databricks/labs/dqx/llm/resources/training_examples.yml index 09c4a46fc..64df41298 100644 --- a/src/databricks/labs/dqx/llm/resources/training_examples.yml +++ b/src/databricks/labs/dqx/llm/resources/training_examples.yml @@ -7,14 +7,16 @@ quality_rules: '[{"criticality":"error","check":{"function":"is_not_null_and_not_empty","arguments":{"column":"product_code"}}}]' reasoning: "Product code is a key identifier and cannot be missing or blank" -- name: "status_in_allowed_values" +- name: "status_in_allowed_values_for_placed_orders" schema_info: columns: - name: "status" type: "integer" - business_description: "Status must be one of the allowed values: 1=Active, 2=Inactive, 3=Pending" - quality_rules: '[{"criticality":"error","check":{"function":"is_in_list","arguments":{"column":"status","allowed":[1,2,3]}}}]' - reasoning: "Restrict status to predefined set of values for data consistency" + - name: "order_id" + type: "string" + business_description: "Status must be one of the allowed values for all placed orders: 1=Active, 2=Inactive, 3=Pending" + quality_rules: '[{"criticality":"error","check":{"function":"is_in_list","arguments":{"column":"status","allowed":[1,2,3]}},"filter": "order_id is not null"}]' + reasoning: "Restrict status to predefined set of values for all placed orders (where order_id not null) for data consistency" - name: "order_quantity_in_range" schema_info: diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 05475ae37..206e04e2f 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -6,17 +6,17 @@ from databricks.sdk import WorkspaceClient from databricks.labs.dqx.base import DQEngineBase -from databricks.labs.dqx.config import LLMModelConfig +from databricks.labs.dqx.config import LLMModelConfig, InputConfig from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.profiler.common import val_maybe_to_str from databricks.labs.dqx.profiler.profiler import DQProfile from databricks.labs.dqx.telemetry import telemetry_logger from databricks.labs.dqx.errors import MissingParameterError -from databricks.labs.dqx.utils import get_column_metadata # Conditional imports for LLM-assisted rules generation try: from databricks.labs.dqx.llm.llm_engine import DQLLMEngine + from databricks.labs.dqx.llm.llm_utils import get_column_metadata LLM_ENABLED = True except ImportError: @@ -90,14 +90,14 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str return dq_rules @telemetry_logger("generator", "generate_dq_rules_ai_assisted") - def generate_dq_rules_ai_assisted(self, user_input: str, table_name: str = "") -> list[dict]: + def generate_dq_rules_ai_assisted(self, user_input: str, input_config: InputConfig | None = None) -> list[dict]: """ Generates data quality rules using LLM based on natural language input. Args: user_input: Natural language description of data quality requirements. - table_name: Optional fully qualified table name. - If not provided, LLM will be used to guess the table schema. + input_config: Optional input config providing input data location as a path or fully qualified table name + to infer schema. If not provided, LLM will be used to guess the table schema. Returns: A list of dictionaries representing the generated data quality rules. @@ -112,7 +112,7 @@ def generate_dq_rules_ai_assisted(self, user_input: str, table_name: str = "") - ) logger.info(f"Generating DQ rules with LLM for input: '{user_input}'") - schema_info = get_column_metadata(self.spark, table_name) if table_name else "" + schema_info = get_column_metadata(self.spark, input_config) if input_config else "" # Generate rules using pre-initialized LLM compiler prediction = self.llm_engine.get_business_rules_with_llm(user_input=user_input, schema_info=schema_info) diff --git a/src/databricks/labs/dqx/profiler/profiler_runner.py b/src/databricks/labs/dqx/profiler/profiler_runner.py index 306e3635b..221fdb6f2 100644 --- a/src/databricks/labs/dqx/profiler/profiler_runner.py +++ b/src/databricks/labs/dqx/profiler/profiler_runner.py @@ -10,6 +10,7 @@ BaseChecksStorageConfig, InstallationChecksStorageConfig, RunConfig, + InputConfig, ) from databricks.labs.dqx.config_serializer import ConfigSerializer from databricks.labs.dqx.engine import DQEngine @@ -79,7 +80,7 @@ def run( if run_config.checks_user_requirements: checks += generator.generate_dq_rules_ai_assisted( - run_config.checks_user_requirements, table_name=run_config.input_config.location + user_input=run_config.checks_user_requirements, input_config=run_config.input_config ) storage_config = InstallationChecksStorageConfig( @@ -135,7 +136,9 @@ def run_for_patterns( logger.info(f"Generated summary statistics: \n{summary_stats}") if run_config.checks_user_requirements: - checks += generator.generate_dq_rules_ai_assisted(run_config.checks_user_requirements, table_name=table) + checks += generator.generate_dq_rules_ai_assisted( + user_input=run_config.checks_user_requirements, input_config=InputConfig(location=table) + ) storage_config = InstallationChecksStorageConfig( location=( diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py index ce6df48a7..3ec8d0b58 100644 --- a/src/databricks/labs/dqx/utils.py +++ b/src/databricks/labs/dqx/utils.py @@ -6,7 +6,7 @@ from typing import Any from fnmatch import fnmatch -from pyspark.sql import Column, SparkSession +from pyspark.sql import Column # Import spark connect column if spark session is created using spark connect try: @@ -284,23 +284,6 @@ def list_tables( raise NotFound("No tables found matching include or exclude criteria") -def get_column_metadata(spark: SparkSession, table_name: str) -> str: - """ - Get the column metadata for a given table. - - Args: - table_name (str): The name of the table to retrieve metadata for. - spark (SparkSession): The Spark session used to access the table. - - Returns: - str: A JSON string containing the column metadata with columns wrapped in a "columns" key. - """ - df = spark.table(table_name) - columns = [{"name": field.name, "type": field.dataType.simpleString()} for field in df.schema.fields] - schema_info = {"columns": columns} - return json.dumps(schema_info) - - def _split_pattern(pattern: str) -> tuple[str, str, str]: """ Splits a wildcard pattern into its catalog, schema, and table components. diff --git a/tests/integration/test_ai_rules_generator.py b/tests/integration/test_ai_rules_generator.py index 034561d3e..cf3a2dc04 100644 --- a/tests/integration/test_ai_rules_generator.py +++ b/tests/integration/test_ai_rules_generator.py @@ -1,14 +1,16 @@ import pyspark.sql.functions as F +from pyspark.sql.types import StructType, StructField, StringType, IntegerType from tests.conftest import TEST_CATALOG from databricks.labs.dqx.engine import DQEngineCore from databricks.labs.dqx.profiler.generator import DQGenerator -from databricks.labs.dqx.config import LLMModelConfig +from databricks.labs.dqx.config import LLMModelConfig, InputConfig from databricks.labs.dqx.check_funcs import make_condition, register_rule USER_INPUT = """ -Username should not start with 's' if age is less than 18. Use exact wording if needed in the generated rule. -All users must have a valid email address. +Username should not start with 's' and should not contain more than 20 letters if user id is provided. Error message must be: "Username should not start with 's' and should not contain more than 20 letters if user id is provided" +Apply validation when user_id is not null. +Users at age 18 or above must have a valid email address. Age should be between 0 and 120. """ @@ -16,13 +18,14 @@ { "check": { "arguments": { - "columns": ["username", "age"], - "expression": "NOT (username LIKE 's%' AND age < 18)", - "msg": "Username should not start with 's' if age is less than 18", + "columns": ["username"], + "expression": "NOT (username LIKE 's%') AND LENGTH(username) <= 20", + "msg": "Username should not start with 's' and should not contain more than 20 letters if user id is provided", }, "function": "sql_expression", }, "criticality": "error", + "filter": "user_id IS NOT NULL", }, { "check": { @@ -30,6 +33,7 @@ "function": "regex_match", }, "criticality": "error", + "filter": "age >= 18", }, { "check": {"arguments": {"column": "age", "max_limit": 120, "min_limit": 0}, "function": "is_in_range"}, @@ -52,7 +56,35 @@ def test_generate_dq_rules_ai_assisted_with_input_table(ws, spark, make_table, m columns=[("user_id", "string"), ("username", "string"), ("email", "string"), ("age", "int")], ) generator = DQGenerator(ws, spark) - actual_checks = generator.generate_dq_rules_ai_assisted(user_input=USER_INPUT, table_name=input_table.full_name) + actual_checks = generator.generate_dq_rules_ai_assisted( + user_input=USER_INPUT, input_config=InputConfig(location=input_table.full_name) + ) + assert actual_checks == EXPECTED_CHECKS + + +def test_generate_dq_rules_ai_assisted_with_input_path(ws, spark, make_directory): + folder = make_directory() + workspace_file_path = str(folder.absolute()) + "/input_data.parquet" + + schema = StructType( + [ + StructField("user_id", StringType(), True), + StructField("username", StringType(), True), + StructField("email", StringType(), True), + StructField("age", IntegerType(), True), + ] + ) + + test_data = [ + ("user1", "john_doe", "john@example.com", 25), + ] + df = spark.createDataFrame(test_data, schema=schema) + df.write.mode("overwrite").parquet(workspace_file_path) + + generator = DQGenerator(ws, spark) + actual_checks = generator.generate_dq_rules_ai_assisted( + user_input=USER_INPUT, input_config=InputConfig(location=workspace_file_path, format="parquet") + ) assert actual_checks == EXPECTED_CHECKS diff --git a/tests/unit/test_llm_utils.py b/tests/unit/test_llm_utils.py index 314e43e6a..caa04a989 100644 --- a/tests/unit/test_llm_utils.py +++ b/tests/unit/test_llm_utils.py @@ -1,12 +1,18 @@ import inspect import json +from unittest.mock import Mock +import pytest import dspy # type: ignore import pyspark.sql.functions as F +from pyspark.sql.types import StructField, StringType, IntegerType + from databricks.labs.dqx.check_funcs import make_condition, register_rule +from databricks.labs.dqx.config import InputConfig from databricks.labs.dqx.llm.llm_utils import ( get_check_function_definitions, create_optimizer_training_set, get_required_check_functions_definitions, + get_column_metadata, ) @@ -148,3 +154,37 @@ def test_get_training_examples_with_custom_check_functions(): ) ] assert filtered_examples + + +@pytest.mark.parametrize( + "location, spark_read_mock_method", + [ + ("catalog.schema.table", "table"), + ("s3://bucket/path/to/data", "load"), + ], +) +def test_get_column_metadata(location, spark_read_mock_method): + mock_spark = Mock() + mock_df = Mock() + mock_schema = Mock() + mock_schema.fields = [ + StructField("customer_id", StringType(), True), + StructField("first_name", StringType(), True), + StructField("last_name", StringType(), True), + StructField("age", IntegerType(), True), + ] + mock_df.schema = mock_schema + + # Dynamically set the mock method (either table or load) + setattr(mock_spark.read.options.return_value, spark_read_mock_method, Mock(return_value=mock_df)) + + result = get_column_metadata(mock_spark, InputConfig(location=location)) + expected_result = { + "columns": [ + {"name": "customer_id", "type": "string"}, + {"name": "first_name", "type": "string"}, + {"name": "last_name", "type": "string"}, + {"name": "age", "type": "int"}, + ] + } + assert result == json.dumps(expected_result) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a45b2643b..a2acc39e9 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1,11 +1,9 @@ -import json from datetime import date, datetime from typing import Any from unittest.mock import Mock import pyspark.sql.functions as F import pytest from pyspark.sql import Column -from pyspark.sql.types import StructField, StringType, IntegerType from databricks.labs.dqx.io import read_input_data, get_reference_dataframes from databricks.labs.dqx.utils import ( @@ -17,7 +15,6 @@ is_simple_column_expression, normalize_bound_args, safe_strip_file_from_path, - get_column_metadata, ) from databricks.labs.dqx.errors import InvalidParameterError, InvalidConfigError from databricks.labs.dqx.config import InputConfig @@ -369,26 +366,3 @@ def test_get_reference_dataframes_with_missing_ref_tables() -> None: ) def test_safe_strip_file_from_path(path: str, expected: str): assert safe_strip_file_from_path(path) == expected - - -def test_column_metadata(): - mock_spark = Mock() - mock_df = Mock() - mock_df.schema.fields = [ - StructField("customer_id", StringType(), True), - StructField("first_name", StringType(), True), - StructField("last_name", StringType(), True), - StructField("age", IntegerType(), True), - ] - mock_spark.table.return_value = mock_df - - result = get_column_metadata(mock_spark, "test_table") - expected_result = { - "columns": [ - {"name": "customer_id", "type": "string"}, - {"name": "first_name", "type": "string"}, - {"name": "last_name", "type": "string"}, - {"name": "age", "type": "int"}, - ] - } - assert result == json.dumps(expected_result)