Aip 99 llmdataqualityoperator

[cetingokhan]

Add LLMDataQualityOperator to apache-airflow-providers-common-ai

This PR introduces LLMDataQualityOperator, a new operator that uses an LLM agent to plan and execute data-quality checks against a SQL database, with optional human-in-the-loop (HITL) approval before any SQL is executed.

What is added

Core operator — LLMDataQualityOperator

• Receives a list of natural-language DQCheckInput definitions.
• Runs an LLM agent that discovers the database schema, writes SQL queries, and selects validators from a registry.
• Evaluates each check and returns a structured DQReport; raises DQCheckFailedError when any check fails.
• Supports require_approval=True: the operator first runs the LLM planning phase (schema discovery + validator assignment, no SQL execution), defers for HITL review in the Airflow UI, and only executes SQL and applies validators after a human approves.

Toolset layer

• BaseDQToolset — abstract base for all DQ toolsets. Exposes list_checks to the agent and declares an output_mode property with two values:
  • "execute" — the toolset runs checks against a live data source and the operator gates on pass/fail via DQReport.
  • "generate" — the toolset produces a configuration artifact (e.g. a Great Expectations suite JSON, a SODA YAML) that a downstream operator can feed to an external DQ framework. The operator returns the config string as its XCom value without performing pass/fail gating. This makes it straightforward to implement GEXDQToolset, SodaDQToolset, or any other framework-backed toolset by subclassing BaseDQToolset and setting output_mode = "generate".
• SQLDQToolset — extends BaseDQToolset (output_mode = "execute") with list_validators and apply_validator. Supports scalar, row-level, and fixed (pre-assigned) validators.
• SQLToolset — gives the agent safe, read-only access to a SQL database via DbApiHook. Provides list_tables, get_schema, query, and check_query tools. SQL validation (sqlglot), row-limit truncation, and retryable-error handling (ModelRetry) are built in.
• DataFusionToolset — registers S3/object-store Parquet/CSV data as a queryable DataFusion table.

Validator registry (utils/dataquality/)

• ValidatorRegistry + @register_validator decorator for defining reusable validator factories with LLM context hints.
• Built-in factories: null_pct_check, row_count_check, duplicate_pct_check, between_check, exact_check.
• Row-level validator support (plain SELECT col FROM table, validator applied per-row with invalid_pct gating).

@task.llm_data_quality decorator

• Wraps LLMDataQualityOperator for use in TaskFlow-style Dags.

Documentation

• Full operator guide at llm_data_quality.rst covering single-table checks, multi-toolset usage, HITL approval flow, custom validators, and the @task.llm_data_quality decorator.

---

Was generative AI tooling used to co-author this PR?

• Yes — GitHub Copilot (Claude Sonnet 4.6, Claude Opus 4.6, GPT-5.3-Codex)

Generated-by: GitHub Copilot (Claude Sonnet 4.6) following the guidelines

[cetingokhan]

Sample Task Log;

.......
[2026-03-05T17:41:42.539957Z] INFO - Registered object store for schema: s3://
[2026-03-05T17:41:42.918575Z] INFO - Registered data source format parquet for table: sales_data
[2026-03-05T17:41:42.920314Z] INFO - Using schema context:
Table: sales_data
Columns: price: int64
area: int64
bedrooms: int64
bathrooms: int64
stories: int64
mainroad: string_view
guestroom: string_view
basement: string_view
hotwaterheating: string_view
airconditioning: string_view
parking: int64
prefarea: string_view
furnishingstatus: string_view
[2026-03-05T17:41:42.947772Z] INFO - DQ plan cache miss — generating via LLM (key: 'dq_plan_v8_535ba38ccff9be9a').
[2026-03-05T17:41:42.948151Z] INFO - Using system prompt:
You are a data-quality SQL expert.

Given a set of named data-quality checks and a database schema, produce a DQPlan that minimises the number of SQL queries executed.

PRIMARY RULE — combine everything on the same table into ONE SELECT:
  All checks that query the same table MUST be merged into a single SELECT
  statement. Each check becomes one output column in that statement.

  CORRECT (one query for two checks on the same table):
    SELECT
      (COUNT(*) FILTER (WHERE email IS NULL) * 100.0 / COUNT(*)) AS null_email_pct,
      COUNT(*) FILTER (WHERE bathrooms < 0)                      AS invalid_bathrooms
    FROM customers

  WRONG (two separate queries for the same table):
    SELECT COUNT(*) FILTER (WHERE email IS NULL) AS null_email_count FROM customers
    SELECT COUNT(*) FILTER (WHERE bathrooms < 0) AS invalid_bathrooms FROM customers

GROUPING STRATEGY:
  Assign a group_id that describes the table being queried (e.g. "sales_data_checks").
  Only split into multiple groups when the checks genuinely require different tables
  or subqueries that cannot be expressed as columns of a single SELECT.

OUTPUT RULES:
  1. Each output column must be aliased to exactly the metric_key of its check.
     Example: ... AS null_email_pct
  2. Each check_name must exactly match the key in the prompts dict.
  3. metric_key values must be valid SQL column aliases (snake_case, no spaces).
  4. Generates only SELECT queries — no INSERT, UPDATE, DELETE, DROP, or DDL.
  5. Use SQL syntax.
  6. Each check must appear in exactly ONE group.
  7. Return a valid DQPlan object. No extra commentary.

Available schema:
Table: sales_data
Columns: price: int64
area: int64
bedrooms: int64
bathrooms: int64
stories: int64
mainroad: string_view
guestroom: string_view
basement: string_view
hotwaterheating: string_view
airconditioning: string_view
parking: int64
prefarea: string_view
furnishingstatus: string_view

[2026-03-05T17:41:42.948518Z] INFO - Using user message:
Generate a DQPlan for the following data-quality checks.

IMPORTANT: All checks that query the same table MUST be combined into a single SELECT with one output column per check.

Checks:
  - check_name="null_mainroad": Check the percentage of rows where mainroad is NULL
  - check_name="invalid_bathrooms": Count rows where bathrooms is negative or NULL
  - check_name="bathroom_data": Count rows where bathrooms is greater than 2
  - check_name="furnishingstatus_values": furnishingstatus column contains only 'furnished', 'semi-furnished', or 'unfurnished'
......
[2026-03-05T17:41:49.436353Z] INFO - Registered data source format parquet for table: sales_data
[2026-03-05T17:41:49.438619Z] INFO - Executing query: SELECT
  (COUNT(*) FILTER (WHERE mainroad IS NULL) * 100.0 / COUNT(*)) AS null_mainroad,
  COUNT(*) FILTER (WHERE bathrooms IS NULL OR bathrooms < 0) AS invalid_bathrooms,
  COUNT(*) FILTER (WHERE bathrooms > 2) AS bathroom_data,
  COUNT(*) FILTER (WHERE furnishingstatus NOT IN ('furnished', 'semi-furnished', 'unfurnished')) AS furnishingstatus_values
FROM sales_data
[2026-03-05T17:41:49.617786Z] INFO - All 4 data-quality check(s) passed.
[2026-03-05T17:41:49.619021Z] INFO - Pushing xcom ti=RuntimeTaskInstance(id=UUID('019cbf17-4648-7359-ab77-9e24023d0c69'), task_id='validate_sales_events', dag_id='example_llm_dq_s3_parquet', run_id='manual__2026-03-05T17:41:39.517609+00:00', try_number=1, dag_version_id=UUID('019cbf09-ae3e-7f97-afb3-cc5ec0a4160f'), map_index=-1, hostname='b483f76e2940', context_carrier={}, task=<Task(LLMDataQualityOperator): validate_sales_events>, bundle_instance=LocalDagBundle(name=dags-folder), max_tries=0, start_date=datetime.datetime(2026, 3, 5, 17, 41, 40, 355724, tzinfo=datetime.timezone.utc), end_date=None, state=<TaskInstanceState.RUNNING: 'running'>, is_mapped=False, rendered_map_index=None, sentry_integration='') 
[2026-03-05T17:41:49.669975Z] INFO - Task instance in success state
[2026-03-05T17:41:49.670192Z] INFO -  Previous state of the Task instance: TaskInstanceState.RUNNING
[2026-03-05T17:41:49.670349Z] INFO - Task operator:<Task(LLMDataQualityOperator): validate_sales_events>

[potiuk]

@cetingokhan Converting to draft — this PR doesn't yet meet our Pull Request quality criteria.

• ❌ Pre-commit / static checks: Failing: CI image checks / Static checks. Run prek run --from-ref main locally to find and fix issues. See Pre-commit / static checks docs.
• ❌ Build docs: Failing: CI image checks / Build documentation (--docs-only). Run breeze build-docs locally to reproduce. See Build docs docs.
• ❌ Provider tests: Failing: provider distributions tests / Compat 3.0.6:P3.10:, provider distributions tests / Compat 3.1.7:P3.10:. Run provider tests with breeze run pytest <provider-test-path> -xvs. See Provider tests docs.
• ❌ Other failing CI checks: Failing: CI image checks / Build documentation (--spellcheck-only). Run prek run --from-ref main locally to reproduce. See static checks docs.

Note: Your branch is 6 commits behind main. Please rebase and push again to get up-to-date CI results.

See the linked criteria for how to fix each item, then mark the PR "Ready for review". This is not a rejection — just an invitation to bring the PR up to standard. No rush.

---

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

[cetingokhan]

Hi @kaxil and @gopidesupavan,

I’ve reached a solid milestone with the LLMDataQualityOperator and wanted to share the current progress. Since LLM-based data quality is a broad topic, I’ve implemented several core features, but I’d love your feedback on whether we should keep this scope or simplify certain parts.

Current Implementation Highlights:

Foundation: The operator is built on top of the LLMOperator base.

Performance Optimization: Instead of running individual queries for every rule, I’ve implemented batch execution. The logic intelligently groups rules by table, DQ type, and column type to minimize overhead.

Approval-Ready Output: When dry_run=True, the operator writes a Markdown-formatted report to XCom. This makes the manual approval process much cleaner and more human-readable in the UI.

Issue Detection (Optional): Users can opt to generate specific SQL queries to identify records that fail validation rules.

Result Collection: If collect_unexpected is enabled, the operator executes the generated SQL for failed rules and lists the problematic records in XCom as JSON (governed by a configurable limit).

Caching Mechanism: To avoid redundant LLM calls, I’ve implemented a way to store and reuse results in Airflow Variables when the quality rules remain unchanged.

My Question:
I feel this covers the essential "Data Quality with LLM" workflow, but I’m wary of over-engineering. Does this feature set feel right for the initial version, or would you prefer a more "stripped-down" approach to start with?

Looking forward to your guidance!

[cetingokhan]

Hi @kaxil and @gopidesupavan ,
Could I get your comments on the scope, the current status of which I've shared below?

Final Scopes

Plan generation & caching

• Takes a prompts dict where each key is a check name and each value is a plain-language description
• Sends all prompts and the target schema to the LLM in one call; the LLM groups related checks into optimised SQL queries to minimise round-trips
• Serialises the generated plan to an Airflow Variable keyed by a hash of prompts + prompt_version + collect_unexpected — repeat runs skip the LLM call entirely

Built-in validators

• Factory functions like null_pct_check, row_count_check, unique_pct_check are ready to use from dq_validation
• Each factory returns a callable that receives the raw scalar metric from the SQL result and returns True (pass) or False (fail)
• Plain lambdas also work: "row_check": lambda v: v >= 1000

Custom aggregate validators with register_validator

• @register_validator("my_check") registers a factory function in the global registry
• The registered validator works the same way as built-ins: receives a raw scalar from the SQL SELECT column and returns a bool
• Useful for domain-specific aggregate thresholds that do not fit the built-in factories

Custom row-level validators with register_validator(row_level=True)

• @register_validator("my_check", row_level=True) marks the validator as row-level
• The LLM generates a plain SELECT for the column instead of an aggregate query
• Each individual row value is passed to the validator callable; Airflow counts how many rows fail
• Result is returned as a RowLevelResult with total, invalid, invalid_pct, and sample_violations
• Pass/fail is decided by comparing invalid_pct against the validator's _max_invalid_pct attribute

dry_run mode

• Set dry_run=True to generate and cache the plan without running any SQL
• Only returns the serialised plan dict immediately without any control results

require_approval — HITL integration

• Set require_approval=True to gate SQL execution on human review
• After plan generation the task defers, surfacing the full plan as a structured markdown review body in the HITL interface
• SQL checks run only after the reviewer approves; rejection raises HITLRejectException
• dry_run=True takes precedence — combining both flags returns the plan dict immediately without requesting approval???

Unexpected row collection

• Set collect_unexpected=True to have the LLM also generate a detail query alongside each validity/format check
• When a check fails the detail query runs and attaches up to unexpected_sample_size violating rows to the report

[Copilot]

Pull request overview

This PR adds an LLM-driven data quality feature set to the providers/common/ai provider: a new LLMDataQualityOperator (plus TaskFlow decorator), supporting models/utilities for plan generation/execution, schema introspection (DB + object storage via DataFusion), and documentation/examples.

Changes:

• Introduces LLMDataQualityOperator with plan caching, optional HITL approval, and “unexpected rows” collection for failed validity checks.
• Adds SQL DQ planning/execution utilities (SQLDQPlanner), pydantic models for plans/results, and a validator registry with built-in validator factories (including row-level validation support).
• Adds unit tests, example DAGs, and provider/docs wiring for the new operator/decorator.

Reviewed changes

Copilot reviewed 18 out of 19 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
uv.lock	Updates locked dependency metadata (incl. oauth extra/authlib).
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py	Implements the new operator, caching, dry-run, HITL approval, and reporting.
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py	Adds LLM-based DQ plan generation and SQL execution (DB + DataFusion), including row-level handling and retry-on-invalid-SQL.
providers/common/ai/src/airflow/providers/common/ai/utils/dq_validation.py	Adds validator registry + built-in validator factories and LLM context injection.
providers/common/ai/src/airflow/providers/common/ai/utils/dq_models.py	Adds plan/result models (including row-level and unexpected-row result types).
providers/common/ai/src/airflow/providers/common/ai/utils/db_schema.py	Adds shared schema introspection utilities (DB + DataFusion).
providers/common/ai/src/airflow/providers/common/ai/decorators/llm_data_quality.py	Adds @task.llm_dq decorator wrapping the operator.
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py	Adds example DAGs demonstrating operator usage.
providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py	Registers new operator/decorator docs and modules.
providers/common/ai/provider.yaml	Registers new operator/decorator docs and modules.
providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py	Broadens execute_complete return type to support subclasses returning richer types.
providers/common/ai/docs/operators/llm_data_quality.rst	Adds operator documentation and examples.
providers/common/ai/tests/unit/common/ai/...	Adds unit tests for planner/models/validators/operator/decorator/schema utils.
docs/spelling_wordlist.txt	Adds spelling exceptions used by new docs/code.

[Copilot]

Pull request overview

Copilot reviewed 19 out of 20 changed files in this pull request and generated 6 comments.

[Copilot]

Pull request overview

Copilot reviewed 19 out of 20 changed files in this pull request and generated 7 comments.

[gopidesupavan]

I like the overall direction here, but I wonder if the long-term authoring model should be more config-driven than Python-driven.

Instead of requiring DAG authors to define many checks as Python validators/factories, could the LLM generate a simple JSON/YAML rule spec from the natural-language prompts, and then have the operator execute that spec with a deterministic engine?

An example qualink https://github.com/gopidesupavan/qualink and https://gopidesupavan.github.io/qualink/guide/yaml-config/

Not sure similarly may be great expection may have something like that..

So that based on the prompts llm generates rules and we can simply execute on them?

I do think a config-first layer for common checks may scale better than growing the Python validator API. Python validators could still remain as an escape hatch for advanced/custom cases.

If you have any other tools may be please suggest something config driven much better for llms easily generate rules and we execute them.. WDYT?

On side note at my org we started discussion to use qualink ( i have developed that 😄 so not insisting that to use ) it performs nicely llm generates rules and can be executed directly on the object stores..

[cetingokhan]

I like the overall direction here, but I wonder if the long-term authoring model should be more config-driven than Python-driven.

Instead of requiring DAG authors to define many checks as Python validators/factories, could the LLM generate a simple JSON/YAML rule spec from the natural-language prompts, and then have the operator execute that spec with a deterministic engine?

An example qualink https://github.com/gopidesupavan/qualink and https://gopidesupavan.github.io/qualink/guide/yaml-config/

Not sure similarly may be great expection may have something like that..

So that based on the prompts llm generates rules and we can simply execute on them?

I do think a config-first layer for common checks may scale better than growing the Python validator API. Python validators could still remain as an escape hatch for advanced/custom cases.

If you have any other tools may be please suggest something config driven much better for llms easily generate rules and we execute them.. WDYT?

On side note at my org we started discussion to use qualink ( i have developed that 😄 so not insisting that to use ) it performs nicely llm generates rules and can be executed directly on the object stores..

Hi @gopidesupavan,

Thanks for the feedback! I’ve spent some time thinking about the config-driven approach versus the current implementation. I’ve checked out Qualink as well—congratulations on that, using DataFusion is a very smart move for performance.

Regarding the operator's direction, I actually started developing this by following your comment on the AIP-99 proposal, where you mentioned that users would express requirements in a prompt, and we would generate queries, translate them into internal rules, and execute them. That’s why I included the execution logic directly—to provide that seamless, all-in-one experience.

In my previous LLM-based DQ projects, I’ve found that forcing an LLM to generate complex YAML or JSON schemas often leads to unexpected hallucinations. Moreover, while dealing with large-scale data, I had to rely on Spark Dataframes via Great Expectations (GEX) to handle the load, and as the rule set grows, it becomes increasingly difficult for the LLM to map natural language to the exact parameters of a specific tool's config. On the other hand, generating SQL is generally more deterministic and reliable for the LLM to handle consistently across different scales.

However, if you have found a reliable way to solve this mapping challenge in your Org, I’d love to hear more about your methodology. If we can consistently avoid hallucinations while generating these rule specs, I’d be more than happy to re-evaluate this and integrate a more config-driven layer.

I believe this operator should stay lean and accessible as a "first-step" tool. By focusing on SQL generation, we keep it engine-agnostic and reduce the barrier to entry for users who don't want to manage a secondary DQ platform immediately. If a user’s needs grow to an enterprise scale where they require complex rule management, they might be better served by a dedicated provider or a specialized engine.

For this specific PR, my suggestion is to keep the initial version simple and focused on high-reliability SQL generation, aligning with the execution flow proposed in AIP-99. We could potentially add an "export" feature later to bridge it with tools like Qualink, SODA or GEX, but I think starting with a deterministic, SQL-first approach will provide a much more stable experience for Airflow users right out of the box, allowing Airflow users to meet most of their needs directly while enabling them to easily add minor custom check extensions for simpler requirements.

I'm ready to pivot the development based on your final call. Since you are leading the direction here, I'll follow your guidance, please let me know how you’d like me to proceed, and I’ll be happy to implement it quickly. :)

Regarding Qualink, I actually think it has great potential for big data scales; for instance, integrating DataFusion-Comet could be a game-changer for Spark environments, and I’d honestly love to support you on that front in the future!

[cetingokhan]

Hi @gopidesupavan and @kaxil,
Since we are at a crossroads between a Python/SQL-driven model and a config-driven (YAML/JSON) approach, I’ve paused development to avoid rework. I’m waiting for your guidance on the final architectural direction before I proceed with the implementation.
Looking forward to your thoughts!

[gopidesupavan]

@cetingokhan sorry for the late reply , i will review this weekend and get back to you.. am currently bit busy with my day job..

[Copilot]

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 5 comments.

[Copilot]

Pull request overview

Copilot reviewed 25 out of 25 changed files in this pull request and generated 10 comments.

[Copilot]

Pull request overview

Copilot reviewed 26 out of 26 changed files in this pull request and generated 12 comments.

[gopidesupavan]

Hey @cetingokhan we are getting new DQ provider #69575 lets wait for this to merge then this pr going to be easy ..llms can use that provider to generate rules . the dq provider comes with skills https://github.com/gopidesupavan/airflow/blob/1047525a4d87eb127af372461441381a95da3b48/providers/common/dataquality/src/airflow/providers/common/dataquality/skills/dataquality-rule-authoring/SKILL.md. so that the agent operator / llm operators should be able to use these skills and generate rules.