You are viewing a plain text version of this content. The canonical link for it is here.

Posted to commits@airflow.apache.org by "cetingokhan (via GitHub)" <gi...@apache.org> on 2026/03/05 21:27:24 UTC

[PR] Aip 99 llmdataqualityoperator [airflow]

cetingokhan opened a new pull request, #62963:
URL: https://github.com/apache/airflow/pull/62963

## AIP - 99 LLMDataQualityOperator

This pull request introduces a new **LLMDataQualityOperator** for generating and executing data-quality checks using natural language prompts and LLMs, along with supporting utilities for database/schema introspection and example usage. The changes add a robust operator for data-quality validation, enable schema context resolution for both relational and object-storage sources.

### How It Works
**Plan Generation (LLM-backed)**: The operator accepts a prompts dict mapping check names to natural-language expectations (e.g. "email_nulls": "Less than 5% of emails should be null"). It introspects the target database schema and sends prompts + schema context to the configured LLM, which produces a DQPlan — a set of optimised SQL query groups.

**Plan Caching**: Generated plans are serialised and stored in Airflow Variable (key: dq_plan_<version>_<sha256[:16]>). Cache key is computed from a sorted serialisation of prompts + prompt_version, making it order-independent and version-bumped when prompts change semantically. This avoids redundant LLM calls on rerun.

**Execution**: Each SQL group is executed against the target DB via a DbApiHook. Results are collected per check name into a results_map.

**Validation**: Each metric value is passed to the corresponding callable in validators. A check passes if no validator is provided (metrics are collected but not gated) or if the validator returns True. Failures record the reason.

**Dry Run Mode**: When dry_run=True, the plan is generated/cached but not executed.

<!--
Thank you for contributing!

Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/

Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.

Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.

In case of an existing issue, reference it using one of the following:

* closes: #ISSUE
* related: #ISSUE
-->

---

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

- [X] Yes
Cloude Sonnet 4.6 & Gemini 3.1 Pro
Filled some of methods scope and tests created via copilot

---

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3030608302


##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_models.py:
##########
@@ -0,0 +1,227 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Pydantic models for the LLM data quality plan and result reporting.
+
+``DQPlan`` is the structured output type requested from the LLM.  The remaining
+dataclasses hold execution results and are never serialised back to the model.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from pydantic import BaseModel, Field, computed_field
+
+
+class DQCheck(BaseModel):
+    """
+    A single data-quality check produced by the LLM.
+
+    :param check_name: Matches the key supplied by the user in ``prompts``.
+    :param metric_key: Column alias used in the generated SQL (e.g. ``null_email_count``).
+        The operator reads ``row[metric_key]`` from the query result.
+    :param group_id: Logical bucket for grouping checks into a single SQL query
+        (e.g. ``customers_null_check_1``, ``orders_validity_1``).
+    :param check_category: Semantic category assigned by the LLM based on the check
+        description.  Used to sub-group checks within a table so that checks of
+        different natures (e.g. null-checks vs. regex validity) land in separate
+        SQL queries.  Allowed values: ``null_check``, ``uniqueness``, ``validity``,
+        ``numeric_range``, ``row_count``, ``string_format``.

Review Comment:
   The `DQCheck` class docstring’s `:param check_category:` section lists allowed values but omits the newly supported `row_level` category (even though the Field description includes it and the planner prompt uses it). Please update the docstring list to include `row_level` to avoid misleading API/docs for consumers and LLM prompt authors.
   ```suggestion
           ``numeric_range``, ``row_count``, ``string_format``, ``row_level``.
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,616 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import AirflowException, Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises AirflowException: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises AirflowException: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise AirflowException(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        plan_hash = _compute_plan_hash(self.prompts, self.prompt_version, self.collect_unexpected)
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, default=None)
+        if cached_json is not None:
+            self.log.info("DQ plan cache hit — key: %r", variable_key)
+            plan = DQPlan.model_validate_json(cached_json)
+            if not plan.plan_hash:
+                plan.plan_hash = plan_hash
+            return plan
+
+        self.log.info("DQ plan cache miss — generating via LLM (key: %r).", variable_key)
+        plan = planner.generate_plan(self.prompts, schema_ctx)
+        plan.plan_hash = plan_hash
+        Variable.set(variable_key, plan.model_dump_json())
+        return plan
+
+    def _validate_results(
+        self,
+        results_map: dict[str, Any],
+        plan: DQPlan,
+    ) -> list[DQCheckResult]:
+        """
+        Apply validators to each metric value and return per-check results.
+
+        For aggregate checks each validator callable receives the raw metric
+        value returned by the database.  For row-level checks, where *value* is
+        a :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`,
+        the pass/fail decision compares ``invalid_pct`` against the validator's
+        ``_max_invalid_pct`` attribute (defaulting to ``0.0`` when absent).
+        Checks without a registered validator are logged and marked as passed.
+
+        :param results_map: ``{check_name: metric_value_or_RowLevelResult}`` as
+            returned by
+            :meth:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner.execute_plan`.
+        :param plan: The DQ plan whose groups and checks drive iteration order.
+        :returns: Per-check
+            :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckResult`
+            list in plan-group order.
+        """
+        check_results: list[DQCheckResult] = []
+
+        for group in plan.groups:
+            for check in group.checks:

Review Comment:
   `_validate_results()` indexes `results_map[check.check_name]` directly. If the planner returns a partial results map (e.g. due to a row-level group being skipped or a missing metric key), this becomes a raw `KeyError` without context. Consider using `results_map.get(...)` and raising a `ValueError`/`AirflowException` that names the missing `check_name` and the group/query that produced it, so failures are diagnosable.
   ```suggestion
               for check in group.checks:
                   if check.check_name not in results_map:
                       group_name = getattr(group, "group_name", None)
                       query = getattr(group, "query", None)
                       group_context = f"group={group_name!r}, query={query!r}" if group_name is not None else f"query={query!r}"
                       raise AirflowException(
                           f"Missing result for check_name={check.check_name!r} while validating data quality plan "
                           f"({group_context})."
                       )
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,836 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        active_checks = [
+            check for check in group.checks if check.row_level and check.check_name in self._row_validators
+        ]
+        for check in group.checks:
+            if check.row_level and check.check_name not in self._row_validators:
+                log.warning("No row-level validator found for check %r — skipping.", check.check_name)
+
+        if not active_checks:
+            return {}
+
+        # counters[check_name] = [total, invalid, sample_violations]
+        counters: dict[str, list[Any]] = {check.check_name: [0, 0, []] for check in active_checks}
+
+        chunk_iter: Iterator[list[dict[str, Any]]]
+        if datafusion_engine is not None:
+            chunk_iter = self._iter_datafusion_row_chunks(datafusion_engine, group.query)
+        else:
+            chunk_iter = self._iter_db_row_chunks(group.query)
+
+        total_rows_seen = 0
+        for chunk in chunk_iter:
+            total_rows_seen += len(chunk)
+            for row in chunk:
+                for check in active_checks:
+                    value = row.get(check.metric_key)
+                    c = counters[check.check_name]
+                    c[0] += 1  # total
+                    try:
+                        passed = bool(self._row_validators[check.check_name](value))
+                    except Exception:
+                        # Treat any exception from the validator as a validation failure.
+                        passed = False
+                    if not passed:
+                        c[1] += 1  # invalid
+                        if len(c[2]) < _MAX_VIOLATION_SAMPLES:
+                            c[2].append(str(value))
+
+        if total_rows_seen == 0:
+            log.warning("Row-level query for group %r returned no rows.", group.group_id)
+
+        results: dict[str, RowLevelResult] = {}
+        for check in active_checks:
+            total, invalid, violations = counters[check.check_name]
+            invalid_pct = invalid / total if total else 0.0
+            results[check.check_name] = RowLevelResult(
+                check_name=check.check_name,
+                total=total,
+                invalid=invalid,
+                invalid_pct=invalid_pct,
+                sample_violations=violations,
+                sample_size=_MAX_VIOLATION_SAMPLES if total else 0,
+            )
+            log.info(
+                "Row-level check %r: %d/%d invalid (%.4f%%)",
+                check.check_name,
+                invalid,
+                total,
+                invalid_pct * 100,
+            )
+
+        return results
+
+    def _iter_db_row_chunks(self, query: str) -> Iterator[list[dict[str, Any]]]:
+        """
+        Execute *query* via the DB hook and yield chunks of row-dicts.
+
+        Uses ``cursor.fetchmany(_ROW_LEVEL_CHUNK_SIZE)`` so at most
+        :data:`_ROW_LEVEL_CHUNK_SIZE` raw rows are held in memory at once.
+        Column names are resolved from ``cursor.description``.
+        """
+        with closing(self._db_hook.get_conn()) as conn:  # type: ignore[union-attr]
+            with closing(conn.cursor()) as cur:
+                cur.execute(query)
+                col_names = [str(desc[0]) for desc in cur.description] if cur.description else []
+                while True:
+                    chunk = cur.fetchmany(_ROW_LEVEL_CHUNK_SIZE)
+                    if not chunk:
+                        break
+                    result_chunk: list[dict[str, Any]] = []
+                    for raw_row in chunk:
+                        if isinstance(raw_row, dict):
+                            result_chunk.append(raw_row)
+                        elif isinstance(raw_row, Sequence) and not isinstance(
+                            raw_row, str | bytes | bytearray
+                        ):
+                            result_chunk.append(dict(zip(col_names, raw_row)))
+                        else:
+                            result_chunk.append({})
+                    yield result_chunk

Review Comment:
   `_iter_db_row_chunks()` silently converts any unexpected row shape into an empty dict (`{}`) and also uses `dict(zip(col_names, raw_row))` without verifying that `cursor.description` exists and matches the tuple length. This can silently drop columns (or produce empty rows) and cause row-level validators to run against `None` values instead of the intended column. Please raise a clear exception for unsupported row types and for tuple/sequence rows when `cursor.description` is missing or the column count doesn’t match the row length, so query/driver issues don’t turn into incorrect DQ results.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,616 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import AirflowException, Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises AirflowException: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises AirflowException: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise AirflowException(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        plan_hash = _compute_plan_hash(self.prompts, self.prompt_version, self.collect_unexpected)
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, default=None)
+        if cached_json is not None:
+            self.log.info("DQ plan cache hit — key: %r", variable_key)
+            plan = DQPlan.model_validate_json(cached_json)
+            if not plan.plan_hash:
+                plan.plan_hash = plan_hash
+            return plan
+
+        self.log.info("DQ plan cache miss — generating via LLM (key: %r).", variable_key)
+        plan = planner.generate_plan(self.prompts, schema_ctx)
+        plan.plan_hash = plan_hash
+        Variable.set(variable_key, plan.model_dump_json())
+        return plan
+
+    def _validate_results(
+        self,
+        results_map: dict[str, Any],
+        plan: DQPlan,
+    ) -> list[DQCheckResult]:
+        """
+        Apply validators to each metric value and return per-check results.
+
+        For aggregate checks each validator callable receives the raw metric
+        value returned by the database.  For row-level checks, where *value* is
+        a :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`,
+        the pass/fail decision compares ``invalid_pct`` against the validator's
+        ``_max_invalid_pct`` attribute (defaulting to ``0.0`` when absent).
+        Checks without a registered validator are logged and marked as passed.
+
+        :param results_map: ``{check_name: metric_value_or_RowLevelResult}`` as
+            returned by
+            :meth:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner.execute_plan`.
+        :param plan: The DQ plan whose groups and checks drive iteration order.
+        :returns: Per-check
+            :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckResult`
+            list in plan-group order.
+        """
+        check_results: list[DQCheckResult] = []
+
+        for group in plan.groups:
+            for check in group.checks:
+                value = results_map[check.check_name]
+                validator = self.validators.get(check.check_name)
+
+                passed = True
+                failure_reason: str | None = None
+
+                if isinstance(value, RowLevelResult):
+                    # Row-level check: evaluate threshold against invalid_pct.
+                    max_pct = getattr(validator, "_max_invalid_pct", 0.0)
+                    passed = value.invalid_pct <= max_pct
+                    if not passed:
+                        failure_reason = (
+                            f"Row-level check failed: {value.invalid}/{value.total} rows invalid "
+                            f"({value.invalid_pct:.4%}), threshold {max_pct:.4%}"
+                        )

Review Comment:
   Row-level checks are always evaluated against `max_pct = getattr(validator, "_max_invalid_pct", 0.0)` even when `validator` is `None`. This makes “no validator provided ⇒ check passes by default” true for aggregate checks but **not** for row-level checks (a missing validator effectively enforces a 0% invalid threshold). Consider explicitly handling `validator is None` in the `RowLevelResult` branch the same way as aggregate checks (warn + mark passed), and/or raising earlier if a row-level result is returned for a check that has no corresponding validator callable.



##########
providers/common/ai/tests/unit/common/ai/utils/test_dq_planner.py:
##########
@@ -0,0 +1,1174 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.utils.dq_models import DQCheck, DQCheckGroup, DQPlan
+from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+from airflow.providers.common.ai.utils.sql_validation import SQLSafetyError
+
+
+def _make_plan(*check_names: str) -> DQPlan:
+    """Helper: build a minimal DQPlan with one group per check."""
+    groups = [
+        DQCheckGroup(
+            group_id="numeric_aggregate",
+            query=f"SELECT COUNT(*) AS {name}_count FROM t",
+            checks=[DQCheck(check_name=name, metric_key=f"{name}_count", group_id="numeric_aggregate")],
+        )
+        for name in check_names
+    ]
+    return DQPlan(groups=groups)
+
+
+def _make_llm_hook(plan: DQPlan) -> MagicMock:
+    """Helper: mock PydanticAIHook that returns *plan* from agent.run_sync."""
+    mock_usage = MagicMock(requests=1, tool_calls=0, input_tokens=100, output_tokens=50, total_tokens=150)
+    mock_result = MagicMock(spec=["output", "all_messages", "usage", "response"])
+    mock_result.output = plan
+    mock_result.all_messages.return_value = []
+    mock_result.usage.return_value = mock_usage
+    mock_result.response.model_name = "test-model"
+    mock_agent = MagicMock(spec=["run_sync"])
+    mock_agent.run_sync.return_value = mock_result
+    mock_hook = MagicMock(spec=PydanticAIHook)
+    mock_hook.create_agent.return_value = mock_agent
+    return mock_hook
+
+
+class TestSQLDQPlannerBuildSchema:
+    def test_returns_manual_schema_context_verbatim(self):
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=None)
+        result = planner.build_schema_context(
+            table_names=None,
+            schema_context="Table: t\nColumns: id INT",
+        )
+        assert result == "Table: t\nColumns: id INT"
+
+    def test_introspects_via_db_hook_when_no_manual_context(self):
+        mock_db_hook = MagicMock()
+        mock_db_hook.get_table_schema.return_value = [{"name": "id", "type": "INT"}]
+
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=mock_db_hook)
+        result = planner.build_schema_context(

Review Comment:
   These tests construct several `MagicMock()` instances without `spec`/`autospec` (e.g. `mock_db_hook = MagicMock()` here). Unspec’d mocks can hide real interface mismatches and make refactors risky. Prefer `MagicMock(spec=...)` (or `patch(..., autospec=True)`) for hooks/agents/cursors so attribute mistakes fail fast.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,836 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue

Review Comment:
   `execute_plan()` treats a group as row-level if **any** check has `row_level=True`, then skips executing the group SQL entirely. If the LLM ever returns a mixed group (some `row_level`, some aggregate) or marks a check `row_level=True` without a corresponding entry in `self._row_validators`, the aggregate checks will be silently dropped and the final `results` map may miss prompt keys (leading to downstream `KeyError`s). Consider validating that groups are homogeneous (all row-level vs none), and raising a clear `ValueError` when a row-level group has zero runnable row validators rather than continuing.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3114032278


##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -119,7 +120,25 @@ def execute_query(self, query: str, max_rows: int | None = None) -> dict[str, li
                 return result
             return df.to_pydict()
         except Exception as e:
-            raise QueryExecutionException(f"Error while executing query: {e}")
+            raise QueryExecutionException(f"Error while executing query: {e}") from e
+
+    def iter_query_row_chunks(self, query: str) -> Iterator[dict[str, list[Any]]]:
+        """
+        Execute *query* and yield one column-dict per RecordBatch, streaming results.
+
+        :param query: SQL SELECT query to execute.
+        :raises QueryExecutionException: On SQL execution errors.
+        """
+        try:
+            df = self.session_context.sql(query)
+            if hasattr(df, "execute_stream"):
+                for batch in df.execute_stream():
+                    yield batch.to_pyarrow().to_pydict()
+            else:
+                for batch in df.collect():
+                    yield batch.to_pyarrow().to_pydict()
+        except Exception as e:
+            raise QueryExecutionException(f"Error while executing query: {e}") from e

Review Comment:
   `iter_query_row_chunks()` is new public behavior on `DataFusionEngine` and is now used for row-level DQ execution, but there are no unit tests covering its streaming/collect fallback behavior or the yielded payload shape. Since this module already has comprehensive unit tests, please add tests for `iter_query_row_chunks()` (including the `execute_stream` path and the `collect` fallback).



##########
providers/common/ai/docs/operators/llm_data_quality.rst:
##########
@@ -0,0 +1,488 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _howto/operator:llm_data_quality:
+
+``LLMDataQualityOperator``
+==========================
+
+Use :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`
+to generate and execute data-quality checks from natural language using an LLM.
+
+Each entry in ``checks`` describes **one** data-quality expectation as a
+:class:`~airflow.providers.common.ai.utils.dq_models.DQCheckInput` object.
+The LLM groups related checks into optimised SQL queries, executes them against the
+target database, and validates each metric.  The task fails if any check does not
+pass, gating downstream tasks on data quality.
+
+.. seealso::
+    :ref:`Connection configuration <howto/connection:pydanticai>`
+
+Basic Usage
+-----------
+
+Provide a ``checks`` list and a target ``db_conn_id``. The operator introspects the
+schema automatically when ``table_names`` is provided:
+
+.. code-block:: python
+
+    from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+    from airflow.providers.common.ai.utils.dq_models import DQCheckInput
+    from airflow.providers.common.ai.utils.dq_validation import null_pct_check, row_count_check
+
+    LLMDataQualityOperator(
+        task_id="validate_orders",
+        llm_conn_id="pydanticai_default",
+        db_conn_id="postgres_default",
+        table_names=["orders", "customers"],
+        checks=[
+            DQCheckInput(
+                name="row_count",
+                description="The orders table must contain at least 1000 rows.",
+                validator=row_count_check(min_count=1000),
+            ),
+            DQCheckInput(
+                name="email_nulls",
+                description="No more than 5% of customer email addresses should be null.",
+                validator=null_pct_check(max_pct=0.05),
+            ),
+        ],
+    )
+
+Validators
+----------
+
+Each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckInput` can carry
+an optional ``validator`` — a callable that receives the raw metric value returned
+by the generated SQL and returns ``True`` (pass) or ``False`` (fail).
+
+Built-in Factories
+~~~~~~~~~~~~~~~~~~
+
+:mod:`~airflow.providers.common.ai.utils.dq_validation` ships ready-made factories
+for the most common thresholds:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 50 20
+
+   * - Factory
+     - Passes when …
+     - Example
+   * - ``null_pct_check(max_pct=…)``
+     - metric ≤ ``max_pct``
+     - ``null_pct_check(max_pct=0.05)``
+   * - ``row_count_check(min_count=…)``
+     - metric ≥ ``min_count``
+     - ``row_count_check(min_count=1000)``
+   * - ``duplicate_pct_check(max_pct=…)``
+     - metric ≤ ``max_pct``
+     - ``duplicate_pct_check(max_pct=0.0)``
+   * - ``between_check(min_val=…, max_val=…)``
+     - ``min_val`` ≤ metric ≤ ``max_val``
+     - ``between_check(min_val=0.0, max_val=1.0)``
+   * - ``exact_check(expected=…)``
+     - metric == ``expected``
+     - ``exact_check(expected=0)``
+
+You can also use plain lambdas for one-off conditions::
+
+    DQCheckInput(
+        name="stale_rows",
+        description="Count rows older than 30 days",
+        validator=lambda v: int(v) < 1000,
+    )
+
+Aggregate checks without a validator are marked as **passed** — metrics are
+still collected and included in the report, but no threshold is enforced.
+Row-level checks are stricter: each row-level check must have a corresponding
+row-level validator, otherwise execution fails fast.
+
+Custom Validators with ``register_validator``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use the :func:`~airflow.providers.common.ai.utils.dq_validation.register_validator`
+decorator to attach an ``llm_context`` hint to your validator factory.  The
+operator injects the hint into the LLM system prompt so the model generates SQL
+that returns the metric format your validator expects:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_custom_validator]
+    :end-before: [END howto_operator_llm_dq_custom_validator]
+

Review Comment:
   The `exampleinclude` directives reference snippet markers (e.g. `howto_operator_llm_dq_custom_validator`, `howto_operator_llm_dq_s3_parquet`, `howto_operator_llm_dq_custom_row_level_validator`, etc.) that do not exist in `example_llm_data_quality.py` (which defines different `[START ...]` tags). This will break the docs build; update the `:start-after:`/`:end-before:` tags (or the example DAG markers) so they match.
   ```suggestion
   .. code-block:: python
   
       @register_validator(llm_context="SQL must return a single percentage between 0 and 1.")
       def max_pct(max_pct: float):
           return lambda value: float(value) <= max_pct
   
       checks = [
           DQCheckInput(
               name="null_pct_check",
               description="Percent of rows with null values",
               validator=max_pct(0.05),
           ),
       ]
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,1122 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import re
+from collections import Counter
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckGroup,
+    DQCheckInput,
+    DQPlan,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import DQValidationToolset
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 name in the provided checks list.
+  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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+_TABLE_NAME_CONSTRAINT_SECTION = """
+TABLE NAME CONSTRAINT:
+  The ONLY table(s) you may reference in FROM clauses are: {table_names}.
+  You MUST use these exact table names in every SQL query you generate.
+  Do NOT rename, abbreviate, alias the table, or invent new table names.
+  Using any table name not listed above will cause a runtime error.
+"""
+
+
+def _extract_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by :func:`build_schema_context`."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    :param toolset: :class:`~airflow.providers.common.ai.utils.dq_validation.DQValidationToolset`
+        that exposes the validator catalog to the LLM and validates its selections.
+        Defaults to a toolset backed by the :data:`~airflow.providers.common.ai.utils.dq_validation.default_registry`.
+    :param fixed_validators: Mapping of ``{check_name: callable}`` for checks where
+        the user preselected a validator.  These checks are excluded from LLM
+        validator-suggestion and the provided callables are used directly.
+    :param max_validator_retries: Maximum number of times the LLM is asked to correct
+        invalid validator selections before the plan is rejected.  Default ``3``.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+        toolset: DQValidationToolset | None = None,
+        fixed_validators: dict[str, Any] | None = None,
+        max_validator_retries: int = 3,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        self._toolset: DQValidationToolset = toolset if toolset is not None else DQValidationToolset()
+        self._fixed_validators: dict[str, Any] = fixed_validators or {}
+        self._max_validator_retries = max_validator_retries
+        self._cached_datafusion_engine: DataFusionEngine | None = None
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def build_catalog_hash(self) -> str:
+        """
+        Return a 16-char SHA-256 fingerprint of the current validator catalog.
+
+        Used by the operator to include a catalog change fingerprint in the plan
+        cache key, ensuring that adding or removing a validator invalidates stale
+        cached plans without requiring a new ``prompt_version``.
+        """
+        section = self._toolset.build_system_prompt_section()
+        return hashlib.sha256(section.encode()).hexdigest()[:16]
+
+    def generate_plan(self, checks: list[DQCheckInput], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the check definitions, schema context, validator catalog,
+        and planning instructions as a structured-output call (``output_type=DQPlan``).
+        In a single response the model produces both the SQL plan and validator
+        selections for each check.
+
+        After generation the method:
+
+        1. Verifies that the returned ``check_names`` exactly match the input names.
+        2. Validates each LLM-suggested validator via the toolset.
+        3. On validation failures, retries up to ``max_validator_retries`` times by
+           continuing the same conversation thread.
+        4. Instantiates validated validator suggestions and merges them with any
+           user-fixed validators.
+
+        :param checks: A list of :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckInput`
+            objects describing the data-quality expectations.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every check name exactly once.
+        :raises ValueError: If validator suggestions cannot be corrected within the retry budget.
+        """
+        check_descriptions = {c.name: c.description for c in checks}
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+            table_names = _extract_table_names(schema_context)
+            if table_names:
+                system_prompt += _TABLE_NAME_CONSTRAINT_SECTION.format(table_names=", ".join(table_names))
+
+        # Inject validator catalog for LLM-driven selection.
+        system_prompt += self._toolset.build_system_prompt_section()
+
+        # Legacy per-validator context (still injected for fixed validators that carry it).
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(checks, fixed_validator_names=set(self._fixed_validators))
+        prompt_fingerprint = hashlib.sha256(f"{system_prompt}\n{user_message}".encode()).hexdigest()[:12]
+
+        log.info(
+            "Generating DQ plan with %d check(s) (system_prompt_chars=%d, user_message_chars=%d, "
+            "prompt_fingerprint=%s).",
+            len(checks),
+            len(system_prompt),
+            len(user_message),
+            prompt_fingerprint,
+        )
+        log.debug("Using system prompt:\n%s", system_prompt)
+        log.debug("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan and validator
+        # correction can continue the same chat thread.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, check_descriptions)
+        self._validate_group_sizes(plan)
+
+        # Validate and potentially retry LLM validator suggestions.
+        plan = self._validate_and_fix_validator_suggestions(plan)
+
+        return plan
+
+    def _validate_and_fix_validator_suggestions(self, plan: DQPlan) -> DQPlan:
+        """
+        Validate LLM-suggested validator names and arguments for each check.
+
+        Checks with a user-fixed validator (in ``self._fixed_validators``) are
+        skipped.  For the rest, each suggestion is tested via
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.DQValidationToolset.validate_suggestion`.
+
+        On failure the LLM is sent a correction prompt in the same conversation
+        thread (up to ``max_validator_retries`` attempts).  If the LLM still
+        returns invalid selections after all retries, an informative
+        :class:`ValueError` is raised.
+
+        :param plan: Plan returned by the initial LLM call.
+        :returns: The plan with validated (and possibly corrected) validator fields.
+        :raises ValueError: If validator suggestions cannot be corrected.
+        """
+        current_plan = plan
+        for attempt in range(1, self._max_validator_retries + 1):
+            errors = self._collect_validator_errors(current_plan)
+            if not errors:
+                break
+
+            # Log all failed checks as errors so they are visible in task logs.
+            for check_name, error_msg in errors.items():
+                log.error(
+                    "Validator suggestion invalid for check %r (attempt %d/%d): %s",
+                    check_name,
+                    attempt,
+                    self._max_validator_retries,
+                    error_msg,
+                )
+
+            if self._plan_agent is None or self._plan_all_messages is None:
+                # No conversation thread to continue — fail immediately.
+                break
+
+            correction_prompt = self._build_validator_correction_prompt(errors)
+            log.info(
+                "Sending validator correction prompt (attempt %d/%d).",
+                attempt,
+                self._max_validator_retries,
+            )
+            result = self._plan_agent.run_sync(
+                correction_prompt,
+                message_history=self._plan_all_messages,
+            )
+            self._plan_all_messages = result.all_messages()
+            current_plan = result.output
+
+        # Final check after loop.
+        errors = self._collect_validator_errors(current_plan)
+        if errors:
+            error_lines = [f"  - {name}: {msg}" for name, msg in sorted(errors.items())]
+            raise ValueError(
+                f"LLM validator suggestions could not be corrected after "
+                f"{self._max_validator_retries} attempt(s). "
+                f"Invalid suggestions:\n" + "\n".join(error_lines)
+            )
+
+        return current_plan
+
+    def _collect_validator_errors(self, plan: DQPlan) -> dict[str, str]:
+        """
+        Return a ``{check_name: error_message}`` dict for every invalid validator suggestion.
+
+        Checks with user-fixed validators are excluded from validation.
+        """
+        errors: dict[str, str] = {}
+        for group in plan.groups:
+            for check in group.checks:
+                if check.check_name in self._fixed_validators:
+                    continue
+                # For non-fixed checks, missing validator_name is invalid and must
+                # trigger correction/failure.
+                suggested_name = check.validator_name or ""
+                ok, msg = self._toolset.validate_suggestion(
+                    check.check_name,
+                    suggested_name,
+                    check.validator_args,
+                )

Review Comment:
   `_collect_validator_errors()` flags missing/"none" validator_name as invalid for every non-fixed check, which forces the LLM to always select a validator even though the prompt text allows "none" and the operator treats aggregate checks without validators as pass-through metric collection. This will either cause repeated correction loops or hard failures for checks where "none" is appropriate. Consider allowing aggregate checks to omit a validator (and only requiring a validator for row-level checks).



##########
providers/common/ai/src/airflow/providers/common/ai/decorators/llm_data_quality.py:
##########
@@ -0,0 +1,164 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+TaskFlow decorator for LLM-driven data-quality checks.
+
+The user writes a function that **returns a list of
+:class:`~airflow.providers.common.ai.utils.dq_models.DQCheckInput`** objects.
+The decorator handles LLM plan generation, plan caching, SQL execution against
+the target database, and metric validation.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Collection, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.compat.sdk import (
+    DecoratedOperator,
+    TaskDecorator,
+    context_merge,
+    determine_kwargs,
+    task_decorator_factory,
+)
+from airflow.sdk.definitions._internal.types import SET_DURING_EXECUTION
+
+if TYPE_CHECKING:
+    from airflow.sdk import Context
+
+
+class _LLMDQDecoratedOperator(DecoratedOperator, LLMDataQualityOperator):
+    """
+    Wraps a callable that returns a list of :class:`DQCheckInput` for LLM data-quality checks.
+
+    The user function is called at execution time to produce the checks list.
+    All other parameters (``llm_conn_id``, ``db_conn_id``, ``table_names``,
+    etc.) are passed through to
+    :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`.
+
+    :param python_callable: A callable that returns a
+        ``list[DQCheckInput]``.
+    :param op_args: Positional arguments for the callable.
+    :param op_kwargs: Keyword arguments for the callable.
+    """
+
+    template_fields: Sequence[str] = (
+        *DecoratedOperator.template_fields,
+        *LLMDataQualityOperator.template_fields,
+    )
+    template_fields_renderers: ClassVar[dict[str, str]] = {
+        **DecoratedOperator.template_fields_renderers,
+    }
+
+    custom_operator_name: str = "@task.llm_dq"
+
+    def __init__(
+        self,
+        *,
+        python_callable: Callable,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            python_callable=python_callable,
+            op_args=op_args,
+            op_kwargs=op_kwargs,
+            checks=SET_DURING_EXECUTION,
+            **kwargs,
+        )
+
+    def execute(self, context: Context) -> Any:
+        context_merge(context, self.op_kwargs)
+        kwargs = determine_kwargs(self.python_callable, self.op_args, context)
+
+        checks = self.python_callable(*self.op_args, **kwargs)
+
+        if not isinstance(checks, list) or not checks:
+            raise TypeError(
+                "The returned value from the @task.llm_dq callable must be a non-empty list[DQCheckInput]."
+            )
+
+        self.checks = checks
+        self.render_template_fields(context)
+        return LLMDataQualityOperator.execute(self, context)

Review Comment:
   `_LLMDQDecoratedOperator.execute()` assigns `self.checks` directly from the user callable without coercing dicts to `DQCheckInput` (unlike `LLMDataQualityOperator.__init__`) and without running the duplicate-name/empty-field validation that the operator normally does. This can lead to runtime attribute errors (e.g. when computing the plan hash) or allow duplicate check names through. Consider coercing each item via `DQCheckInput.coerce(...)` and calling the operator’s `_validate_checks()` after setting `self.checks`.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/db_schema.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Shared database hook and schema introspection utilities.
+
+These helpers are used by both :class:`~airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator`
+and :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator` to
+avoid code duplication while keeping both operators decoupled from each other.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from airflow.providers.common.compat.sdk import BaseHook
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+# SQLAlchemy dialect_name → sqlglot dialect mapping for names that differ.
+SQLALCHEMY_TO_SQLGLOT_DIALECT: dict[str, str] = {
+    "postgresql": "postgres",
+    "mssql": "tsql",
+}
+
+
+def get_db_hook(db_conn_id: str) -> DbApiHook:
+    """
+    Return a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` for *db_conn_id*.
+
+    :param db_conn_id: Airflow connection ID that resolves to a ``DbApiHook``.
+    :raises ValueError: If the connection does not resolve to a ``DbApiHook``.
+    """
+    # Lazy load to avoid hard dependency on common.sql
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+    connection = BaseHook.get_connection(db_conn_id)
+    hook = connection.get_hook()
+    if not isinstance(hook, DbApiHook):
+        raise ValueError(
+            f"Connection {db_conn_id!r} does not provide a DbApiHook. Got {type(hook).__name__}."
+        )
+    return hook
+
+
+def resolve_dialect(db_hook: DbApiHook | None, explicit_dialect: str | None) -> str | None:
+    """
+    Resolve the SQL dialect from an explicit parameter or a database hook.
+
+    Normalises SQLAlchemy dialect names to sqlglot equivalents
+    (e.g. ``postgresql`` → ``postgres``).
+
+    :param db_hook: Database hook to read ``dialect_name`` from when *explicit_dialect* is absent.
+    :param explicit_dialect: Caller-supplied dialect string; takes priority over the hook.
+    :return: Resolved dialect string, or ``None`` when neither source provides one.
+    """
+    raw = explicit_dialect
+    if not raw and db_hook and hasattr(db_hook, "dialect_name"):
+        candidate = db_hook.dialect_name
+        raw = candidate if isinstance(candidate, str) else None
+    if raw:
+        return SQLALCHEMY_TO_SQLGLOT_DIALECT.get(raw, raw)
+    return None
+
+
+def build_schema_context(
+    *,
+    db_hook: DbApiHook | None,
+    table_names: list[str] | None,
+    schema_context: str | None,
+    datasource_config: DataSourceConfig | None,
+) -> str:
+    """
+    Return a schema description string suitable for inclusion in an LLM prompt.
+
+    Resolution order:
+    1. *schema_context* — returned as-is when provided (manual override).
+    2. DB introspection via *db_hook* + *table_names*.
+    3. Object-storage introspection via *datasource_config*.
+    4. Empty string when none of the above are available.
+
+    :param db_hook: Hook used for relational-database schema introspection.
+    :param table_names: Table names to introspect via *db_hook*.
+    :param schema_context: Manual schema description; bypasses introspection when set.
+    :param datasource_config: DataFusion datasource config for object-storage schema.
+    :raises ValueError: If *table_names* are provided but none yield schema information.
+    """
+    if schema_context:
+        return schema_context
+
+    if (db_hook and table_names) or datasource_config:
+        return _introspect_schemas(
+            db_hook=db_hook,
+            table_names=table_names,
+            datasource_config=datasource_config,
+        )
+
+    return ""

Review Comment:
   `build_schema_context()` returns an empty string when `table_names` is provided but `db_hook` and `datasource_config` are both missing, so a misconfiguration silently disables schema introspection. Since `_introspect_schemas()` explicitly raises for `table_names` without a DB hook, consider raising a `ValueError` in `build_schema_context()` whenever `table_names` is set but `db_hook` is None (unless `schema_context` is provided).



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_validation.py:
##########
@@ -0,0 +1,632 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Built-in and custom validator factories for :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`.
+
+Each factory returns a ``Callable[[Any], bool]`` and can be passed as the
+``validator`` argument of a :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckInput`
+to force a specific validator for that check, bypassing LLM selection.
+Factories are intentionally decoupled from the operator so they can be tested and
+composed independently.
+
+Custom validators registered with :func:`register_validator` are exposed to the LLM
+via the validator catalog so the model can select them automatically.
+
+Usage::
+
+    from airflow.providers.common.ai.utils.dq_models import DQCheckInput
+    from airflow.providers.common.ai.utils.dq_validation import (
+        null_pct_check,
+        row_count_check,
+        register_validator,
+    )
+
+    # Fixed validators — LLM is not asked to select a validator for these checks.
+    checks = [
+        DQCheckInput(
+            name="email_nulls",
+            description="Check for null emails",
+            validator=null_pct_check(max_pct=0.05),
+        ),
+        DQCheckInput(
+            name="min_customers",
+            description="Ensure at least 1000 rows",
+            validator=row_count_check(min_count=1000),
+        ),
+    ]
+
+
+    # Custom validator with LLM context — LLM can select this automatically.
+    @register_validator(
+        "freshness_check",
+        llm_context=(
+            "Compute hours since the most recent row. "
+            "SQL pattern: EXTRACT(EPOCH FROM (NOW() - MAX(ts_col))) / 3600.0. "
+            "Returns a DOUBLE representing hours elapsed."
+        ),
+        check_category="freshness",
+    )
+    def freshness_check(*, max_hours: float):
+        def _check(value):
+            return float(value) <= max_hours
+
+        return _check
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+log = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ValidatorEntry:
+    """
+    Metadata for a registered validator factory.
+
+    :param factory: Callable that returns a ``Callable[[Any], bool]`` validator.
+    :param llm_context: Optional hint injected into the LLM system prompt so
+        the model knows what SQL metric format this validator expects.
+    :param check_category: Optional custom check category.  When set, the LLM
+        is instructed to use this category for grouping.
+    :param row_level: When ``True`` the LLM is instructed to generate a plain
+        ``SELECT pk, col FROM table`` (no aggregation).  The planner fetches
+        every row and applies the validator callable to each column value,
+        then reports ``{total, invalid, invalid_pct, sample_violations}``.
+    """
+
+    factory: Callable[..., Callable[[Any], bool]]
+    llm_context: str = ""
+    check_category: str = ""
+    row_level: bool = False
+
+
+class ValidatorRegistry:
+    """
+    Registry for reusable validator factories with optional LLM context.
+
+    Validators registered here can carry an ``llm_context`` string that the
+    operator automatically injects into the LLM system prompt, guiding the
+    model to produce SQL that returns the metric format the validator expects.
+
+    A module-level :data:`default_registry` instance is available.  Use the
+    convenience decorator :func:`register_validator` to register into it.
+    """
+
+    def __init__(self) -> None:
+        self._entries: dict[str, ValidatorEntry] = {}
+
+    def register(
+        self,
+        name: str,
+        *,
+        llm_context: str = "",
+        check_category: str = "",
+        row_level: bool = False,
+    ) -> Callable[[Callable[..., Callable[[Any], bool]]], Callable[..., Callable[[Any], bool]]]:
+        """
+        Return a decorator that registers a validator factory under *name*.
+
+        :param name: Unique name for this validator.
+        :param llm_context: SQL generation hint injected into the LLM prompt.
+        :param check_category: Custom check category for LLM grouping.
+        :param row_level: When ``True``, the LLM generates a plain SELECT
+            returning raw row values instead of an aggregate query.  The
+            planner applies the validator to each row and aggregates results.
+        :raises ValueError: If *name* is already registered.
+        """
+        if name in self._entries:
+            raise ValueError(
+                f"Validator {name!r} is already registered. "
+                "Use a different name or unregister the existing one first."
+            )
+
+        def _decorator(
+            factory: Callable[..., Callable[[Any], bool]],
+        ) -> Callable[..., Callable[[Any], bool]]:
+            # Wrap the factory so every closure it returns carries introspection
+            # attributes used by the operator and planner.
+            @functools.wraps(factory)
+            def _wrapped_factory(*args: Any, **kwargs: Any) -> Callable[[Any], bool]:
+                closure = factory(*args, **kwargs)
+                arg_parts = [repr(a) for a in args]
+                kwarg_parts = [f"{k}={v!r}" for k, v in sorted(kwargs.items())]
+                call_str = f"{name}({', '.join(arg_parts + kwarg_parts)})"
+                if not hasattr(closure, "_validator_name"):
+                    closure._validator_name = name  # type: ignore[attr-defined]
+                if not hasattr(closure, "_row_level"):
+                    closure._row_level = row_level  # type: ignore[attr-defined]
+                if not hasattr(closure, "_validator_display"):
+                    closure._validator_display = call_str  # type: ignore[attr-defined]
+                for k, v in sorted(kwargs.items()):
+                    if not hasattr(closure, f"_{k}"):
+                        setattr(closure, f"_{k}", v)  # e.g. _max_pct, _min_count
+                return closure
+
+            _wrapped_factory._validator_name = name  # type: ignore[attr-defined]
+            _wrapped_factory._llm_context = llm_context  # type: ignore[attr-defined]
+            _wrapped_factory._check_category = check_category  # type: ignore[attr-defined]
+            _wrapped_factory._row_level = row_level  # type: ignore[attr-defined]
+            _wrapped_factory._validator_display = name  # type: ignore[attr-defined]
+            _wrapped_factory.__name__ = factory.__name__
+            _wrapped_factory.__qualname__ = factory.__qualname__
+            _wrapped_factory.__doc__ = factory.__doc__
+
+            self._entries[name] = ValidatorEntry(
+                factory=_wrapped_factory,
+                llm_context=llm_context,
+                check_category=check_category,
+                row_level=row_level,
+            )
+            return _wrapped_factory
+
+        return _decorator
+
+    def get(self, name: str) -> ValidatorEntry:
+        """
+        Return the :class:`ValidatorEntry` for *name*.
+
+        :raises KeyError: If *name* is not registered.
+        """
+        try:
+            return self._entries[name]
+        except KeyError:
+            raise KeyError(
+                f"Validator {name!r} is not registered. Available validators: {sorted(self._entries)}"
+            ) from None
+
+    def list_validators(self) -> list[str]:
+        """Return sorted list of all registered validator names."""
+        return sorted(self._entries)
+
+    def is_row_level(self, validator: Callable[[Any], bool]) -> bool:
+        """
+        Return ``True`` when *validator* was produced by a row-level factory.
+
+        Checks the ``_row_level`` attribute set by the factory closure and,
+        as a fallback, the registry entry for the factory name.
+        """
+        if hasattr(validator, "_row_level"):
+            return bool(validator._row_level)
+        factory_name: str | None = getattr(validator, "_validator_name", None)
+        if factory_name and factory_name in self._entries:
+            return self._entries[factory_name].row_level
+        return False
+
+    def build_llm_context(self, validators: dict[str, Callable[[Any], bool]]) -> str:
+        """
+        Collect LLM context strings from all validators that carry one.
+
+        Aggregate and row-level validators are emitted in separate sections so
+        the LLM knows which checks require raw-row SELECTs vs. aggregate queries.
+
+        Checks three sources in order for each validator callable:
+
+        1. Registry entry (if the callable's factory was registered).
+        2. ``_llm_context`` attribute on the callable itself.
+        3. ``llm_context`` attribute on the callable itself.
+
+        :param validators: The ``{check_name: callable}`` dict from the operator.
+        :returns: Combined context string ready for injection into the system prompt,
+            or empty string if no validator carries context.
+        """
+        aggregate_lines: list[str] = []
+        row_level_lines: list[str] = []
+
+        for check_name, validator in validators.items():
+            context = self._resolve_llm_context(validator)
+            if not context:
+                continue
+            if self.is_row_level(validator):
+                row_level_lines.append(f"  - {check_name}: {context}")
+            else:
+                aggregate_lines.append(f"  - {check_name}: {context}")
+
+        if not aggregate_lines and not row_level_lines:
+            return ""
+
+        parts: list[str] = []
+        if aggregate_lines:
+            parts.append(
+                "\nCUSTOM VALIDATOR CONTEXT:\n"
+                "  The following checks have specific metric requirements.\n"
+                "  Generate SQL that returns values matching these descriptions:\n"
+                + "\n".join(aggregate_lines)
+            )
+        if row_level_lines:
+            parts.append(
+                "\nROW-LEVEL CHECKS:\n"
+                "  The following checks require ROW-LEVEL validation.\n"
+                "  For each, generate a SELECT that returns the primary key column(s)\n"
+                "  and the column(s) to validate — do NOT aggregate.\n"
+                "  Set check.row_level = true on these DQCheck entries.\n"
+                "  The Python-side validator will inspect each returned value:\n" + "\n".join(row_level_lines)
+            )
+        return "\n".join(parts) + "\n"
+
+    def _resolve_llm_context(self, validator: Callable[[Any], bool]) -> str:
+        """Resolve LLM context from registry entries or callable attributes."""
+        # Check registry by factory name attribute.
+        factory_name: str | None = getattr(validator, "_validator_name", None)
+        if factory_name and factory_name in self._entries:
+            entry = self._entries[factory_name]
+            if entry.llm_context:
+                return entry.llm_context
+
+        # Fallback: attribute on the callable itself.
+        for attr in ("_llm_context", "llm_context"):
+            context = getattr(validator, attr, None)
+            if context and isinstance(context, str):
+                return context
+
+        return ""
+
+    def unregister(self, name: str) -> None:
+        """
+        Remove a validator from the registry.
+
+        :raises KeyError: If *name* is not registered.
+        """
+        try:
+            del self._entries[name]
+        except KeyError:
+            raise KeyError(f"Validator {name!r} is not registered.") from None
+
+
+default_registry = ValidatorRegistry()
+
+
+def register_validator(
+    name: str,
+    *,
+    llm_context: str = "",
+    check_category: str = "",
+    row_level: bool = False,
+) -> Callable[[Callable[..., Callable[[Any], bool]]], Callable[..., Callable[[Any], bool]]]:
+    """
+    Register a validator factory in the :data:`default_registry`.
+
+    Use as a decorator on a factory function::
+
+        @register_validator(
+            "freshness_check",
+            llm_context="Compute hours since MAX(updated_at). Returns DOUBLE.",
+            check_category="freshness",
+        )
+        def freshness_check(*, max_hours: float):
+            def _check(value):
+                return float(value) <= max_hours
+
+            return _check
+
+    For row-level validation (e.g. TCKN formula)::
+
+        @register_validator(
+            "tckn_check",
+            llm_context="ROW-LEVEL: SELECT pk, tckn_col FROM table. No aggregation.",
+            check_category="row_level",
+            row_level=True,
+        )
+        def tckn_check(*, max_invalid_pct: float = 0.0):
+            def _check_row(value): ...
+
+            return _check_row
+
+    :param name: Unique name for this validator.
+    :param llm_context: SQL generation hint injected into the LLM prompt.
+    :param check_category: Custom check category for LLM grouping.
+    :param row_level: When ``True``, the LLM generates a plain SELECT returning
+        raw column values.  The planner validates each row with the callable.
+    """
+    return default_registry.register(
+        name, llm_context=llm_context, check_category=check_category, row_level=row_level
+    )
+
+
+@register_validator(
+    "null_pct_check",
+    llm_context="Returns null percentage as a float between 0.0 and 1.0. SQL pattern: COUNT(CASE WHEN col IS NULL THEN 1 END) * 1.0 / COUNT(*).",
+    check_category="null_check",
+)
+def null_pct_check(*, max_pct: float) -> Callable[[Any], bool]:
+    """
+    Return a validator that passes when ``value <= max_pct``.
+
+    :param max_pct: Maximum allowed null percentage (0.0 – 1.0).
+    :raises TypeError: When the metric value cannot be converted to ``float``.
+    """
+
+    def _check(value: Any) -> bool:
+        try:
+            return float(value) <= max_pct
+        except (TypeError, ValueError) as exc:
+            raise TypeError(
+                f"null_pct_check(max_pct={max_pct!r}): expected a numeric value, got {value!r}"
+            ) from exc
+
+    return _check
+
+
+@register_validator(
+    "row_count_check",
+    llm_context="Returns an integer row count. SQL pattern: COUNT(*).",
+    check_category="row_count",
+)
+def row_count_check(*, min_count: int) -> Callable[[Any], bool]:
+    """
+    Return a validator that passes when ``value >= min_count``.
+
+    :param min_count: Minimum required row count.
+    :raises TypeError: When the metric value cannot be converted to ``int``.
+    """
+
+    def _check(value: Any) -> bool:
+        try:
+            return int(value) >= min_count
+        except (TypeError, ValueError) as exc:
+            raise TypeError(
+                f"row_count_check(min_count={min_count!r}): expected an integer value, got {value!r}"
+            ) from exc
+
+    return _check
+
+
+@register_validator(
+    "duplicate_pct_check",
+    llm_context="Returns duplicate percentage as a float between 0.0 and 1.0. SQL pattern: (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / COUNT(*).",
+    check_category="uniqueness",
+)
+def duplicate_pct_check(*, max_pct: float) -> Callable[[Any], bool]:
+    """
+    Return a validator that passes when ``value <= max_pct``.
+
+    :param max_pct: Maximum allowed duplicate percentage (0.0 – 1.0).
+    :raises TypeError: When the metric value cannot be converted to ``float``.
+    """
+
+    def _check(value: Any) -> bool:
+        try:
+            return float(value) <= max_pct
+        except (TypeError, ValueError) as exc:
+            raise TypeError(
+                f"duplicate_pct_check(max_pct={max_pct!r}): expected a numeric value, got {value!r}"
+            ) from exc
+
+    return _check
+
+
+@register_validator(
+    "between_check",
+    llm_context="Returns a numeric value that will be compared against inclusive bounds. SQL should return a single DOUBLE or INTEGER metric.",
+    check_category="numeric_range",
+)
+def between_check(*, min_val: float, max_val: float) -> Callable[[Any], bool]:
+    """
+    Return a validator that passes when ``min_val <= value <= max_val``.
+
+    :param min_val: Inclusive lower bound.
+    :param max_val: Inclusive upper bound.
+    :raises ValueError: When *min_val* > *max_val*.
+    :raises TypeError: When the metric value cannot be converted to ``float``.
+    """
+    if min_val > max_val:
+        raise ValueError(f"between_check: min_val ({min_val!r}) must be <= max_val ({max_val!r})")
+
+    def _check(value: Any) -> bool:
+        try:
+            return min_val <= float(value) <= max_val
+        except (TypeError, ValueError) as exc:
+            raise TypeError(
+                f"between_check(min_val={min_val!r}, max_val={max_val!r}): "
+                f"expected a numeric value, got {value!r}"
+            ) from exc
+
+    return _check
+
+
+@register_validator(
+    "exact_check",
+    llm_context="Returns a value that must exactly equal an expected constant. SQL should return a single scalar metric.",
+    check_category="validity",
+)
+def exact_check(*, expected: Any) -> Callable[[Any], bool]:
+    """
+    Return a validator that passes when ``value == expected``.
+
+    .. note::
+        Comparison uses Python's ``==`` operator without type coercion.
+        ``Decimal(0) == 0`` passes (Python numeric promotion), but
+        ``"0" == 0`` does not.  The behaviour depends on the DB driver's
+        Python type for the returned column.
+
+    :param expected: The exact value the metric must equal.
+    """
+
+    def _check(value: Any) -> bool:
+        return value == expected
+
+    return _check
+
+
+class DQValidationToolset:
+    """
+    Validator catalog manager that exposes registered validators to the LLM.
+
+    Serves two purposes:
+
+    1. **Prompt generation** — :meth:`build_system_prompt_section` produces a
+       text block that is injected into the LLM system prompt so the model knows
+       which validators are available, their parameters, and what SQL metric each
+       one expects.
+
+    2. **Suggestion validation** — :meth:`validate_suggestion` verifies that the
+       name and arguments proposed by the LLM can actually instantiate a callable
+       without error, before the plan is accepted.
+
+    :param registry: Validator registry to expose.  Defaults to
+        :data:`default_registry`.
+    """
+
+    def __init__(self, registry: ValidatorRegistry | None = None) -> None:
+        self._registry: ValidatorRegistry = registry if registry is not None else default_registry
+
+    def build_system_prompt_section(self) -> str:
+        """
+        Return a system-prompt block listing all registered validators.
+
+        The block is appended to the SQL-planning prompt and tells the LLM
+        which validators exist, their parameter signatures, and what SQL
+        metric each one expects.  The LLM fills ``validator_name`` and
+        ``validator_args`` on each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        based on this catalogue.
+        """
+        lines: list[str] = [
+            "",
+            "AVAILABLE VALIDATORS:",
+            "  For each DQCheck, select the most appropriate validator from the list below.",
+            "  Fill the validator_name field with the exact name and validator_args with the",
+            "  required keyword arguments as a JSON object.",
+            '  If no validator is appropriate for a check, set validator_name to "none".',
+            "  If the check has a [FIXED VALIDATOR] annotation, leave validator_name as null.",
+            "",
+        ]
+
+        aggregate_entries: list[str] = []
+        row_level_entries: list[str] = []
+
+        for name in self._registry.list_validators():
+            entry = self._registry.get(name)
+            sig_str = self._format_signature(entry.factory)
+            category = entry.check_category or "—"
+            block = (
+                f'  - name: "{name}"\n'
+                f"    category: {category}\n"
+                f"    row_level: {str(entry.row_level).lower()}\n"
+                f"    parameters: {sig_str}\n"
+                f"    description: {entry.llm_context or '(no description)'}"
+            )
+            if entry.row_level:
+                row_level_entries.append(block)
+            else:
+                aggregate_entries.append(block)
+
+        if aggregate_entries:
+            lines.append("  Aggregate validators (SQL returns a single scalar metric):")
+            lines.extend(aggregate_entries)
+            lines.append("")
+
+        if row_level_entries:
+            lines.append("  Row-level validators (SQL returns raw column values per row, no aggregation):")
+            lines.extend(row_level_entries)
+            lines.append("")
+
+        lines += [
+            "  IMPORTANT: validator_args must contain ONLY the keyword argument names shown",
+            "  above — no extra keys, no positional args.  The argument values must match the",
+            "  expected Python types (float, int, etc.).",
+            "",
+        ]
+        return "\n".join(lines)
+
+    def validate_suggestion(
+        self,
+        check_name: str,
+        validator_name: str,
+        validator_args: dict[str, Any],
+    ) -> tuple[bool, str]:
+        """
+        Verify that *validator_name* exists and can be instantiated with *validator_args*.
+
+        :param check_name: The check this suggestion belongs to (used in error messages).
+        :param validator_name: Name of the validator factory as registered.
+        :param validator_args: Keyword arguments to pass to the factory.
+        :returns: ``(True, "")`` on success; ``(False, error_message)`` on failure.
+        """
+        if not validator_name or validator_name.lower() == "none":
+            return False, (
+                f"Check {check_name!r}: LLM did not suggest a validator (validator_name is null or 'none')."
+            )

Review Comment:
   `validate_suggestion()` currently treats `validator_name` being empty or "none" as an error. This contradicts the system-prompt text in `build_system_prompt_section()` (which instructs the LLM to set "none" when no validator fits) and the operator/docs behavior that aggregate checks may have no validator. Consider treating "none"/empty as a valid "no validator" case (and only enforcing non-empty validator selection for row-level checks).
   ```suggestion
           Empty values and ``"none"`` are treated as a valid "no validator selected"
           case and do not produce an error.
   
           :param check_name: The check this suggestion belongs to (used in error messages).
           :param validator_name: Name of the validator factory as registered, or ``"none"``.
           :param validator_args: Keyword arguments to pass to the factory.
           :returns: ``(True, "")`` on success; ``(False, error_message)`` on failure.
           """
           if not validator_name or validator_name.lower() == "none":
               return True, ""
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,786 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections import Counter
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheck,
+    DQCheckFailedError,
+    DQCheckGroup,
+    DQCheckInput,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import DQValidationToolset, default_registry
+from airflow.providers.common.compat.sdk import Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+def _describe_validator(validator: Callable[[Any], bool]) -> str:
+    """Return a human-readable validator label for failure messages."""
+    display = getattr(validator, "_validator_display", None)
+    if isinstance(display, str) and display:
+        return display
+    validator_name = getattr(validator, "_validator_name", None)
+    if isinstance(validator_name, str) and validator_name:
+        return validator_name
+    validator_name = getattr(validator, "__name__", None)
+    if isinstance(validator_name, str) and validator_name:
+        return validator_name
+    return repr(validator)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``checks`` describes **one** data-quality expectation.  The LLM
+    groups related checks into optimised SQL queries, selects the most appropriate
+    validator for each check from the registered catalog, executes the SQL against
+    the target database, and applies the validators.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Optionally, supply a fixed ``validator`` on a :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckInput`
+    to bypass LLM validator selection for that specific check.
+
+    Generated SQL plans (including LLM-chosen validators) are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional ``validator`` keys).
+        Each entry describes one data-quality expectation.  Names must be unique.
+        Example::
+
+            from airflow.providers.common.ai.utils.dq_models import DQCheckInput
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            checks = [
+                DQCheckInput(name="email_nulls", description="Check for null email addresses"),
+                DQCheckInput(
+                    name="row_count",
+                    description="Ensure at least 1000 rows exist",
+                    validator=row_count_check(min_count=1000),  # fixed — LLM skips this one
+                ),
+            ]
+
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a
+        :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when checks change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.
+        ``dry_run=True`` takes precedence.
+
+    When approval is granted Airflow resumes the task by calling
+    :meth:`execute_complete` with the approved plan JSON.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            checks if not isinstance(checks, list) else [DQCheckInput.coerce(c) for c in checks]
+        )
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_checks()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        The plan is generated with a single LLM call that simultaneously selects
+        validators from the registry **and** produces the SQL for each check.
+        Checks that have a user-supplied fixed validator bypass LLM selection.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.
+        :raises DQCheckFailedError: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        fixed_validators = self._collect_fixed_validators()
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=self.validator_contexts,
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+            fixed_validators=fixed_validators,
+        )
+
+    @cached_property
+    def validator_contexts(self) -> str:
+        """Return validator-specific LLM context rendered from fixed validators only."""
+        fixed = self._collect_fixed_validators()
+        return default_registry.build_llm_context(fixed)
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises DQCheckFailedError: If any data-quality check fails.
+        """
+        effective_validators = self._resolve_effective_validators_from_plan(plan)
+        planner.set_row_validators(
+            {name: fn for name, fn in effective_validators.items() if default_registry.is_row_level(fn)}
+        )
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan, effective_validators)
+
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    "value": r.value.to_dict() if isinstance(r.value, RowLevelResult) else r.value,
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise DQCheckFailedError(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category | Validator |",
+            "|---|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            validator_label = self._describe_validator_for_check(check)
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} | {validator_label} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        all_validators = self._resolve_effective_validators()
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % | Validator |",
+            "|---|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = all_validators.get(check.check_name)
+            max_pct = (
+                self._resolve_row_level_max_invalid_pct(
+                    check.check_name,
+                    validator,
+                    default_when_missing=None,
+                    warn_on_missing=False,
+                )
+                if validator is not None
+                else None
+            )
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            validator_label = self._describe_validator_for_check(check)
+            lines.append(

Review Comment:
   In the HITL markdown for row-level groups, `Max invalid %` is derived only from user-fixed validator callables (`_resolve_effective_validators()`), so for LLM-suggested row-level validators the threshold will display as `—` even though the plan includes `validator_args` (e.g. `max_invalid_pct`). Consider deriving the displayed threshold from `check.validator_args` when the callable isn’t available yet, so reviewers can see the effective threshold before approving.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4187884616

   > 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.
   
   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!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4440929715

Hello @gopidesupavan and @kaxil

A bit of context on the design evolution and one open concern I'd appreciate feedback on.

### Architecture shift from the previous version
The earlier version of this operator handled everything internally — schema discovery, SQL execution, and validation were all wired directly into the operator class. This PR aligns it with the pattern used by other common.ai operators: the operator is intentionally thin and orchestrates work through toolsets. This meant dropping some features I had explored (collecting multiple metrics from a single SQL query, grouping checks by category, etc.), but I think the result is a much cleaner and more approachable architecture. A big part of that is possible because of the toolset abstraction that was already in place — hat tip to everyone who built that foundation.

### Open concern: durable and missing Airflow config

Currently, if durable=True is set on the operator but no durable backend is configured in airflow.cfg, the operator raises an error at runtime. The default is effectively "no config = no durable support", which makes durable=True an opt-in that silently breaks without clear guidance.
One option worth considering: default the durable backend to file-based storage when no explicit config is present, so durable=True works out of the box without requiring any Airflow configuration. This would make the opt-in path much more forgiving for new users.

### General note

This PR went through quite a few iterations — each one a learning exercise. The current state is as lean and extensible as I could make it. I know you're busy, but whenever you get a chance, I'd genuinely welcome any feedback on missing features or patterns that don't sit right with you.

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "potiuk (via GitHub)" <gi...@apache.org>.

potiuk commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4043129954

   @cetingokhan This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
   
   **Issues found:**
   - :x: **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](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst).
   - :x: **Build docs**: Failing: CI image checks / Build documentation (--docs-only). Run `breeze build-docs` locally to reproduce. See [Build docs docs](https://github.com/apache/airflow/blob/main/contributing-docs/16_documentation.rst).
   - :x: **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](https://github.com/apache/airflow/blob/main/contributing-docs/12_provider_distributions.rst).
   - :x: **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](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst).
   
   > **Note:** Your branch is **6 commits behind `main`**. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.
   
   **What to do next:**
   - The comment informs you what you need to do.
   - Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
   - Maintainers will then proceed with a normal review.
   
   Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "potiuk (via GitHub)" <gi...@apache.org>.

potiuk commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4476910524

@cetingokhan — Removing the `ready for maintainer review` label and converting back to draft. **CI has failed** since the label was added (`1` failing checks).

The label's contract is that the PR is ready for maintainer review — a regression like this means the PR temporarily isn't. Check the failing checks, fix the regression, push, then mark "Ready for review" again to re-enter the queue.

See the [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).

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](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3032074895


##########
providers/common/ai/tests/unit/common/ai/utils/test_dq_planner.py:
##########
@@ -0,0 +1,1174 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.utils.dq_models import DQCheck, DQCheckGroup, DQPlan
+from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+from airflow.providers.common.ai.utils.sql_validation import SQLSafetyError
+
+
+def _make_plan(*check_names: str) -> DQPlan:
+    """Helper: build a minimal DQPlan with one group per check."""
+    groups = [
+        DQCheckGroup(
+            group_id="numeric_aggregate",
+            query=f"SELECT COUNT(*) AS {name}_count FROM t",
+            checks=[DQCheck(check_name=name, metric_key=f"{name}_count", group_id="numeric_aggregate")],
+        )
+        for name in check_names
+    ]
+    return DQPlan(groups=groups)
+
+
+def _make_llm_hook(plan: DQPlan) -> MagicMock:
+    """Helper: mock PydanticAIHook that returns *plan* from agent.run_sync."""
+    mock_usage = MagicMock(requests=1, tool_calls=0, input_tokens=100, output_tokens=50, total_tokens=150)
+    mock_result = MagicMock(spec=["output", "all_messages", "usage", "response"])
+    mock_result.output = plan
+    mock_result.all_messages.return_value = []
+    mock_result.usage.return_value = mock_usage
+    mock_result.response.model_name = "test-model"
+    mock_agent = MagicMock(spec=["run_sync"])
+    mock_agent.run_sync.return_value = mock_result
+    mock_hook = MagicMock(spec=PydanticAIHook)
+    mock_hook.create_agent.return_value = mock_agent
+    return mock_hook
+
+
+class TestSQLDQPlannerBuildSchema:
+    def test_returns_manual_schema_context_verbatim(self):
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=None)
+        result = planner.build_schema_context(
+            table_names=None,
+            schema_context="Table: t\nColumns: id INT",
+        )
+        assert result == "Table: t\nColumns: id INT"
+
+    def test_introspects_via_db_hook_when_no_manual_context(self):
+        mock_db_hook = MagicMock()
+        mock_db_hook.get_table_schema.return_value = [{"name": "id", "type": "INT"}]
+
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=mock_db_hook)
+        result = planner.build_schema_context(

Review Comment:
   All MagicMock() instances for hooks and engines now use spec=DbApiHook / spec=DataFusionEngine so attribute mismatches fail fast.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,836 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue

Review Comment:
   Added homogeneity validation in execute_plan(): raises ValueError listing the conflicting check names when a group mixes row_level=True and row_level=False checks.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_models.py:
##########
@@ -0,0 +1,227 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Pydantic models for the LLM data quality plan and result reporting.
+
+``DQPlan`` is the structured output type requested from the LLM.  The remaining
+dataclasses hold execution results and are never serialised back to the model.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from pydantic import BaseModel, Field, computed_field
+
+
+class DQCheck(BaseModel):
+    """
+    A single data-quality check produced by the LLM.
+
+    :param check_name: Matches the key supplied by the user in ``prompts``.
+    :param metric_key: Column alias used in the generated SQL (e.g. ``null_email_count``).
+        The operator reads ``row[metric_key]`` from the query result.
+    :param group_id: Logical bucket for grouping checks into a single SQL query
+        (e.g. ``customers_null_check_1``, ``orders_validity_1``).
+    :param check_category: Semantic category assigned by the LLM based on the check
+        description.  Used to sub-group checks within a table so that checks of
+        different natures (e.g. null-checks vs. regex validity) land in separate
+        SQL queries.  Allowed values: ``null_check``, ``uniqueness``, ``validity``,
+        ``numeric_range``, ``row_count``, ``string_format``.

Review Comment:
   Added row_level to the allowed values list in the :param check_category: docstring.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,836 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        active_checks = [
+            check for check in group.checks if check.row_level and check.check_name in self._row_validators
+        ]
+        for check in group.checks:
+            if check.row_level and check.check_name not in self._row_validators:
+                log.warning("No row-level validator found for check %r — skipping.", check.check_name)
+
+        if not active_checks:
+            return {}
+
+        # counters[check_name] = [total, invalid, sample_violations]
+        counters: dict[str, list[Any]] = {check.check_name: [0, 0, []] for check in active_checks}
+
+        chunk_iter: Iterator[list[dict[str, Any]]]
+        if datafusion_engine is not None:
+            chunk_iter = self._iter_datafusion_row_chunks(datafusion_engine, group.query)
+        else:
+            chunk_iter = self._iter_db_row_chunks(group.query)
+
+        total_rows_seen = 0
+        for chunk in chunk_iter:
+            total_rows_seen += len(chunk)
+            for row in chunk:
+                for check in active_checks:
+                    value = row.get(check.metric_key)
+                    c = counters[check.check_name]
+                    c[0] += 1  # total
+                    try:
+                        passed = bool(self._row_validators[check.check_name](value))
+                    except Exception:
+                        # Treat any exception from the validator as a validation failure.
+                        passed = False
+                    if not passed:
+                        c[1] += 1  # invalid
+                        if len(c[2]) < _MAX_VIOLATION_SAMPLES:
+                            c[2].append(str(value))
+
+        if total_rows_seen == 0:
+            log.warning("Row-level query for group %r returned no rows.", group.group_id)
+
+        results: dict[str, RowLevelResult] = {}
+        for check in active_checks:
+            total, invalid, violations = counters[check.check_name]
+            invalid_pct = invalid / total if total else 0.0
+            results[check.check_name] = RowLevelResult(
+                check_name=check.check_name,
+                total=total,
+                invalid=invalid,
+                invalid_pct=invalid_pct,
+                sample_violations=violations,
+                sample_size=_MAX_VIOLATION_SAMPLES if total else 0,
+            )
+            log.info(
+                "Row-level check %r: %d/%d invalid (%.4f%%)",
+                check.check_name,
+                invalid,
+                total,
+                invalid_pct * 100,
+            )
+
+        return results
+
+    def _iter_db_row_chunks(self, query: str) -> Iterator[list[dict[str, Any]]]:
+        """
+        Execute *query* via the DB hook and yield chunks of row-dicts.
+
+        Uses ``cursor.fetchmany(_ROW_LEVEL_CHUNK_SIZE)`` so at most
+        :data:`_ROW_LEVEL_CHUNK_SIZE` raw rows are held in memory at once.
+        Column names are resolved from ``cursor.description``.
+        """
+        with closing(self._db_hook.get_conn()) as conn:  # type: ignore[union-attr]
+            with closing(conn.cursor()) as cur:
+                cur.execute(query)
+                col_names = [str(desc[0]) for desc in cur.description] if cur.description else []
+                while True:
+                    chunk = cur.fetchmany(_ROW_LEVEL_CHUNK_SIZE)
+                    if not chunk:
+                        break
+                    result_chunk: list[dict[str, Any]] = []
+                    for raw_row in chunk:
+                        if isinstance(raw_row, dict):
+                            result_chunk.append(raw_row)
+                        elif isinstance(raw_row, Sequence) and not isinstance(
+                            raw_row, str | bytes | bytearray
+                        ):
+                            result_chunk.append(dict(zip(col_names, raw_row)))
+                        else:
+                            result_chunk.append({})
+                    yield result_chunk

Review Comment:
   Now raises ValueError when cursor.description is None, when tuple length doesn't match column count, or when an unsupported row type is encountered — query/driver issues surface immediately instead of silently producing empty rows.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        For each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        in the group that has ``row_level=True``, the callable stored in
+        :attr:`_row_validators` is applied to ``row[check.metric_key]`` for
+        every fetched row.  A
+        :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`
+        is built from the aggregate counts and a sample of failing values.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        # Fetch all rows from the plain SELECT generated by the LLM.
+        if datafusion_engine is not None:
+            col_lists: dict[str, list[Any]] = datafusion_engine.execute_query(group.query)
+            if not col_lists:
+                log.warning("Row-level query for group %r returned no data.", group.group_id)
+                return {}
+            num_rows = max(len(v) for v in col_lists.values())
+            # Transpose column-oriented result to a list of row dicts.
+            rows: list[dict[str, Any]] = [
+                {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
+                for i in range(num_rows)
+            ]
+        else:
+            raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+            if not raw_rows:
+                log.warning("Row-level query for group %r returned no rows.", group.group_id)
+                return {}
+            # Normalise to list[dict] regardless of whether hook returns dicts or sequences.
+            if isinstance(raw_rows[0], dict):
+                rows = list(raw_rows)  # type: ignore[arg-type]
+            else:
+                col_names = [check.metric_key for check in group.checks]
+                rows = [
+                    dict(zip(col_names, row))
+                    if isinstance(row, Sequence) and not isinstance(row, str | bytes | bytearray)
+                    else {}
+                    for row in raw_rows
+                ]
+
+        results: dict[str, RowLevelResult] = {}
+
+        for check in group.checks:
+            if not check.row_level:
+                continue
+            validator = self._row_validators.get(check.check_name)
+            if validator is None:
+                log.warning(
+                    "No row-level validator found for check %r — skipping.",
+                    check.check_name,
+                )
+                continue
+
+            # sample_size governs only the violation sample list, not which rows are fetched.
+            # The SQL LIMIT (set by the LLM) controls how many rows are actually fetched.
+            sample_size = self._row_level_sample_size or len(rows)
+            total = 0

Review Comment:
   I'm looking for a more stable solution.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        For each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        in the group that has ``row_level=True``, the callable stored in
+        :attr:`_row_validators` is applied to ``row[check.metric_key]`` for
+        every fetched row.  A
+        :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`
+        is built from the aggregate counts and a sample of failing values.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        # Fetch all rows from the plain SELECT generated by the LLM.
+        if datafusion_engine is not None:
+            col_lists: dict[str, list[Any]] = datafusion_engine.execute_query(group.query)
+            if not col_lists:
+                log.warning("Row-level query for group %r returned no data.", group.group_id)
+                return {}
+            num_rows = max(len(v) for v in col_lists.values())
+            # Transpose column-oriented result to a list of row dicts.
+            rows: list[dict[str, Any]] = [
+                {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
+                for i in range(num_rows)
+            ]
+        else:
+            raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+            if not raw_rows:
+                log.warning("Row-level query for group %r returned no rows.", group.group_id)
+                return {}
+            # Normalise to list[dict] regardless of whether hook returns dicts or sequences.
+            if isinstance(raw_rows[0], dict):
+                rows = list(raw_rows)  # type: ignore[arg-type]
+            else:
+                col_names = [check.metric_key for check in group.checks]
+                rows = [
+                    dict(zip(col_names, row))
+                    if isinstance(row, Sequence) and not isinstance(row, str | bytes | bytearray)
+                    else {}
+                    for row in raw_rows
+                ]
+
+        results: dict[str, RowLevelResult] = {}
+
+        for check in group.checks:
+            if not check.row_level:
+                continue
+            validator = self._row_validators.get(check.check_name)
+            if validator is None:
+                log.warning(
+                    "No row-level validator found for check %r — skipping.",
+                    check.check_name,
+                )
+                continue
+
+            # sample_size governs only the violation sample list, not which rows are fetched.
+            # The SQL LIMIT (set by the LLM) controls how many rows are actually fetched.
+            sample_size = self._row_level_sample_size or len(rows)
+            total = 0

Review Comment:
   Fixed by switching to cursor-based chunk iteration (fetchmany(_ROW_LEVEL_CHUNK_SIZE)). sample_violations is now capped at _MAX_VIOLATION_SAMPLES = 100, independent of row_level_sample_size.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,616 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import AirflowException, Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises AirflowException: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises AirflowException: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise AirflowException(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        plan_hash = _compute_plan_hash(self.prompts, self.prompt_version, self.collect_unexpected)
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, default=None)
+        if cached_json is not None:
+            self.log.info("DQ plan cache hit — key: %r", variable_key)
+            plan = DQPlan.model_validate_json(cached_json)
+            if not plan.plan_hash:
+                plan.plan_hash = plan_hash
+            return plan
+
+        self.log.info("DQ plan cache miss — generating via LLM (key: %r).", variable_key)
+        plan = planner.generate_plan(self.prompts, schema_ctx)
+        plan.plan_hash = plan_hash
+        Variable.set(variable_key, plan.model_dump_json())
+        return plan
+
+    def _validate_results(
+        self,
+        results_map: dict[str, Any],
+        plan: DQPlan,
+    ) -> list[DQCheckResult]:
+        """
+        Apply validators to each metric value and return per-check results.
+
+        For aggregate checks each validator callable receives the raw metric
+        value returned by the database.  For row-level checks, where *value* is
+        a :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`,
+        the pass/fail decision compares ``invalid_pct`` against the validator's
+        ``_max_invalid_pct`` attribute (defaulting to ``0.0`` when absent).
+        Checks without a registered validator are logged and marked as passed.
+
+        :param results_map: ``{check_name: metric_value_or_RowLevelResult}`` as
+            returned by
+            :meth:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner.execute_plan`.
+        :param plan: The DQ plan whose groups and checks drive iteration order.
+        :returns: Per-check
+            :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckResult`
+            list in plan-group order.
+        """
+        check_results: list[DQCheckResult] = []
+
+        for group in plan.groups:
+            for check in group.checks:
+                value = results_map[check.check_name]
+                validator = self.validators.get(check.check_name)
+
+                passed = True
+                failure_reason: str | None = None
+
+                if isinstance(value, RowLevelResult):
+                    # Row-level check: evaluate threshold against invalid_pct.
+                    max_pct = getattr(validator, "_max_invalid_pct", 0.0)
+                    passed = value.invalid_pct <= max_pct
+                    if not passed:
+                        failure_reason = (
+                            f"Row-level check failed: {value.invalid}/{value.total} rows invalid "
+                            f"({value.invalid_pct:.4%}), threshold {max_pct:.4%}"
+                        )

Review Comment:
   Added explicit validator is None guard in the RowLevelResult branch: logs a warning and marks the check as passed, consistent with aggregate check behaviour.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,616 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import AirflowException, Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises AirflowException: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises AirflowException: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise AirflowException(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        plan_hash = _compute_plan_hash(self.prompts, self.prompt_version, self.collect_unexpected)
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, default=None)
+        if cached_json is not None:
+            self.log.info("DQ plan cache hit — key: %r", variable_key)
+            plan = DQPlan.model_validate_json(cached_json)
+            if not plan.plan_hash:
+                plan.plan_hash = plan_hash
+            return plan
+
+        self.log.info("DQ plan cache miss — generating via LLM (key: %r).", variable_key)
+        plan = planner.generate_plan(self.prompts, schema_ctx)
+        plan.plan_hash = plan_hash
+        Variable.set(variable_key, plan.model_dump_json())
+        return plan
+
+    def _validate_results(
+        self,
+        results_map: dict[str, Any],
+        plan: DQPlan,
+    ) -> list[DQCheckResult]:
+        """
+        Apply validators to each metric value and return per-check results.
+
+        For aggregate checks each validator callable receives the raw metric
+        value returned by the database.  For row-level checks, where *value* is
+        a :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`,
+        the pass/fail decision compares ``invalid_pct`` against the validator's
+        ``_max_invalid_pct`` attribute (defaulting to ``0.0`` when absent).
+        Checks without a registered validator are logged and marked as passed.
+
+        :param results_map: ``{check_name: metric_value_or_RowLevelResult}`` as
+            returned by
+            :meth:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner.execute_plan`.
+        :param plan: The DQ plan whose groups and checks drive iteration order.
+        :returns: Per-check
+            :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckResult`
+            list in plan-group order.
+        """
+        check_results: list[DQCheckResult] = []
+
+        for group in plan.groups:
+            for check in group.checks:

Review Comment:
   Replaced direct index with results_map.get(); raises ValueError naming the missing check_name and its group_id when the planner returns a partial results map.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3070301493


##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,633 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckFailedError,
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises DQCheckFailedError: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises DQCheckFailedError: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise DQCheckFailedError(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        plan_hash = _compute_plan_hash(
+            self.prompts, self.prompt_version, self.collect_unexpected, self.row_level_sample_size
+        )
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, default=None)

Review Comment:
   I updated with; Variable.get(variable_key, None).
   
   Because of:
   - Some Variable implementations and autospec mocks do not accept the default_var keyword.
   - Using a positional second argument is signature-agnostic and works across implementations.
   - This avoids the runtime/mock TypeError we saw.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3027461412


##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,615 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import AirflowException, Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",

Review Comment:
   `row_level_sample_size` is a user-facing parameter (and affects the generated SQL via LIMIT), but it isn’t included in `template_fields`. This prevents templating it with DAG context (similar to `unexpected_sample_size`). Consider adding `row_level_sample_size` to `template_fields` for consistency.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        For each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        in the group that has ``row_level=True``, the callable stored in
+        :attr:`_row_validators` is applied to ``row[check.metric_key]`` for
+        every fetched row.  A
+        :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`
+        is built from the aggregate counts and a sample of failing values.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        # Fetch all rows from the plain SELECT generated by the LLM.
+        if datafusion_engine is not None:
+            col_lists: dict[str, list[Any]] = datafusion_engine.execute_query(group.query)
+            if not col_lists:
+                log.warning("Row-level query for group %r returned no data.", group.group_id)
+                return {}
+            num_rows = max(len(v) for v in col_lists.values())
+            # Transpose column-oriented result to a list of row dicts.
+            rows: list[dict[str, Any]] = [
+                {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
+                for i in range(num_rows)
+            ]
+        else:
+            raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+            if not raw_rows:
+                log.warning("Row-level query for group %r returned no rows.", group.group_id)
+                return {}
+            # Normalise to list[dict] regardless of whether hook returns dicts or sequences.
+            if isinstance(raw_rows[0], dict):
+                rows = list(raw_rows)  # type: ignore[arg-type]
+            else:
+                col_names = [check.metric_key for check in group.checks]
+                rows = [
+                    dict(zip(col_names, row))
+                    if isinstance(row, Sequence) and not isinstance(row, str | bytes | bytearray)
+                    else {}

Review Comment:
   Tuple-shaped DB rows are normalized by zipping them to `col_names = [check.metric_key for check in group.checks]`. This will mis-map values whenever the row-level SELECT includes primary key columns (as required by the prompt) or any extra columns, because tuple positions won’t align with the metric keys. To make row-level validation reliable with `DbApiHook.get_records()` (which commonly returns tuples), derive column names from cursor description (e.g. `self._db_hook.last_description`) or enforce/obtain dict-shaped rows for these queries.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,615 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import AirflowException, Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises AirflowException: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises AirflowException: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise AirflowException(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        plan_hash = _compute_plan_hash(self.prompts, self.prompt_version, self.collect_unexpected)
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, default=None)
+        if cached_json is not None:
+            self.log.info("DQ plan cache hit — key: %r", variable_key)
+            plan = DQPlan.model_validate_json(cached_json)
+            if not plan.plan_hash:
+                plan.plan_hash = plan_hash
+            return plan
+
+        self.log.info("DQ plan cache miss — generating via LLM (key: %r).", variable_key)
+        plan = planner.generate_plan(self.prompts, schema_ctx)
+        plan.plan_hash = plan_hash
+        Variable.set(variable_key, plan.model_dump_json())
+        return plan
+
+    def _validate_results(
+        self,
+        results_map: dict[str, Any],
+        plan: DQPlan,
+    ) -> list[DQCheckResult]:
+        """
+        Apply validators to each metric value and return per-check results.
+
+        For aggregate checks each validator callable receives the raw metric
+        value returned by the database.  For row-level checks, where *value* is
+        a :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`,
+        the pass/fail decision compares ``invalid_pct`` against the validator's
+        ``_max_invalid_pct`` attribute (defaulting to ``0.0`` when absent).
+        Checks without a registered validator are logged and marked as passed.
+
+        :param results_map: ``{check_name: metric_value_or_RowLevelResult}`` as
+            returned by
+            :meth:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner.execute_plan`.
+        :param plan: The DQ plan whose groups and checks drive iteration order.
+        :returns: Per-check
+            :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckResult`
+            list in plan-group order.
+        """
+        check_results: list[DQCheckResult] = []
+
+        for group in plan.groups:
+            for check in group.checks:
+                value = results_map[check.check_name]
+                validator = self.validators.get(check.check_name)
+
+                passed = True
+                failure_reason: str | None = None
+
+                if isinstance(value, RowLevelResult):
+                    # Row-level check: evaluate threshold against invalid_pct.
+                    max_pct = getattr(validator, "_max_invalid_pct", 0.0)
+                    passed = value.invalid_pct <= max_pct
+                    if not passed:
+                        failure_reason = (
+                            f"Row-level check failed: {value.invalid}/{value.total} rows invalid "
+                            f"({value.invalid_pct:.4%}), threshold {max_pct:.4%}"
+                        )
+                elif validator is not None:
+                    try:
+                        passed = bool(validator(value))
+                    except Exception as exc:
+                        passed = False
+                        failure_reason = str(exc)
+
+                    if not passed and failure_reason is None:
+                        failure_reason = f"{repr(validator)} returned False"
+                else:
+                    self.log.warning(
+                        "No validator found for check %r (metric key: %r). Marking as passed by default.",
+                        check.check_name,
+                        check.metric_key,
+                    )
+
+                check_results.append(
+                    DQCheckResult(
+                        check_name=check.check_name,
+                        metric_key=check.metric_key,
+                        value=value,
+                        passed=passed,
+                        failure_reason=failure_reason,
+                    )
+                )
+
+        return check_results
+
+    def _collect_row_validators(self) -> dict[str, Callable[[Any], bool]]:
+        """Return the subset of validators that are row-level."""
+        return {name: fn for name, fn in self.validators.items() if default_registry.is_row_level(fn)}
+
+    @staticmethod
+    def _attach_unexpected(
+        check_results: list[DQCheckResult],
+        unexpected_map: dict[str, UnexpectedResult],
+    ) -> None:
+        """Attach :class:`UnexpectedResult` objects to their corresponding check results."""
+        for result in check_results:
+            unexpected = unexpected_map.get(result.check_name)
+            if unexpected is not None:
+                result.unexpected = unexpected
+
+    def _validate_validator_keys(self) -> None:
+        """
+        Raise :class:`ValueError` if any validator key is absent from *prompts*.
+
+        Skips validation when *prompts* is not yet a dict — this happens when the
+        operator is constructed via the ``@task.llm_dq`` decorator and *prompts*
+        is set to ``SET_DURING_EXECUTION`` at init time.  The decorator calls this
+        method again after the callable has populated ``self.prompts``.
+        """
+        if not isinstance(self.prompts, dict):
+            return
+        unknown = set(self.validators.keys()) - set(self.prompts.keys())
+        if unknown:
+            raise ValueError(
+                f"validators contains keys that are not present in prompts: {sorted(unknown)}. "
+                "Every validator key must correspond to a prompt key."
+            )
+
+    def _validate_prompts(self) -> None:
+        """
+        Raise :class:`ValueError` if *prompts* is an empty dict.
+
+        Skips validation when *prompts* is not yet a dict (decorator use-case).
+        """
+        if isinstance(self.prompts, dict) and not self.prompts:
+            raise ValueError("prompts must not be empty. Provide at least one check_name → description pair.")
+
+    @cached_property
+    def db_hook(self) -> DbApiHook | None:
+        """Return a DbApiHook when *db_conn_id* is configured, or ``None``."""
+        if not self.db_conn_id:
+            return None
+        return get_db_hook(self.db_conn_id)
+
+    def execute_complete(
+        self, context: Context, generated_output: str, event: dict[str, Any]
+    ) -> dict[str, Any]:
+        """
+        Resume after human approval and execute the data-quality checks.
+
+        Called automatically by Airflow when the HITL trigger fires.  The
+        ``generated_output`` is the JSON-serialised
+        :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan` that was
+        deferred for review.
+
+        :param context: Airflow task context.
+        :param generated_output: JSON string of the approved
+            :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+        :param event: Trigger event payload from the HITL reviewer.
+        :raises HITLRejectException: If the reviewer rejected the plan.
+        :raises HITLTimeoutError: If the approval timed out.
+        :raises AirflowException: If any data-quality check fails after approval.
+        """
+        approved_json = super().execute_complete(context, generated_output, event)
+        plan = DQPlan.model_validate_json(approved_json)
+        planner = self._build_planner()
+        return self._run_checks_and_report(context, planner, plan)
+
+
+def _compute_plan_hash(
+    prompts: dict[str, str],
+    prompt_version: str | None,
+    collect_unexpected: bool = False,
+) -> str:
+    """
+    Return a short, stable hash of *prompts*, *prompt_version* and *collect_unexpected*.
+
+    Sorted serialisation ensures the hash is order-independent.
+    The result is prefixed with the version tag so cache keys are human-readable.
+    ``collect_unexpected`` is included because enabling it changes the generated
+    plan (unexpected_query fields are populated).
+    """
+    payload = json.dumps(sorted(prompts.items()), sort_keys=True)
+    if collect_unexpected:
+        payload += ":unexpected=1"
+    digest = hashlib.sha256(payload.encode()).hexdigest()[:16]
+    version_tag = prompt_version or "default"
+    key = f"{version_tag}_{digest}"
+    return key[:_PLAN_VARIABLE_KEY_MAX_LEN]

Review Comment:
   `_compute_plan_hash()` truncates the full `{version_tag}_{digest}` string to `_PLAN_VARIABLE_KEY_MAX_LEN`. If `prompt_version` is very long, the digest can be truncated away entirely, causing cache-key collisions between different prompt payloads. Consider truncating `version_tag` while always preserving the full digest suffix to avoid collisions.
   ```suggestion
       suffix = f"_{digest}"
       max_version_len = _PLAN_VARIABLE_KEY_MAX_LEN - len(suffix)
       if max_version_len <= 0:
           # Fallback: key length limit is smaller than or equal to the digest suffix.
           # Preserve as much of the suffix (and thus the digest) as possible.
           return suffix[-_PLAN_VARIABLE_KEY_MAX_LEN:]
       if len(version_tag) > max_version_len:
           version_tag = version_tag[:max_version_len]
       return f"{version_tag}{suffix}"
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        For each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        in the group that has ``row_level=True``, the callable stored in
+        :attr:`_row_validators` is applied to ``row[check.metric_key]`` for
+        every fetched row.  A
+        :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`
+        is built from the aggregate counts and a sample of failing values.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        # Fetch all rows from the plain SELECT generated by the LLM.
+        if datafusion_engine is not None:
+            col_lists: dict[str, list[Any]] = datafusion_engine.execute_query(group.query)
+            if not col_lists:
+                log.warning("Row-level query for group %r returned no data.", group.group_id)
+                return {}
+            num_rows = max(len(v) for v in col_lists.values())
+            # Transpose column-oriented result to a list of row dicts.
+            rows: list[dict[str, Any]] = [
+                {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
+                for i in range(num_rows)
+            ]
+        else:
+            raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+            if not raw_rows:
+                log.warning("Row-level query for group %r returned no rows.", group.group_id)
+                return {}
+            # Normalise to list[dict] regardless of whether hook returns dicts or sequences.
+            if isinstance(raw_rows[0], dict):
+                rows = list(raw_rows)  # type: ignore[arg-type]
+            else:
+                col_names = [check.metric_key for check in group.checks]
+                rows = [
+                    dict(zip(col_names, row))
+                    if isinstance(row, Sequence) and not isinstance(row, str | bytes | bytearray)
+                    else {}
+                    for row in raw_rows
+                ]
+
+        results: dict[str, RowLevelResult] = {}
+
+        for check in group.checks:
+            if not check.row_level:
+                continue
+            validator = self._row_validators.get(check.check_name)
+            if validator is None:
+                log.warning(
+                    "No row-level validator found for check %r — skipping.",
+                    check.check_name,
+                )
+                continue
+
+            # sample_size governs only the violation sample list, not which rows are fetched.
+            # The SQL LIMIT (set by the LLM) controls how many rows are actually fetched.
+            sample_size = self._row_level_sample_size or len(rows)
+            total = 0

Review Comment:
   `sample_size = self._row_level_sample_size or len(rows)` can make `sample_violations` grow to O(N) (potentially millions of strings) on full scans, which is a significant memory risk. Consider using a small dedicated cap for `sample_violations` (e.g. 10/100) independent of `row_level_sample_size`, or introduce a separate parameter for the violation sample size.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        For each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        in the group that has ``row_level=True``, the callable stored in
+        :attr:`_row_validators` is applied to ``row[check.metric_key]`` for
+        every fetched row.  A
+        :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`
+        is built from the aggregate counts and a sample of failing values.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        # Fetch all rows from the plain SELECT generated by the LLM.
+        if datafusion_engine is not None:
+            col_lists: dict[str, list[Any]] = datafusion_engine.execute_query(group.query)
+            if not col_lists:
+                log.warning("Row-level query for group %r returned no data.", group.group_id)
+                return {}
+            num_rows = max(len(v) for v in col_lists.values())
+            # Transpose column-oriented result to a list of row dicts.
+            rows: list[dict[str, Any]] = [
+                {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
+                for i in range(num_rows)
+            ]
+        else:
+            raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+            if not raw_rows:
+                log.warning("Row-level query for group %r returned no rows.", group.group_id)
+                return {}
+            # Normalise to list[dict] regardless of whether hook returns dicts or sequences.
+            if isinstance(raw_rows[0], dict):
+                rows = list(raw_rows)  # type: ignore[arg-type]
+            else:
+                col_names = [check.metric_key for check in group.checks]
+                rows = [
+                    dict(zip(col_names, row))
+                    if isinstance(row, Sequence) and not isinstance(row, str | bytes | bytearray)
+                    else {}
+                    for row in raw_rows
+                ]
+
+        results: dict[str, RowLevelResult] = {}
+
+        for check in group.checks:
+            if not check.row_level:
+                continue
+            validator = self._row_validators.get(check.check_name)
+            if validator is None:
+                log.warning(
+                    "No row-level validator found for check %r — skipping.",
+                    check.check_name,
+                )
+                continue
+
+            # sample_size governs only the violation sample list, not which rows are fetched.
+            # The SQL LIMIT (set by the LLM) controls how many rows are actually fetched.
+            sample_size = self._row_level_sample_size or len(rows)
+            total = 0
+            invalid = 0
+            sample_violations: list[str] = []
+
+            for row in rows:
+                value = row.get(check.metric_key)
+                total += 1
+                try:
+                    passed = bool(validator(value))
+                except Exception:
+                    # Treat any exception from the validator as a validation failure.
+                    passed = False
+                if not passed:
+                    invalid += 1
+                    if len(sample_violations) < sample_size:
+                        sample_violations.append(str(value))
+
+            invalid_pct = invalid / total if total else 0.0
+            results[check.check_name] = RowLevelResult(
+                check_name=check.check_name,
+                total=total,
+                invalid=invalid,
+                invalid_pct=invalid_pct,
+                sample_violations=sample_violations,
+                sample_size=sample_size,
+            )
+            log.info(
+                "Row-level check %r: %d/%d invalid (%.4f%%)",
+                check.check_name,
+                invalid,
+                total,
+                invalid_pct * 100,
+            )
+
+        return results
+
+    def _build_datafusion_engine(self) -> DataFusionEngine:
+        """
+        Instantiate a DataFusionEngine and register *datasource_config*.
+
+        Ensures the datasource table name is available to all group queries.
+
+        :raises AirflowOptionalProviderFeatureException: When the ``datafusion`` package
+            is not installed.
+        """
+        try:
+            # Lazy load dependency for worker isolation and optional feature degradation
+            from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+        except ImportError as e:
+            from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+            raise AirflowOptionalProviderFeatureException(e)
+
+        engine = DataFusionEngine()
+        engine.register_datasource(self._datasource_config)  # type: ignore[arg-type]
+        return engine
+
+    def _run_datafusion_group(self, engine: DataFusionEngine, group: DQCheckGroup) -> dict[str, Any]:
+        """
+        Execute *group.query* via DataFusion and return the first row as a dict.
+
+        :class:`~airflow.providers.common.sql.datafusion.engine.DataFusionEngine`
+        returns column-oriented results ``{col: [val, ...]}``.  This method
+        converts the first row to a flat ``{col: val}`` dict.
+        """
+        col_lists: dict[str, list[Any]] = engine.execute_query(group.query)
+        if not col_lists or not any(col_lists.values()):
+            return {}
+        return {col: vals[0] for col, vals in col_lists.items() if vals}
+
+    def _run_db_group(self, group: DQCheckGroup) -> dict[str, Any]:
+        """Execute *group.query* via the relational DB hook and return the first row as a dict."""
+        rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+        if not rows:
+            log.warning("Query for group %r returned no rows.", group.group_id)
+            return {}
+
+        metric_keys = [check.metric_key for check in group.checks]
+        raw_row = rows[0]
+        if isinstance(raw_row, dict):
+            return raw_row
+
+        if isinstance(raw_row, Sequence) and not isinstance(raw_row, str | bytes | bytearray):
+            if len(raw_row) != len(metric_keys):
+                raise ValueError(
+                    f"Query for group {group.group_id!r} returned {len(raw_row)} value(s) but "
+                    f"{len(metric_keys)} metric key(s) are required ({metric_keys})."
+                )
+            return dict(zip(metric_keys, raw_row, strict=True))
+
+        raise ValueError(
+            f"Unsupported row type from DbApiHook.get_records for group {group.group_id!r}: "
+            f"{type(raw_row).__name__}. Expected dict or sequence."
+        )
+
+    def _validate_or_fix_group(self, group: DQCheckGroup) -> DQCheckGroup:
+        """
+        Validate *group*'s SQL query and return the (possibly corrected) group.
+
+        If validation passes immediately, the original group is returned unchanged.
+        On :class:`~airflow.providers.common.ai.utils.sql_validation.SQLSafetyError`,
+        the failing query and the error message are sent back to the LLM as a
+        follow-up message in the same conversation.  This repeats up to
+        ``max_sql_retries`` times before re-raising.
+        """
+        try:
+            _validate_sql(
+                group.query,
+                allowed_types=DEFAULT_ALLOWED_TYPES,
+                dialect=self._validation_dialect,
+            )
+            return group
+        except SQLSafetyError as exc:
+            return self._retry_fix_group(group, exc)
+
+    def _retry_fix_group(self, group: DQCheckGroup, initial_error: SQLSafetyError) -> DQCheckGroup:
+        """
+        Ask the LLM to fix *group*'s query by continuing the planning conversation.
+
+        Sends up to ``max_sql_retries`` correction requests.  Each message includes
+        the failing query, the validation error, and a reminder of the SQL rules so
+        the model has full context without needing to re-read the original prompt.
+        """
+        if self._plan_agent is None or self._plan_all_messages is None:
+            raise initial_error
+
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        last_error: SQLSafetyError = initial_error
+        current_group = group
+
+        for attempt in range(1, self._max_sql_retries + 1):
+            log.warning(
+                "SQL validation failed for group %r (attempt %d/%d): %s",
+                current_group.group_id,
+                attempt,
+                self._max_sql_retries,
+                last_error,
+            )
+
+            fix_prompt = self._build_fix_prompt(
+                group=current_group,
+                last_error=last_error,
+                dialect_label=dialect_label,
+                is_datafusion=self._is_datafusion,
+            )
+
+            result = self._plan_agent.run_sync(
+                fix_prompt,
+                message_history=self._plan_all_messages,
+            )
+            self._plan_all_messages = result.all_messages()
+            corrected_plan: DQPlan = result.output
+
+            fixed_group = self._extract_group_by_id(
+                corrected_plan=corrected_plan,
+                target_group_id=current_group.group_id,
+            )
+            if fixed_group is None:
+                log.warning(
+                    "LLM did not return group %r in correction attempt %d — retrying.",
+                    current_group.group_id,
+                    attempt,
+                )
+                continue
+
+            try:
+                _validate_sql(
+                    fixed_group.query,
+                    allowed_types=DEFAULT_ALLOWED_TYPES,
+                    dialect=self._validation_dialect,
+                )
+                log.info(
+                    "SQL for group %r corrected successfully on attempt %d.",
+                    fixed_group.group_id,
+                    attempt,
+                )
+                return fixed_group
+            except SQLSafetyError as new_exc:
+                last_error = new_exc
+                current_group = fixed_group
+
+        raise SQLSafetyError(
+            f"SQL for group {group.group_id!r} could not be corrected after "
+            f"{self._max_sql_retries} attempt(s). Last error: {last_error}"
+        )

Review Comment:
   When `max_sql_retries` is 0, `_retry_fix_group()` skips the loop and raises a new `SQLSafetyError` stating it “could not be corrected after 0 attempt(s)”, rather than re-raising the original validation error. Consider short-circuiting when `max_sql_retries <= 0` to preserve the original exception and keep the error message accurate.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        For each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        in the group that has ``row_level=True``, the callable stored in
+        :attr:`_row_validators` is applied to ``row[check.metric_key]`` for
+        every fetched row.  A
+        :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`
+        is built from the aggregate counts and a sample of failing values.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        # Fetch all rows from the plain SELECT generated by the LLM.
+        if datafusion_engine is not None:
+            col_lists: dict[str, list[Any]] = datafusion_engine.execute_query(group.query)
+            if not col_lists:
+                log.warning("Row-level query for group %r returned no data.", group.group_id)
+                return {}
+            num_rows = max(len(v) for v in col_lists.values())
+            # Transpose column-oriented result to a list of row dicts.
+            rows: list[dict[str, Any]] = [
+                {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
+                for i in range(num_rows)
+            ]
+        else:
+            raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+            if not raw_rows:
+                log.warning("Row-level query for group %r returned no rows.", group.group_id)
+                return {}

Review Comment:
   For DB row-level execution, returning `{}` when the query yields no rows means `execute_plan()` will not populate results for the plan’s row-level check(s), leading to `KeyError` during operator validation. Consider returning a `RowLevelResult` per check with `total=0` (or raising a clear exception) instead of returning an empty dict.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+

Review Comment:
   The planning prompt’s allowed `check_category` list doesn’t include `row_level`, but the docs/examples and `DQCheck.check_category` description use `row_level` for row-level checks. Because the prompt later requires every check to use an allowed category, this can cause the LLM to generate plans that violate the rules. Consider either adding `row_level` to the allowed list (and describing its grouping expectations) or updating docs/model to keep `check_category` within the existing categories and rely solely on `row_level=True`.
   ```suggestion
       - string_format   — length, encoding, whitespace, or character-set checks
       - row_level       — per-row or anomaly checks that evaluate individual records
   
     Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,806 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        # Populated by generate_plan; used by _retry_fix_group to continue the conversation.
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            datafusion_engine = self._build_datafusion_engine()
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        For each :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`
+        in the group that has ``row_level=True``, the callable stored in
+        :attr:`_row_validators` is applied to ``row[check.metric_key]`` for
+        every fetched row.  A
+        :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`
+        is built from the aggregate counts and a sample of failing values.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        # Fetch all rows from the plain SELECT generated by the LLM.
+        if datafusion_engine is not None:
+            col_lists: dict[str, list[Any]] = datafusion_engine.execute_query(group.query)
+            if not col_lists:
+                log.warning("Row-level query for group %r returned no data.", group.group_id)
+                return {}
+            num_rows = max(len(v) for v in col_lists.values())
+            # Transpose column-oriented result to a list of row dicts.
+            rows: list[dict[str, Any]] = [
+                {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
+                for i in range(num_rows)
+            ]
+        else:
+            raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
+            if not raw_rows:
+                log.warning("Row-level query for group %r returned no rows.", group.group_id)
+                return {}
+            # Normalise to list[dict] regardless of whether hook returns dicts or sequences.
+            if isinstance(raw_rows[0], dict):
+                rows = list(raw_rows)  # type: ignore[arg-type]
+            else:
+                col_names = [check.metric_key for check in group.checks]
+                rows = [
+                    dict(zip(col_names, row))
+                    if isinstance(row, Sequence) and not isinstance(row, str | bytes | bytearray)
+                    else {}
+                    for row in raw_rows
+                ]

Review Comment:
   For DataFusion row-level execution, returning `{}` when `execute_query()` yields no columns means `execute_plan()` will not populate results for the plan’s row-level check(s). The operator later expects every `check_name` to be present and will raise `KeyError`. Consider returning a `RowLevelResult` per check with `total=0` (or raising a clear exception) instead of returning an empty dict.
   ```suggestion
                   rows: list[dict[str, Any]] = []
               else:
                   num_rows = max(len(v) for v in col_lists.values())
                   # Transpose column-oriented result to a list of row dicts.
                   rows = [
                       {col: vals[i] if i < len(vals) else None for col, vals in col_lists.items()}
                       for i in range(num_rows)
                   ]
           else:
               raw_rows = self._db_hook.get_records(group.query)  # type: ignore[union-attr]
               if not raw_rows:
                   log.warning("Row-level query for group %r returned no rows.", group.group_id)
                   rows = []
               else:
                   # Normalise to list[dict] regardless of whether hook returns dicts or sequences.
                   if isinstance(raw_rows[0], dict):
                       rows = list(raw_rows)  # type: ignore[arg-type]
                   else:
                       col_names = [check.metric_key for check in group.checks]
                       rows = [
                           dict(zip(col_names, row))
                           if isinstance(row, Sequence) and not isinstance(row, str | bytes | bytearray)
                           else {}
                           for row in raw_rows
                       ]
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4176546929

   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


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4051051339

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!

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4214988980

   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!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4007934653

   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>


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3032145799


##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_validation.py:
##########
@@ -0,0 +1,475 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Built-in and custom validator factories for :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`.
+
+Each factory returns a ``Callable[[Any], bool]`` that can be placed directly in the
+``validators`` dict alongside plain lambdas.  They are intentionally decoupled from
+the operator so they can be tested and composed independently.
+
+Custom validators can be registered with :func:`register_validator` so that the
+operator automatically injects their ``llm_context`` into the LLM system prompt,
+guiding SQL generation for custom metric types.
+
+Usage::
+
+    from airflow.providers.common.ai.utils.dq_validation import (
+        null_pct_check,
+        row_count_check,
+        duplicate_pct_check,
+        between_check,
+        exact_check,
+        register_validator,
+    )
+
+    # Built-in validators
+    validators = {
+        "email_nulls": null_pct_check(max_pct=0.05),
+        "min_customers": row_count_check(min_count=1000),
+        "dup_emails": duplicate_pct_check(max_pct=0.01),
+        "amount_range": between_check(min_val=0.0, max_val=1_000_000.0),
+        "active_flag": exact_check(expected=0),
+        "negative_rev": lambda v: v == 0,
+    }
+
+
+    # Custom validator with LLM context
+    @register_validator(
+        "freshness_check",
+        llm_context=(
+            "Compute hours since the most recent row. "
+            "SQL pattern: EXTRACT(EPOCH FROM (NOW() - MAX(ts_col))) / 3600.0. "
+            "Returns a DOUBLE representing hours elapsed."
+        ),
+        check_category="freshness",
+    )
+    def freshness_check(*, max_hours: float):
+        def _check(value):
+            return float(value) <= max_hours
+
+        return _check
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+# ------------------------------------------------------------------
+# Validator registry — allows custom validators with LLM context
+# ------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class ValidatorEntry:
+    """
+    Metadata for a registered validator factory.
+
+    :param factory: Callable that returns a ``Callable[[Any], bool]`` validator.
+    :param llm_context: Optional hint injected into the LLM system prompt so
+        the model knows what SQL metric format this validator expects.
+    :param check_category: Optional custom check category.  When set, the LLM
+        is instructed to use this category for grouping.
+    :param row_level: When ``True`` the LLM is instructed to generate a plain
+        ``SELECT pk, col FROM table`` (no aggregation).  The planner fetches
+        every row and applies the validator callable to each column value,
+        then reports ``{total, invalid, invalid_pct, sample_violations}``.
+    """
+
+    factory: Callable[..., Callable[[Any], bool]]
+    llm_context: str = ""
+    check_category: str = ""
+    row_level: bool = False
+
+
+class ValidatorRegistry:
+    """
+    Registry for reusable validator factories with optional LLM context.
+
+    Validators registered here can carry an ``llm_context`` string that the
+    operator automatically injects into the LLM system prompt, guiding the
+    model to produce SQL that returns the metric format the validator expects.
+
+    A module-level :data:`default_registry` instance is available.  Use the
+    convenience decorator :func:`register_validator` to register into it.
+    """
+
+    def __init__(self) -> None:
+        self._entries: dict[str, ValidatorEntry] = {}
+
+    def register(
+        self,
+        name: str,
+        *,
+        llm_context: str = "",
+        check_category: str = "",
+        row_level: bool = False,
+    ) -> Callable[[Callable[..., Callable[[Any], bool]]], Callable[..., Callable[[Any], bool]]]:
+        """
+        Return a decorator that registers a validator factory under *name*.
+
+        :param name: Unique name for this validator.
+        :param llm_context: SQL generation hint injected into the LLM prompt.
+        :param check_category: Custom check category for LLM grouping.
+        :param row_level: When ``True``, the LLM generates a plain SELECT
+            returning raw row values instead of an aggregate query.  The
+            planner applies the validator to each row and aggregates results.
+        :raises ValueError: If *name* is already registered.
+        """
+        if name in self._entries:
+            raise ValueError(
+                f"Validator {name!r} is already registered. "
+                "Use a different name or unregister the existing one first."
+            )
+
+        def _decorator(
+            factory: Callable[..., Callable[[Any], bool]],
+        ) -> Callable[..., Callable[[Any], bool]]:
+            self._entries[name] = ValidatorEntry(
+                factory=factory,
+                llm_context=llm_context,
+                check_category=check_category,
+                row_level=row_level,
+            )
+
+            # Wrap the factory so every closure it returns carries the full set of
+            # introspection attributes.  Factory authors no longer need to stamp these
+            # manually — the decorator handles it automatically.
+            def _wrapped_factory(*args: Any, **kwargs: Any) -> Callable[[Any], bool]:
+                closure = factory(*args, **kwargs)
+                # Build a human-readable call representation.
+                arg_parts = [repr(a) for a in args]
+                kwarg_parts = [f"{k}={v!r}" for k, v in kwargs.items()]
+                call_str = f"{name}({', '.join(arg_parts + kwarg_parts)})"
+                # Stamp only when the factory has not set the attribute explicitly.
+                if not hasattr(closure, "_validator_name"):
+                    closure._validator_name = name  # type: ignore[attr-defined]
+                if not hasattr(closure, "_row_level"):
+                    closure._row_level = row_level  # type: ignore[attr-defined]
+                for k, v in kwargs.items():
+                    if not hasattr(closure, f"_{k}"):
+                        setattr(closure, f"_{k}", v)  # e.g. _max_pct, _min_count
+                if ".<locals>." in closure.__qualname__:
+                    closure.__qualname__ = call_str
+                if "__repr__" not in closure.__dict__:
+                    closure.__repr__ = lambda: call_str  # type: ignore[method-assign]
+                return closure
+
+            # Preserve factory identity attributes on the wrapper.
+            _wrapped_factory._validator_name = name  # type: ignore[attr-defined]
+            _wrapped_factory._llm_context = llm_context  # type: ignore[attr-defined]
+            _wrapped_factory._check_category = check_category  # type: ignore[attr-defined]
+            _wrapped_factory._row_level = row_level  # type: ignore[attr-defined]
+            _wrapped_factory.__name__ = factory.__name__
+            _wrapped_factory.__qualname__ = factory.__qualname__
+            _wrapped_factory.__doc__ = factory.__doc__
+
+            # Update the registry entry to point to the wrapped factory.
+            self._entries[name] = ValidatorEntry(
+                factory=_wrapped_factory,
+                llm_context=llm_context,
+                check_category=check_category,
+                row_level=row_level,
+            )
+            return _wrapped_factory
+
+        return _decorator
+
+    def get(self, name: str) -> ValidatorEntry:
+        """
+        Return the :class:`ValidatorEntry` for *name*.
+
+        :raises KeyError: If *name* is not registered.
+        """
+        try:
+            return self._entries[name]
+        except KeyError:
+            raise KeyError(
+                f"Validator {name!r} is not registered. Available validators: {sorted(self._entries)}"
+            ) from None
+
+    def list_validators(self) -> list[str]:
+        """Return sorted list of all registered validator names."""
+        return sorted(self._entries)
+
+    def is_row_level(self, validator: Callable[[Any], bool]) -> bool:
+        """
+        Return ``True`` when *validator* was produced by a row-level factory.
+
+        Checks the ``_row_level`` attribute set by the factory closure and,
+        as a fallback, the registry entry for the factory name.
+        """
+        if getattr(validator, "_row_level", False):
+            return True
+        factory_name: str | None = getattr(validator, "_validator_name", None)
+        if factory_name and factory_name in self._entries:
+            return self._entries[factory_name].row_level
+        return False

Review Comment:
   `ValidatorRegistry.is_row_level()` currently treats an explicit `_row_level=False` on the validator callable the same as “attribute missing”, and then falls back to the registry entry (which may be `row_level=True`). This means factory authors cannot override row-level behavior on a per-closure basis, and a validator that explicitly sets `_row_level = False` could still be treated as row-level.
   
   Consider distinguishing “missing attribute” vs “present but False” (e.g., `if hasattr(validator, '_row_level'): return bool(validator._row_level)`; only fall back to registry when the attribute is absent).



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -121,6 +122,20 @@ def execute_query(self, query: str, max_rows: int | None = None) -> dict[str, li
         except Exception as e:
             raise QueryExecutionException(f"Error while executing query: {e}")
 
+    def iter_query_row_chunks(self, query: str) -> Iterator[dict[str, list[Any]]]:
+        """
+        Execute *query* and yield one column-dict per RecordBatch, streaming results.
+
+        :param query: SQL SELECT query to execute.
+        :raises QueryExecutionException: On SQL execution errors.
+        """
+        try:
+            df = self.session_context.sql(query)
+            for batch in df:
+                yield batch.to_pyarrow().to_pydict()
+        except Exception as e:
+            raise QueryExecutionException(f"Error while executing query: {e}")

Review Comment:
   `iter_query_row_chunks()` iterates over the DataFusion DataFrame (`for batch in df:`). In the DataFusion Python API, a `DataFrame` is typically not iterable; results are usually retrieved via methods like `collect()` / `to_pyarrow_table()` / `to_pydict()`. If `DataFrame` isn’t iterable in the pinned version, this method will always raise at runtime, breaking the row-level DataFusion execution path.
   
   Please switch to a supported batching API for the installed DataFusion version, and consider exception chaining (`raise ... from e`) so the original stack trace is preserved.



##########
providers/common/ai/tests/unit/common/ai/utils/test_db_schema.py:
##########
@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context
+
+
+class TestBuildSchemaContext:
+    def test_raises_when_table_names_given_without_db_hook(self):
+        with pytest.raises(ValueError, match="table_names requires db_conn_id"):
+            build_schema_context(
+                db_hook=None,
+                table_names=["customers"],
+                schema_context=None,
+                datasource_config=MagicMock(),
+            )
+
+    def test_uses_bulk_schema_fetch_when_available(self):
+        mock_db_hook = MagicMock()
+        mock_db_hook.get_table_schemas.return_value = {
+            "customers": [{"name": "id", "type": "INT"}],
+            "orders": [{"name": "order_id", "type": "INT"}],
+        }

Review Comment:
   These tests use unspec'd `MagicMock()` for the DB hook (`mock_db_hook = MagicMock()`), which can hide real API mismatches (e.g. a renamed `get_table_schema(s)` method) until runtime. Prefer `MagicMock(spec=DbApiHook)` / `create_autospec` (and then stub only the needed methods) so the test suite catches interface drift earlier.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/db_schema.py:
##########
@@ -0,0 +1,206 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Shared database hook and schema introspection utilities.
+
+These helpers are used by both :class:`~airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator`
+and :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator` to
+avoid code duplication while keeping both operators decoupled from each other.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from airflow.providers.common.compat.sdk import BaseHook
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+# SQLAlchemy dialect_name → sqlglot dialect mapping for names that differ.
+SQLALCHEMY_TO_SQLGLOT_DIALECT: dict[str, str] = {
+    "postgresql": "postgres",
+    "mssql": "tsql",
+}
+
+
+def get_db_hook(db_conn_id: str) -> DbApiHook:
+    """
+    Return a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` for *db_conn_id*.
+
+    :param db_conn_id: Airflow connection ID that resolves to a ``DbApiHook``.
+    :raises ValueError: If the connection does not resolve to a ``DbApiHook``.
+    """
+    # Lazy load to avoid hard dependency on common.sql
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+    connection = BaseHook.get_connection(db_conn_id)
+    hook = connection.get_hook()
+    if not isinstance(hook, DbApiHook):
+        raise ValueError(
+            f"Connection {db_conn_id!r} does not provide a DbApiHook. Got {type(hook).__name__}."
+        )
+    return hook
+
+
+def resolve_dialect(db_hook: DbApiHook | None, explicit_dialect: str | None) -> str | None:
+    """
+    Resolve the SQL dialect from an explicit parameter or a database hook.
+
+    Normalises SQLAlchemy dialect names to sqlglot equivalents
+    (e.g. ``postgresql`` → ``postgres``).
+
+    :param db_hook: Database hook to read ``dialect_name`` from when *explicit_dialect* is absent.
+    :param explicit_dialect: Caller-supplied dialect string; takes priority over the hook.
+    :return: Resolved dialect string, or ``None`` when neither source provides one.
+    """
+    raw = explicit_dialect
+    if not raw and db_hook and hasattr(db_hook, "dialect_name"):
+        candidate = db_hook.dialect_name
+        raw = candidate if isinstance(candidate, str) else None
+    if raw:
+        return SQLALCHEMY_TO_SQLGLOT_DIALECT.get(raw, raw)
+    return None
+
+
+def build_schema_context(
+    *,
+    db_hook: DbApiHook | None,
+    table_names: list[str] | None,
+    schema_context: str | None,
+    datasource_config: DataSourceConfig | None,
+) -> str:
+    """
+    Return a schema description string suitable for inclusion in an LLM prompt.
+
+    Resolution order:
+    1. *schema_context* — returned as-is when provided (manual override).
+    2. DB introspection via *db_hook* + *table_names*.
+    3. Object-storage introspection via *datasource_config*.
+    4. Empty string when none of the above are available.
+
+    :param db_hook: Hook used for relational-database schema introspection.
+    :param table_names: Table names to introspect via *db_hook*.
+    :param schema_context: Manual schema description; bypasses introspection when set.
+    :param datasource_config: DataFusion datasource config for object-storage schema.
+    :raises ValueError: If *table_names* are provided but none yield schema information.
+    :raises ValueError: If *datasource_config* is ``None`` and no DB tables are available.

Review Comment:
   Docstring says `build_schema_context()` raises a `ValueError` when `datasource_config` is `None` and no DB tables are available, but the implementation returns an empty string when neither `(db_hook and table_names)` nor `datasource_config` is provided. Please align the docstring with actual behavior (or raise in that case if that’s the intended contract).
   ```suggestion
   
   ```



##########
providers/common/ai/tests/unit/common/ai/utils/test_dq_planner.py:
##########
@@ -0,0 +1,1305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.utils.dq_models import DQCheck, DQCheckGroup, DQPlan
+from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+from airflow.providers.common.ai.utils.sql_validation import SQLSafetyError
+from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+
+def _make_plan(*check_names: str) -> DQPlan:
+    """Helper: build a minimal DQPlan with one group per check."""
+    groups = [
+        DQCheckGroup(
+            group_id="numeric_aggregate",
+            query=f"SELECT COUNT(*) AS {name}_count FROM t",
+            checks=[DQCheck(check_name=name, metric_key=f"{name}_count", group_id="numeric_aggregate")],
+        )
+        for name in check_names
+    ]
+    return DQPlan(groups=groups)
+
+
+def _make_llm_hook(plan: DQPlan) -> MagicMock:
+    """Helper: mock PydanticAIHook that returns *plan* from agent.run_sync."""
+    mock_usage = MagicMock(requests=1, tool_calls=0, input_tokens=100, output_tokens=50, total_tokens=150)
+    mock_result = MagicMock(spec=["output", "all_messages", "usage", "response"])
+    mock_result.output = plan
+    mock_result.all_messages.return_value = []
+    mock_result.usage.return_value = mock_usage
+    mock_result.response.model_name = "test-model"
+    mock_agent = MagicMock(spec=["run_sync"])
+    mock_agent.run_sync.return_value = mock_result
+    mock_hook = MagicMock(spec=PydanticAIHook)
+    mock_hook.create_agent.return_value = mock_agent
+    return mock_hook
+
+
+class TestSQLDQPlannerBuildSchema:
+    def test_returns_manual_schema_context_verbatim(self):
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=None)
+        result = planner.build_schema_context(
+            table_names=None,
+            schema_context="Table: t\nColumns: id INT",
+        )
+        assert result == "Table: t\nColumns: id INT"
+
+    def test_introspects_via_db_hook_when_no_manual_context(self):
+        mock_db_hook = MagicMock()
+        mock_db_hook.get_table_schema.return_value = [{"name": "id", "type": "INT"}]
+
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=mock_db_hook)
+        result = planner.build_schema_context(

Review Comment:
   These tests use unspec'd `MagicMock()` for a DB hook (e.g. `mock_db_hook = MagicMock()`), which can mask interface mismatches (typos/missing methods) and make refactors riskier. Prefer `MagicMock(spec=DbApiHook)` (or `create_autospec`) for hooks/engines/cursors so invalid attribute access fails fast.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,873 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        self._cached_datafusion_engine: DataFusionEngine | None = None
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            if self._cached_datafusion_engine is None:
+                self._cached_datafusion_engine = self._build_datafusion_engine()
+            datafusion_engine = self._cached_datafusion_engine
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            row_level_flags = {check.row_level for check in group.checks}
+            if len(row_level_flags) > 1:
+                mixed = [(c.check_name, c.row_level) for c in group.checks]
+                raise ValueError(
+                    f"Group {group.group_id!r} contains both row-level and aggregate checks. "
+                    f"Every check in a group must be homogeneous. Checks: {mixed}"
+                )
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        active_checks = [
+            check for check in group.checks if check.row_level and check.check_name in self._row_validators
+        ]
+        for check in group.checks:
+            if check.row_level and check.check_name not in self._row_validators:
+                log.warning("No row-level validator found for check %r — skipping.", check.check_name)
+
+        if not active_checks:
+            return {}
+
+        # counters[check_name] = [total, invalid, sample_violations]
+        counters: dict[str, list[Any]] = {check.check_name: [0, 0, []] for check in active_checks}
+
+        chunk_iter: Iterator[list[dict[str, Any]]]
+        if datafusion_engine is not None:
+            chunk_iter = self._iter_datafusion_row_chunks(datafusion_engine, group.query)
+        else:
+            chunk_iter = self._iter_db_row_chunks(group.query)
+
+        total_rows_seen = 0
+        for chunk in chunk_iter:
+            total_rows_seen += len(chunk)
+            for row in chunk:
+                for check in active_checks:
+                    value = row.get(check.metric_key)
+                    c = counters[check.check_name]
+                    c[0] += 1  # total
+                    try:
+                        passed = bool(self._row_validators[check.check_name](value))

Review Comment:
   Row-level validation uses `row.get(check.metric_key)` without verifying the `metric_key` column is actually present in the query result. If the LLM generates a row-level SELECT that omits/aliases the metric column incorrectly, this will silently validate `None` values and produce incorrect DQ results.
   
   Recommend failing fast (at least once per chunk) when `check.metric_key` is absent from the row dict / cursor description for a row-level group, similar to the aggregate-path metric-key validation.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,633 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckFailedError,
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises DQCheckFailedError: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises DQCheckFailedError: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,

Review Comment:
   The serialized output for row-level checks omits `sample_size`, but the docs for `RowLevelResult` output describe it as part of the returned `value` dict. Either include `sample_size` in this serialization or update the docs to match what the operator actually returns.
   ```suggestion
                               "invalid_pct": r.value.invalid_pct,
                               "sample_size": r.value.sample_size,
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4217182887

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


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4274379568

   am reviewing it now..


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3105702926


##########
providers/common/ai/src/airflow/providers/common/ai/utils/db_schema.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Shared database hook and schema introspection utilities.
+
+These helpers are used by both :class:`~airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator`
+and :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator` to
+avoid code duplication while keeping both operators decoupled from each other.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from airflow.providers.common.compat.sdk import BaseHook
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+# SQLAlchemy dialect_name → sqlglot dialect mapping for names that differ.
+SQLALCHEMY_TO_SQLGLOT_DIALECT: dict[str, str] = {
+    "postgresql": "postgres",
+    "mssql": "tsql",
+}
+
+
+def get_db_hook(db_conn_id: str) -> DbApiHook:
+    """
+    Return a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` for *db_conn_id*.
+
+    :param db_conn_id: Airflow connection ID that resolves to a ``DbApiHook``.
+    :raises ValueError: If the connection does not resolve to a ``DbApiHook``.
+    """
+    # Lazy load to avoid hard dependency on common.sql
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+    connection = BaseHook.get_connection(db_conn_id)
+    hook = connection.get_hook()
+    if not isinstance(hook, DbApiHook):
+        raise ValueError(
+            f"Connection {db_conn_id!r} does not provide a DbApiHook. Got {type(hook).__name__}."
+        )
+    return hook
+
+
+def resolve_dialect(db_hook: DbApiHook | None, explicit_dialect: str | None) -> str | None:
+    """
+    Resolve the SQL dialect from an explicit parameter or a database hook.
+
+    Normalises SQLAlchemy dialect names to sqlglot equivalents
+    (e.g. ``postgresql`` → ``postgres``).
+
+    :param db_hook: Database hook to read ``dialect_name`` from when *explicit_dialect* is absent.
+    :param explicit_dialect: Caller-supplied dialect string; takes priority over the hook.
+    :return: Resolved dialect string, or ``None`` when neither source provides one.
+    """
+    raw = explicit_dialect
+    if not raw and db_hook and hasattr(db_hook, "dialect_name"):
+        candidate = db_hook.dialect_name
+        raw = candidate if isinstance(candidate, str) else None
+    if raw:
+        return SQLALCHEMY_TO_SQLGLOT_DIALECT.get(raw, raw)
+    return None
+
+
+def build_schema_context(
+    *,
+    db_hook: DbApiHook | None,
+    table_names: list[str] | None,
+    schema_context: str | None,
+    datasource_config: DataSourceConfig | None,
+) -> str:
+    """
+    Return a schema description string suitable for inclusion in an LLM prompt.
+
+    Resolution order:
+    1. *schema_context* — returned as-is when provided (manual override).
+    2. DB introspection via *db_hook* + *table_names*.
+    3. Object-storage introspection via *datasource_config*.
+    4. Empty string when none of the above are available.
+
+    :param db_hook: Hook used for relational-database schema introspection.
+    :param table_names: Table names to introspect via *db_hook*.
+    :param schema_context: Manual schema description; bypasses introspection when set.
+    :param datasource_config: DataFusion datasource config for object-storage schema.
+    :raises ValueError: If *table_names* are provided but none yield schema information.
+    """
+    if schema_context:
+        return schema_context
+
+    if (db_hook and table_names) or datasource_config:
+        return _introspect_schemas(
+            db_hook=db_hook,
+            table_names=table_names,
+            datasource_config=datasource_config,
+        )
+
+    return ""
+
+
+def _introspect_schemas(
+    *,
+    db_hook: DbApiHook | None,
+    table_names: list[str] | None,
+    datasource_config: DataSourceConfig | None,
+) -> str:
+    """Build schema context by introspecting tables and/or object-storage sources."""
+    parts: list[str] = []
+    table_to_columns: dict[str, list[dict[str, str]]] = {}
+
+    if table_names and db_hook is None:
+        raise ValueError("table_names requires db_conn_id so table schema can be introspected.")

Review Comment:
   lets validate this early in the operator init, instead of validating after task run in execute.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3232750926


##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,763 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for running data-quality checks from natural language using an LLM agent."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQPlan,
+    DQReport,
+)
+from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
+    from airflow.sdk import Context
+
+_DQ_SYSTEM_PROMPT = """\
+You are a data-quality expert. Evaluate the user's data-quality checks against a \
+live data source using the tools available to you.
+
+WORKFLOW:
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL CHECKS
+      procedure below instead of writing an aggregate query.
+   b. Otherwise, write a SELECT query that computes the relevant metric for this check.
+   c. Optionally call ``check_query`` to validate SQL syntax before executing.
+   d. Execute the query using the ``query`` tool.
+   e. You MUST call ``apply_validator`` for EVERY check — never skip this step, even for
+      fixed validators (``has_fixed_validator: true``).  Use ``validator_name: "fixed"`` for those.
+5. Return a ``DQReport`` with one ``DQCheckResult`` per check — every check must appear exactly once.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+  Example: -- check: null_order_id
+- For conditional counts, use CASE expressions (FILTER WHERE is not universally supported):
+    CORRECT:   COUNT(CASE WHEN col IS NULL THEN 1 END)
+    INCORRECT: COUNT(*) FILTER (WHERE col IS NULL)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- For duplicate percentages:
+    (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / NULLIF(COUNT(*), 0)
+- For float division, cast to avoid integer truncation:
+    CAST(numerator AS DOUBLE) / CAST(denominator AS DOUBLE)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+VALIDATOR SELECTION:
+- ``list_validators`` returns each validator's name, parameter signature, and description.
+  Read the parameter names and types before calling ``apply_validator``.
+- Pass the required ``validator_args`` as a JSON object matching the parameter signature exactly.
+- If a parameter has no default, it is REQUIRED — always include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — you must pass ``{"max_pct": <value>}``.
+  If you are unsure of a threshold, use a safe default (e.g. ``0.05`` for percentage checks).
+- NEVER pass empty ``{}`` args for a validator that has required parameters.
+- For checks marked ``has_fixed_validator: true`` in ``list_checks``, call ``apply_validator``
+  with ``validator_name: "fixed"`` — the pre-assigned validator runs automatically, no args needed.
+- If no suitable validator exists, pass ``validator_name: "none"`` — the check passes by default
+  but its metric value is still recorded.
+
+DQREPORT OUTPUT RULES:
+- ``check_name`` must exactly match the name returned by ``list_checks`` (no abbreviation).
+- Each check must appear in the report exactly once.
+- Set ``passed: false`` and provide a clear ``failure_reason`` for every failed check.
+- Set ``metric_key`` to the SQL column alias you used for this check's metric value.
+- Set ``sql_query`` to the exact SQL statement you executed for this check (including the leading comment).
+- Set ``validator_info`` to ``{"name": "<validator_name>", "args": {<validator_args>}}`` when a validator
+  was applied, or ``null`` when no validator was used (``validator_name: "none"``).
+
+CHECK CATEGORIES (use to guide SQL and validator selection):
+  null_check    — null / missing value counts or percentages
+  uniqueness    — duplicate detection, cardinality checks
+  validity      — regex / format / pattern matching on string columns
+  numeric_range — range, bounds, or statistical checks on numeric columns
+  row_count     — total row counts or existence checks
+  string_format — length, encoding, whitespace, or character-set checks
+  row_level     — per-row anomaly checks; validator receives a list, not an aggregate
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true`` for a check):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>`` that
+  returns one value per row — no GROUP BY, no COUNT, no CASE expressions.
+  Example: -- check: customer_email_format\n  SELECT email FROM customers
+- After executing, extract the column values into a Python list:
+    value = [row["<column>"] for row in result["rows"]]
+  where ``result`` is the JSON object returned by the ``query`` tool.
+- Call ``apply_validator`` with ``value`` set to that list and ``validator_name: "fixed"``.
+- The validator evaluates each item and returns a ``value`` payload containing a
+    ``RowLevelResult`` summary (total/invalid/invalid_pct/sample_violations/sample_size).
+- Set ``metric_key`` to the SQL column name (e.g. ``"email"``).
+- Set ``value`` in the DQCheckResult to the ``value`` object returned by ``apply_validator``.
+"""
+
+_DQ_PLAN_SYSTEM_PROMPT = """\
+You are a data-quality expert.  Your task is to plan *how* to evaluate each
+data-quality check: write the SQL query and choose the validator for each check.
+A human reviewer will inspect your plan before any SQL is executed.
+
+WORKFLOW (PLANNING MODE — no SQL execution, no apply_validator):
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to understand the data model.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL
+      procedure below.  Otherwise, write an aggregate SELECT query for the metric.
+   b. Optionally call ``check_query`` to validate your SQL syntax — but do NOT execute it.
+   c. Select the appropriate validator name and arguments:
+      - For ``has_fixed_validator: true`` checks: use ``validator_name: "fixed"``, empty args.
+      - If no validator fits: use ``validator_name: "none"``.
+5. Return a ``DQPlan`` with one ``DQCheckPlan`` per check — every check must appear exactly once.
+
+IMPORTANT: The ``query`` tool is NOT available.  Do NOT try to execute SQL.
+Do NOT call ``apply_validator`` — it is not available in planning mode.
+SQL will be executed after the human reviewer approves the plan.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+- For conditional counts, use CASE expressions:
+    COUNT(CASE WHEN col IS NULL THEN 1 END)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true``):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>``.
+- Set ``row_level: true`` in the DQCheckPlan output.
+- Set ``metric_key`` to the column name (e.g. ``"email"``).
+
+DQCheckPlan FIELDS:
+- ``check_name``: exact name from ``list_checks`` (no abbreviation).
+- ``sql_query``: the SQL statement to execute later (with leading ``-- check: <name>`` comment).
+- ``metric_key``: the SQL column alias / name to read as the primary metric value.
+- ``row_level``: ``true`` for row-level checks (SELECT returns one row per record), ``false`` for aggregate.
+- ``validator_name``: ``"fixed"``, a registered validator name, or ``"none"``.
+- ``validator_args``: kwargs for the validator factory, or ``{}`` for ``"fixed"``/``"none"``.
+
+VALIDATOR ARGS RULES (critical — missing args cause hard failures in Phase 2):
+- ``list_validators`` returns a ``parameters`` field for each validator showing its exact
+  parameter names, types, and defaults.  Read this carefully before setting ``validator_args``.
+- If a parameter has no default (required), you MUST include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — always set ``{"max_pct": <value>}``.
+- If you are unsure of the right threshold, pick a reasonable default (e.g. 0.05 for pct checks).
+- NEVER leave ``validator_args`` as ``{}`` for a validator that has required parameters.
+- Only use ``{}`` when ``validator_name`` is ``"fixed"`` or ``"none"``.
+"""
+
+_DIALECT_SQL_NOTES: dict[str, str] = {
+    "postgres": (
+        "  - Regex match: `col ~ 'pattern'`; not match: `col !~ 'pattern'` (case-sensitive).\n"
+        "    Case-insensitive variants: `~*` and `!~*`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)` (no CAST needed).\n"
+        "  - String functions: LENGTH(), LOWER(), TRIM(), SUBSTRING().\n"
+    ),
+    "mysql": (
+        "  - Regex match: `col REGEXP 'pattern'`; not match: `col NOT REGEXP 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - String length (multibyte-safe): CHAR_LENGTH(col); byte length: LENGTH(col).\n"
+    ),
+    "sqlite": (
+        "  - No native regex operator; use LIKE or GLOB:\n"
+        "    `col NOT LIKE '%@%'` or `col GLOB '*@*.*'`.\n"
+        "  - Float division: `CAST(numerator AS REAL) / NULLIF(denominator, 0)`.\n"
+    ),
+    "tsql": (
+        "  - No regex operator; use LIKE or PATINDEX:\n"
+        "    `PATINDEX('%@%.%', col) = 0` means no match.\n"
+        "  - Float division: `CAST(numerator AS FLOAT) / NULLIF(denominator, 0)`.\n"
+        "  - Use TOP n instead of LIMIT n; no OFFSET without ORDER BY.\n"
+        "  - String length: LEN(col) (not LENGTH).\n"
+    ),
+    "bigquery": (
+        "  - Regex: `REGEXP_CONTAINS(col, r'pattern')` — returns TRUE if match.\n"
+        "    Negate with `NOT REGEXP_CONTAINS(...)`.\n"
+        "  - Safe division: `SAFE_DIVIDE(numerator, denominator)` (NULL on zero denominator).\n"
+        "  - Quote reserved names with backticks.\n"
+    ),
+    "snowflake": (
+        "  - Regex: `REGEXP_LIKE(col, 'pattern')` (full-string match) or `col RLIKE 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - Unquoted identifiers are uppercased; use double-quotes when case matters.\n"
+    ),
+}
+
+
+def _extract_schema_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by build_schema_context."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Run data-quality checks described in natural language using an LLM agent.
+
+    The agent discovers the database schema, writes and executes SQL queries,
+    applies validators, and produces a
+    :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.  The task
+    fails when any check does not pass, gating downstream tasks on data quality.
+
+    Supply the data-source toolset and DQ toolset together in ``toolsets``::
+
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+        LLMDataQualityOperator(
+            task_id="quality_check",
+            checks=[
+                DQCheckInput(name="email_nulls", description="Check for null emails"),
+                DQCheckInput(
+                    name="row_count",
+                    description="At least 1000 rows",
+                    validator=row_count_check(min_count=1000),
+                ),
+            ],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers"]),
+                SQLDQToolset(),
+            ],
+        )
+
+    When ``toolsets`` is omitted, the operator auto-creates
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` from
+    ``db_conn_id`` and ``table_names``.
+
+    For config-generation backends (e.g. ``SodaDQToolset``), the operator
+    returns the generated config string as its XCom value instead of a report.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional
+        ``validator`` keys).  Names must be unique.
+    :param toolsets: Pydantic-AI toolsets for the agent.  Must include exactly
+        one :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`
+        subclass alongside a data-source toolset.  When ``None``, the operator
+        auto-creates toolsets from ``db_conn_id``.
+    :param db_conn_id: Connection ID for auto-creating an
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`.
+        Ignored when ``toolsets`` is provided.
+    :param table_names: Tables passed to the auto-created
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` as
+        ``allowed_tables``.  Ignored when ``toolsets`` is provided.
+    :param schema_context: Additional schema description injected into the
+        system prompt.  Useful when the data source cannot be introspected
+        at runtime.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        toolsets: list[AbstractToolset] | None = None,
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Pop operator-specific params that must not reach BaseOperator.__init__.
+        # Using kwargs.pop (rather than named params) is necessary because Airflow's
+        # apply_defaults metaclass captures the full **kwargs dict and uses it for
+        # task-map serialization; named params in the dict are NOT removed from the
+        # snapshot, so they would be re-injected on deserialization and reach
+        # BaseOperator as unknown kwargs.
+        durable: bool = kwargs.pop("durable", False)
+
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "Run the data-quality checks.")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            [DQCheckInput.coerce(c) for c in checks] if isinstance(checks, list) else checks  # type: ignore[assignment]
+        )
+        self.toolsets = toolsets
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.durable = durable
+
+        self._validate_checks()
+
+        if not toolsets and db_conn_id is None:
+            raise ValueError("Either toolsets or db_conn_id must be provided.")
+
+    def execute(self, context: Context) -> Any:
+        """
+        Run the LLM agent to execute all data-quality checks.
+
+        When ``require_approval=True``, the task runs in two phases:
+
+        1. **Phase 1** — the LLM discovers schema, executes SQL queries, and selects
+           validators (but does *not* call ``apply_validator``).  The resulting plan
+           (SQL queries + validator choices) is shown to a human reviewer.
+        2. **Phase 2** (in :meth:`execute_complete`) — after approval, validators are
+           applied in pure Python and the :class:`~...DQReport` is produced.
+
+        :returns: :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`
+            as a dict when the DQ toolset is in ``"execute"`` mode, or a config
+            string when in ``"generate"`` mode.
+        :raises DQCheckFailedError: When any check fails in ``"execute"`` mode.
+        """
+        if self.require_approval:
+            plan, _ = self._run_plan_phase(context)
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                subject=f"Review DQ plan for task `{self.task_id}`",
+                body=self._build_plan_approval_body(plan),
+            )
+        return self._run_and_report(context)
+
+    def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any:
+        """
+        Phase 2: execute SQL and apply validators after human approval.
+
+        Called automatically by Airflow when the HITL trigger fires.  The base
+        class validates the approval decision (raises on reject or timeout).
+        Then each SQL query from the approved :class:`~...DQPlan` is executed in
+        pure Python, the metric values are extracted, and validators are applied
+        — no LLM calls.
+
+        :param context: Airflow task context.
+        :param generated_output: The :class:`~...DQPlan` JSON that was deferred.
+        :param event: Trigger event payload.
+        :returns: Same as :meth:`execute`.
+        :raises HITLRejectException: If the reviewer rejected the plan.
+        """
+        approved_output = super().execute_complete(context, generated_output, event)  # type: ignore[misc]
+        plan = DQPlan.model_validate_json(approved_output)
+
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        report = self._execute_plan_validators(plan, dq_toolset, toolsets)
+        context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+        if not report.passed:
+            raise DQCheckFailedError(report.failure_summary)
+        return report.model_dump()
+
+    # ------------------------------------------------------------------
+    # Core execution helpers
+    # ------------------------------------------------------------------
+
+    def _run_plan_phase(self, context: Context) -> tuple[DQPlan, BaseDQToolset]:
+        """
+        Phase 1 of the two-phase approval flow.
+
+        Runs the LLM agent in planning mode: schema-discovery tools
+        (``list_tables``, ``get_schema``, ``check_query``) are available, but
+        the ``query`` and ``apply_validator`` tools are hidden.  The agent
+        writes SQL strings and selects validators for each check, then outputs
+        a :class:`~...DQPlan` — without executing any SQL.
+
+        :returns: ``(plan, dq_toolset)`` — the plan for the approval body and the
+            toolset (already configured with checks) for Phase 2.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset._planning_mode = True  # omit apply_validator from tool list
+        dq_toolset.set_checks(self.checks)
+
+        # Hide the SQL 'query' tool so the LLM cannot execute DQ queries.
+        for ts in toolsets:
+            if ts is not dq_toolset and hasattr(ts, "_query"):
+                ts._planning_mode = True  # type: ignore[attr-defined]
+
+        instructions = self._build_plan_system_prompt(toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _ = self._build_durable_agent(DQPlan, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=DQPlan,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay (plan phase): model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        return result.output, dq_toolset
+
+    def _run_and_report(self, context: Context) -> Any:
+        """
+        Resolve toolsets, run the LLM agent, push XCom, and handle failures.
+
+        This is the single place where the agent executes regardless of whether
+        the task took the direct path or the approval-gated path.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        output_type: type = DQReport if dq_toolset.output_mode == "execute" else str
+        instructions = self._build_system_prompt(dq_toolset, toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _storage = self._build_durable_agent(output_type, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=output_type,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay: model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        output = result.output
+
+        if dq_toolset.output_mode == "execute":
+            report: DQReport = output
+            context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+            if not report.passed:
+                raise DQCheckFailedError(report.failure_summary)
+            return report.model_dump()
+
+        return output
+
+    def _build_durable_agent(
+        self,
+        output_type: type,
+        instructions: str,
+        toolsets: list[AbstractToolset],
+    ) -> tuple[Any, Any, Any]:
+        """Build a pydantic-ai Agent with CachingModel and CachingToolset wrappers."""
+        from pydantic_ai import Agent
+
+        from airflow.providers.common.ai.durable.caching_model import CachingModel
+        from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
+        from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+        from airflow.providers.common.ai.durable.storage import DurableStorage
+
+        plan_hash = self._compute_plan_hash(instructions)
+        storage = DurableStorage(cache_id=plan_hash)
+        counter = DurableStepCounter()
+
+        wrapped_model = CachingModel(self.llm_hook.get_conn(), storage=storage, counter=counter)
+        cached_toolsets = [CachingToolset(ts, storage=storage, counter=counter) for ts in toolsets]
+
+        agent = Agent(
+            wrapped_model,
+            output_type=output_type,
+            instructions=instructions,
+            toolsets=cached_toolsets,
+            **self.agent_params,
+        )
+        return agent, counter, storage
+
+    def _compute_plan_hash(self, instructions: str) -> str:
+        """Return a stable short hash of the check plan for cross-run cache keying."""
+        checks_data = sorted(
+            [{"name": c.name, "description": c.description} for c in self.checks],
+            key=lambda x: x["name"],
+        )
+        payload = json.dumps(
+            {
+                "checks": checks_data,
+                "schema_context": self.schema_context or "",
+                "instructions": instructions,
+            },
+            sort_keys=True,
+        )
+        digest = hashlib.sha256(payload.encode()).hexdigest()[:16]
+        return f"dq_{digest}"

Review Comment:
   PR text updated with latest development



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4952507702

   Hey @cetingokhan  we are getting new DQ provider https://github.com/apache/airflow/pull/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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4590861612

   Am going to look at this today.. apologies was not around recent days..


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3066477478


##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -121,6 +122,20 @@ def execute_query(self, query: str, max_rows: int | None = None) -> dict[str, li
         except Exception as e:
             raise QueryExecutionException(f"Error while executing query: {e}")
 
+    def iter_query_row_chunks(self, query: str) -> Iterator[dict[str, list[Any]]]:
+        """
+        Execute *query* and yield one column-dict per RecordBatch, streaming results.
+
+        :param query: SQL SELECT query to execute.
+        :raises QueryExecutionException: On SQL execution errors.
+        """
+        try:
+            df = self.session_context.sql(query)
+            for batch in df:

Review Comment:
   The DataFusion Python `DataFrame` object is not guaranteed to be directly iterable across supported versions; iteration may fail and break row-level execution. Prefer using documented APIs that return `RecordBatch`es (e.g., `df.execute_stream()` / `df.collect()` depending on the DataFusion version) and iterate those batches, yielding `to_pydict()` per batch.
   ```suggestion
               if hasattr(df, "execute_stream"):
                   batches = df.execute_stream()
               else:
                   batches = df.collect()
               for batch in batches:
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_validation.py:
##########
@@ -0,0 +1,475 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Built-in and custom validator factories for :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`.
+
+Each factory returns a ``Callable[[Any], bool]`` that can be placed directly in the
+``validators`` dict alongside plain lambdas.  They are intentionally decoupled from
+the operator so they can be tested and composed independently.
+
+Custom validators can be registered with :func:`register_validator` so that the
+operator automatically injects their ``llm_context`` into the LLM system prompt,
+guiding SQL generation for custom metric types.
+
+Usage::
+
+    from airflow.providers.common.ai.utils.dq_validation import (
+        null_pct_check,
+        row_count_check,
+        duplicate_pct_check,
+        between_check,
+        exact_check,
+        register_validator,
+    )
+
+    # Built-in validators
+    validators = {
+        "email_nulls": null_pct_check(max_pct=0.05),
+        "min_customers": row_count_check(min_count=1000),
+        "dup_emails": duplicate_pct_check(max_pct=0.01),
+        "amount_range": between_check(min_val=0.0, max_val=1_000_000.0),
+        "active_flag": exact_check(expected=0),
+        "negative_rev": lambda v: v == 0,
+    }
+
+
+    # Custom validator with LLM context
+    @register_validator(
+        "freshness_check",
+        llm_context=(
+            "Compute hours since the most recent row. "
+            "SQL pattern: EXTRACT(EPOCH FROM (NOW() - MAX(ts_col))) / 3600.0. "
+            "Returns a DOUBLE representing hours elapsed."
+        ),
+        check_category="freshness",
+    )
+    def freshness_check(*, max_hours: float):
+        def _check(value):
+            return float(value) <= max_hours
+
+        return _check
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+# ------------------------------------------------------------------
+# Validator registry — allows custom validators with LLM context
+# ------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class ValidatorEntry:
+    """
+    Metadata for a registered validator factory.
+
+    :param factory: Callable that returns a ``Callable[[Any], bool]`` validator.
+    :param llm_context: Optional hint injected into the LLM system prompt so
+        the model knows what SQL metric format this validator expects.
+    :param check_category: Optional custom check category.  When set, the LLM
+        is instructed to use this category for grouping.
+    :param row_level: When ``True`` the LLM is instructed to generate a plain
+        ``SELECT pk, col FROM table`` (no aggregation).  The planner fetches
+        every row and applies the validator callable to each column value,
+        then reports ``{total, invalid, invalid_pct, sample_violations}``.
+    """
+
+    factory: Callable[..., Callable[[Any], bool]]
+    llm_context: str = ""
+    check_category: str = ""
+    row_level: bool = False
+
+
+class ValidatorRegistry:
+    """
+    Registry for reusable validator factories with optional LLM context.
+
+    Validators registered here can carry an ``llm_context`` string that the
+    operator automatically injects into the LLM system prompt, guiding the
+    model to produce SQL that returns the metric format the validator expects.
+
+    A module-level :data:`default_registry` instance is available.  Use the
+    convenience decorator :func:`register_validator` to register into it.
+    """
+
+    def __init__(self) -> None:
+        self._entries: dict[str, ValidatorEntry] = {}
+
+    def register(
+        self,
+        name: str,
+        *,
+        llm_context: str = "",
+        check_category: str = "",
+        row_level: bool = False,
+    ) -> Callable[[Callable[..., Callable[[Any], bool]]], Callable[..., Callable[[Any], bool]]]:
+        """
+        Return a decorator that registers a validator factory under *name*.
+
+        :param name: Unique name for this validator.
+        :param llm_context: SQL generation hint injected into the LLM prompt.
+        :param check_category: Custom check category for LLM grouping.
+        :param row_level: When ``True``, the LLM generates a plain SELECT
+            returning raw row values instead of an aggregate query.  The
+            planner applies the validator to each row and aggregates results.
+        :raises ValueError: If *name* is already registered.
+        """
+        if name in self._entries:
+            raise ValueError(
+                f"Validator {name!r} is already registered. "
+                "Use a different name or unregister the existing one first."
+            )
+
+        def _decorator(
+            factory: Callable[..., Callable[[Any], bool]],
+        ) -> Callable[..., Callable[[Any], bool]]:
+            self._entries[name] = ValidatorEntry(
+                factory=factory,
+                llm_context=llm_context,
+                check_category=check_category,
+                row_level=row_level,
+            )
+
+            # Wrap the factory so every closure it returns carries the full set of
+            # introspection attributes.  Factory authors no longer need to stamp these
+            # manually — the decorator handles it automatically.
+            def _wrapped_factory(*args: Any, **kwargs: Any) -> Callable[[Any], bool]:
+                closure = factory(*args, **kwargs)
+                # Build a human-readable call representation.
+                arg_parts = [repr(a) for a in args]
+                kwarg_parts = [f"{k}={v!r}" for k, v in kwargs.items()]

Review Comment:
   The generated `call_str` relies on `kwargs.items()` iteration order, which can make `__qualname__` / `__repr__()` output unstable across call sites (especially if kwargs are assembled dynamically). Sorting `kwargs` by key when building `kwarg_parts` makes the representation deterministic and improves debuggability/log consistency.
   ```suggestion
                   kwarg_parts = [f"{k}={v!r}" for k, v in sorted(kwargs.items())]
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,633 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckFailedError,
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises DQCheckFailedError: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=default_registry.build_llm_context(self.validators),
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises DQCheckFailedError: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise DQCheckFailedError(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        plan_hash = _compute_plan_hash(
+            self.prompts, self.prompt_version, self.collect_unexpected, self.row_level_sample_size
+        )
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, default=None)

Review Comment:
   `Variable.get` in Airflow’s public API uses `default_var` (not `default`). Using `default=` can raise a `TypeError` at runtime depending on which `Variable` implementation is in use. Prefer `Variable.get(variable_key, default_var=None)` (or positional second arg) for compatibility.
   ```suggestion
           cached_json = Variable.get(variable_key, default_var=None)
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,873 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        self._cached_datafusion_engine: DataFusionEngine | None = None
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            if self._cached_datafusion_engine is None:
+                self._cached_datafusion_engine = self._build_datafusion_engine()
+            datafusion_engine = self._cached_datafusion_engine
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            row_level_flags = {check.row_level for check in group.checks}
+            if len(row_level_flags) > 1:
+                mixed = [(c.check_name, c.row_level) for c in group.checks]
+                raise ValueError(
+                    f"Group {group.group_id!r} contains both row-level and aggregate checks. "
+                    f"Every check in a group must be homogeneous. Checks: {mixed}"
+                )
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        active_checks = [
+            check for check in group.checks if check.row_level and check.check_name in self._row_validators
+        ]
+        for check in group.checks:
+            if check.row_level and check.check_name not in self._row_validators:
+                log.warning("No row-level validator found for check %r — skipping.", check.check_name)
+
+        if not active_checks:
+            return {}
+
+        # counters[check_name] = [total, invalid, sample_violations]
+        counters: dict[str, list[Any]] = {check.check_name: [0, 0, []] for check in active_checks}
+
+        chunk_iter: Iterator[list[dict[str, Any]]]
+        if datafusion_engine is not None:
+            chunk_iter = self._iter_datafusion_row_chunks(datafusion_engine, group.query)
+        else:
+            chunk_iter = self._iter_db_row_chunks(group.query)
+
+        total_rows_seen = 0
+        for chunk in chunk_iter:
+            total_rows_seen += len(chunk)
+            for row in chunk:
+                for check in active_checks:
+                    value = row.get(check.metric_key)
+                    c = counters[check.check_name]
+                    c[0] += 1  # total
+                    try:
+                        passed = bool(self._row_validators[check.check_name](value))
+                    except Exception:
+                        log.debug(
+                            "Validator for check %r raised on value %r — treating row as invalid.",
+                            check.check_name,
+                            value,
+                            exc_info=True,
+                        )
+                        passed = False
+                    if not passed:
+                        c[1] += 1  # invalid
+                        if len(c[2]) < _MAX_VIOLATION_SAMPLES:
+                            c[2].append(str(value))
+
+        if total_rows_seen == 0:
+            log.warning("Row-level query for group %r returned no rows.", group.group_id)
+
+        results: dict[str, RowLevelResult] = {}
+        for check in active_checks:
+            total, invalid, violations = counters[check.check_name]
+            invalid_pct = invalid / total if total else 0.0
+            results[check.check_name] = RowLevelResult(
+                check_name=check.check_name,
+                total=total,
+                invalid=invalid,
+                invalid_pct=invalid_pct,
+                sample_violations=violations,
+                sample_size=_MAX_VIOLATION_SAMPLES if total else 0,

Review Comment:
   `RowLevelResult.sample_size` is documented as a maximum/cap, but here it becomes `0` when `total == 0` and otherwise always `_MAX_VIOLATION_SAMPLES` regardless of how many samples were actually collected. This makes the field ambiguous for consumers. A clearer contract is to set `sample_size=len(sample_violations)` (actual sample count) and, if needed, add a separate constant/field for the cap; alternatively update the model/docstring to explicitly define this as 'cap used (0 when no rows)'.
   ```suggestion
                   sample_size=len(violations),
   ```



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,873 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        self._cached_datafusion_engine: DataFusionEngine | None = None
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+
+        agent = self._llm_hook.create_agent(
+            output_type=DQPlan, instructions=system_prompt, **self._agent_params
+        )
+        result = agent.run_sync(user_message)
+        log_run_summary(log, result)
+
+        # Persist the agent and full conversation so execute_plan can continue
+        # the same chat thread when asking for SQL corrections.
+        self._plan_agent = agent
+        self._plan_all_messages = result.all_messages()
+
+        plan: DQPlan = result.output
+
+        self._validate_plan_coverage(plan, prompts)
+        self._validate_group_sizes(plan)
+        return plan
+
+    def execute_plan(self, plan: DQPlan) -> dict[str, Any]:
+        """
+        Execute every SQL group in *plan* and return a flat ``{check_name: value}`` map.
+
+        Each group's query is safety-validated via
+        :func:`~airflow.providers.common.ai.utils.sql_validation.validate_sql` before
+        execution.  The first row of each result-set is used; each column corresponds
+        to the ``metric_key`` of one :class:`~airflow.providers.common.ai.utils.dq_models.DQCheck`.
+
+        :param plan: Plan produced by :meth:`generate_plan`.
+        :raises ValueError: If neither *db_hook* nor *datasource_config* was supplied.
+        :raises SQLSafetyError: If a generated query fails AST validation even after
+            ``max_sql_retries`` LLM correction attempts.
+        :raises ValueError: If a query result does not contain an expected metric column.
+        """
+        if self._db_hook is None and self._datasource_config is None:
+            raise ValueError("Either db_conn_id or datasource_config is required to execute the DQ plan.")
+
+        datafusion_engine: DataFusionEngine | None = None
+        if self._db_hook is None:
+            if self._cached_datafusion_engine is None:
+                self._cached_datafusion_engine = self._build_datafusion_engine()
+            datafusion_engine = self._cached_datafusion_engine
+
+        results: dict[str, Any] = {}
+
+        for raw_group in plan.groups:
+            group = self._validate_or_fix_group(raw_group)
+            log.debug("Executing DQ group %r:\n%s", group.group_id, group.query)
+
+            row_level_flags = {check.row_level for check in group.checks}
+            if len(row_level_flags) > 1:
+                mixed = [(c.check_name, c.row_level) for c in group.checks]
+                raise ValueError(
+                    f"Group {group.group_id!r} contains both row-level and aggregate checks. "
+                    f"Every check in a group must be homogeneous. Checks: {mixed}"
+                )
+
+            # Row-level checks and aggregate checks are mutually exclusive within a group
+            # because the LLM places them in separate groups based on the system prompt.
+            if any(check.row_level for check in group.checks):
+                row_level_results = self._execute_row_level_group(group, datafusion_engine)
+                results.update(row_level_results)
+                continue
+
+            if datafusion_engine is not None:
+                row = self._run_datafusion_group(datafusion_engine, group)
+            else:
+                row = self._run_db_group(group)
+
+            for check in group.checks:
+                if check.metric_key not in row:
+                    raise ValueError(
+                        f"Query for group {group.group_id!r} did not return "
+                        f"column {check.metric_key!r} required by check {check.check_name!r}. "
+                        f"Available columns: {list(row.keys())}"
+                    )
+                results[check.check_name] = row[check.metric_key]
+
+        return results
+
+    def _execute_row_level_group(
+        self,
+        group: DQCheckGroup,
+        datafusion_engine: DataFusionEngine | None,
+    ) -> dict[str, RowLevelResult]:
+        """
+        Apply row-level validators to every row returned by *group.query*.
+
+        :param group: A plan group whose checks all have ``row_level=True``.
+        :param datafusion_engine: Active DataFusion engine or ``None`` for DB.
+        :returns: ``{check_name: RowLevelResult}`` for every row-level check.
+        """
+        active_checks = [
+            check for check in group.checks if check.row_level and check.check_name in self._row_validators
+        ]
+        for check in group.checks:
+            if check.row_level and check.check_name not in self._row_validators:
+                log.warning("No row-level validator found for check %r — skipping.", check.check_name)
+
+        if not active_checks:
+            return {}
+
+        # counters[check_name] = [total, invalid, sample_violations]
+        counters: dict[str, list[Any]] = {check.check_name: [0, 0, []] for check in active_checks}
+
+        chunk_iter: Iterator[list[dict[str, Any]]]
+        if datafusion_engine is not None:
+            chunk_iter = self._iter_datafusion_row_chunks(datafusion_engine, group.query)
+        else:
+            chunk_iter = self._iter_db_row_chunks(group.query)
+
+        total_rows_seen = 0
+        for chunk in chunk_iter:
+            total_rows_seen += len(chunk)
+            for row in chunk:
+                for check in active_checks:
+                    value = row.get(check.metric_key)
+                    c = counters[check.check_name]
+                    c[0] += 1  # total
+                    try:
+                        passed = bool(self._row_validators[check.check_name](value))
+                    except Exception:
+                        log.debug(
+                            "Validator for check %r raised on value %r — treating row as invalid.",
+                            check.check_name,
+                            value,
+                            exc_info=True,
+                        )
+                        passed = False
+                    if not passed:
+                        c[1] += 1  # invalid
+                        if len(c[2]) < _MAX_VIOLATION_SAMPLES:
+                            c[2].append(str(value))
+
+        if total_rows_seen == 0:
+            log.warning("Row-level query for group %r returned no rows.", group.group_id)
+
+        results: dict[str, RowLevelResult] = {}
+        for check in active_checks:
+            total, invalid, violations = counters[check.check_name]
+            invalid_pct = invalid / total if total else 0.0
+            results[check.check_name] = RowLevelResult(
+                check_name=check.check_name,
+                total=total,
+                invalid=invalid,
+                invalid_pct=invalid_pct,
+                sample_violations=violations,
+                sample_size=_MAX_VIOLATION_SAMPLES if total else 0,
+            )
+            log.info(
+                "Row-level check %r: %d/%d invalid (%.4f%%)",
+                check.check_name,
+                invalid,
+                total,
+                invalid_pct * 100,
+            )
+
+        return results
+
+    def _iter_db_row_chunks(self, query: str) -> Iterator[list[dict[str, Any]]]:
+        """
+        Execute *query* via the DB hook and yield chunks of row-dicts.
+
+        Uses ``cursor.fetchmany(_ROW_LEVEL_CHUNK_SIZE)`` so at most
+        :data:`_ROW_LEVEL_CHUNK_SIZE` raw rows are held in memory at once.
+        Column names are resolved from ``cursor.description``.
+        """
+        with closing(self._db_hook.get_conn()) as conn:  # type: ignore[union-attr]
+            with closing(conn.cursor()) as cur:
+                cur.execute(query)
+                if cur.description is None:
+                    raise ValueError(
+                        "cursor.description is None after executing row-level query — "
+                        "the driver returned no column metadata."
+                    )
+                col_names = [str(desc[0]) for desc in cur.description]
+                while True:
+                    chunk = cur.fetchmany(_ROW_LEVEL_CHUNK_SIZE)
+                    if not chunk:
+                        break
+                    result_chunk: list[dict[str, Any]] = []
+                    for raw_row in chunk:
+                        if isinstance(raw_row, dict):
+                            result_chunk.append(raw_row)
+                        elif isinstance(raw_row, Sequence) and not isinstance(
+                            raw_row, str | bytes | bytearray
+                        ):
+                            if len(raw_row) != len(col_names):
+                                raise ValueError(
+                                    f"Row has {len(raw_row)} value(s) but cursor.description has "
+                                    f"{len(col_names)} column(s) ({col_names}). "
+                                    f"Possible driver or query mismatch."
+                                )
+                            result_chunk.append(dict(zip(col_names, raw_row)))
+                        else:
+                            raise ValueError(
+                                f"Unsupported row type {type(raw_row).__name__!r} returned by "
+                                f"DB cursor. Expected dict or sequence."
+                            )
+                    yield result_chunk
+
+    @staticmethod
+    def _iter_datafusion_row_chunks(engine: DataFusionEngine, query: str) -> Iterator[list[dict[str, Any]]]:
+        """
+        Execute *query* via DataFusion and yield one batch of row-dicts per RecordBatch.
+
+        Delegates to
+        :meth:`~airflow.providers.common.sql.datafusion.engine.DataFusionEngine.iter_query_row_chunks`
+        which streams results one Arrow RecordBatch at a time.  Each batch is transposed
+        from column-oriented format to a list of row dicts before yielding.
+        """
+        for col_dict in engine.iter_query_row_chunks(query):
+            if not col_dict:
+                continue
+            col_names = list(col_dict.keys())
+            n = len(next(iter(col_dict.values())))

Review Comment:
   This assumes all column lists in `col_dict` have identical lengths and uses the first column to determine `n`. If a batch contains columns with mismatched lengths (or an unexpected empty column), this can silently drop data or raise `IndexError`. Consider validating that all columns have the same length (and handling the all-empty case) before transposing.
   ```suggestion
               col_lengths = {col: len(values) for col, values in col_dict.items()}
               unique_lengths = set(col_lengths.values())
               if len(unique_lengths) != 1:
                   raise ValueError(
                       "DataFusion row chunk has mismatched column lengths: "
                       f"{col_lengths}. Cannot transpose column-oriented batch to rows."
                   )
               n = unique_lengths.pop()
               if n == 0:
                   yield []
                   continue
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4289550953

Hi @gopidesupavan ,
Your suggestion was very logical; we sometimes forget the capabilities of an LLM, and I think it's very sensible to ask the LLM for this :)
I've structured the entire system to only receive prompts (checks), but I've also left the optional validator choice as optional to prevent the LLM from handling specific needs. If a specific validator isn't specified, we leave the selection of the appropriate validator to the LLM.

I didn't include the toolset logic in your suggestion because, since I support custom validation generation, the validation domain is quite large, and I felt that managing it under a toolset would make things even more difficult. However, if we still want to separate it and put it under the toolset, I added it as a separate class in dq_validation.py (if you want, we can separate it and put it under the toolset folder?).

After your suggestions, a fairly large refactoring process took place. There might be some areas that need simplification. Before expending more effort, it's more important that we are all on the same page regarding the concept and usage. Therefore, if there are no different suggestions or corrections regarding the use of the operator in the final stage, I will proceed with the simplification phase. Also, as part of the simplification, I plan to create an abstract called BaseOutputGenerator or something similar and move the existing SQL generation under SQLOutputGenerator. This way, extensions can be added for different control methods in the future.

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.

github-actions[bot] commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-5433213466

   This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4274506838

   and here is just pseudo diff
   
   
   ```
   diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/dq_models.py b/providers/common/ai/src/airflow/providers/common/ai/utils/dq_models.py
   @@
   +from typing import Any, Literal
   +
   +class ValidationSpec(BaseModel):
   +    kind: Literal[
   +        "null_pct",
   +        "row_count_min",
   +        "duplicate_pct",
   +        "exact_value",
   +        "numeric_between",
   +        "regex_row_level",
   +    ]
   +    params: dict[str, Any] = Field(default_factory=dict)
   +
    class DQCheck(BaseModel):
   -    check_name: str
   +    name: str
   +    description: str
        metric_key: str
        group_id: str
        check_category: str = ""
   +    validation: ValidationSpec
        unexpected_query: str | None = None
        row_level: bool = False
   
   diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/dq_validation.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/dq_validation.py
   new file mode 100644
   @@
   +from __future__ import annotations
   +import json
   +from typing import Any
   +from pydantic_ai.tools import ToolDefinition
   +from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
   +from pydantic_core import SchemaValidator, core_schema
   +
   +_ARGS = SchemaValidator(core_schema.any_schema())
   +
   +class DQValidationToolset(AbstractToolset[Any]):
   +    @property
   +    def id(self) -> str:
   +        return "dq-validation"
   +
   +    async def get_tools(self, ctx):
   +        defs = [
   +            ToolDefinition(
   +                name="null_pct_check",
   +                description="Null or missing percentage check. Returns validation kind + params.",
   +                parameters_json_schema={"type": "object", "properties": {"max_pct": {"type": "number"}}, "required": ["max_pct"]},
   +                sequential=True,
   +            ),
   +            ToolDefinition(
   +                name="row_count_min_check",
   +                description="Minimum row count check.",
   +                parameters_json_schema={"type": "object", "properties": {"min_count": {"type": "integer"}}, "required": ["min_count"]},
   +                sequential=True,
   +            ),
   +            ToolDefinition(
   +                name="duplicate_pct_check",
   +                description="Duplicate percentage check.",
   +                parameters_json_schema={"type": "object", "properties": {"max_pct": {"type": "number"}}, "required": ["max_pct"]},
   +                sequential=True,
   +            ),
   +            ToolDefinition(
   +                name="regex_row_level_check",
   +                description="Row-level regex validation with max invalid percentage.",
   +                parameters_json_schema={"type": "object", "properties": {"pattern": {"type": "string"}, "max_invalid_pct": {"type": "number"}}, "required": ["pattern", "max_invalid_pct"]},
   +                sequential=True,
   +            ),
   +        ]
   +        return {
   +            d.name: ToolsetTool(toolset=self, tool_def=d, max_retries=1, args_validator=_ARGS)
   +            for d in defs
   +        }
   +
   +    async def call_tool(self, name, tool_args, ctx, tool):
   +        mapping = {
   +            "null_pct_check": {"kind": "null_pct", "check_category": "null_check", "row_level": False},
   +            "row_count_min_check": {"kind": "row_count_min", "check_category": "row_count", "row_level": False},
   +            "duplicate_pct_check": {"kind": "duplicate_pct", "check_category": "uniqueness", "row_level": False},
   +            "regex_row_level_check": {"kind": "regex_row_level", "check_category": "row_level", "row_level": True},
   +        }
   +        return json.dumps({"kind": mapping[name]["kind"], "params": tool_args, "check_category": mapping[name]["check_category"], "row_level": mapping[name]["row_level"]})
   
   
   diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py b/providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py
   @@
   -from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
   +from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
   +from airflow.providers.common.ai.toolsets.dq_validation import DQValidationToolset
   +from airflow.providers.common.ai.toolsets.sql import SQLToolset
   @@
   -    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
   +    def generate_plan(self, checks: list[dict[str, str]], schema_context: str) -> DQPlan:
   @@
   -        system_prompt = _PLANNING_SYSTEM_PROMPT.format(...)
   +        system_prompt = """
   +You are a SQL data-quality planner.
   +For each requested check:
   +1. Inspect schema with SQL tools if needed.
   +2. Call exactly one DQ validation tool to choose the validation kind.
   +3. Generate SQL that returns a metric matching that validation kind.
   +4. Group by (table, check_category).
   +5. Return only a DQPlan.
   +"""
   @@
   -        user_message = self._build_user_message(prompts)
   +        user_message = self._build_user_message(checks)
   @@
   -        agent = self._llm_hook.create_agent(output_type=DQPlan, instructions=system_prompt, **self._agent_params)
   +        agent = self._llm_hook.create_agent(
   +            output_type=DQPlan,
   +            instructions=system_prompt,
   +            toolsets=[
   +                SQLToolset(db_conn_id=self._db_hook.conn_name_attr, allowed_tables=_extract_table_names(schema_context) or None),
   +                DQValidationToolset(),
   +            ],
   +            **self._agent_params,
   +        )
   @@
   -    def _build_user_message(prompts: dict[str, str]) -> str:
   -        for check_name, description in prompts.items():
   -            lines.append(f'  - check_name="{check_name}": {description}')
   +    def _build_user_message(checks: list[dict[str, str]]) -> str:
   +        for check in checks:
   +            lines.append(f'  - name="{check["name"]}": {check["description"]}')
   
   
   diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py
   @@
   -from airflow.providers.common.ai.utils.dq_validation import default_registry
   @@
   -        prompts: dict[str, str],
   +        checks: list[dict[str, str]],
   @@
   -        validators: dict[str, Callable[[Any], bool]] | None = None,
            dialect: str | None = None,
   @@
   -        self.prompts = prompts
   -        self.validators = validators or {}
   +        self.checks = checks
   @@
   -        self._validate_prompts()
   -        self._validate_validator_keys()
   +        self._validate_checks()
   @@
   -        plan = self._load_or_generate_plan(planner, schema_ctx)
   +        plan = self._load_or_generate_plan(planner, schema_ctx)
   @@
   -        return default_registry.build_llm_context(self.validators)
   +        return ""
   @@
   -        plan_hash = _compute_plan_hash(self.prompts, ...)
   +        plan_hash = _compute_plan_hash(self.checks, ...)
   @@
   -        plan = planner.generate_plan(self.prompts, schema_ctx)
   +        plan = planner.generate_plan(self.checks, schema_ctx)
   @@
   -                validator = self.validators.get(check.check_name)
   -                ...
   -                passed = bool(validator(value))
   +                passed = self._apply_validation(value, check.validation)
   @@
   +    def _apply_validation(self, value: Any, spec) -> bool:
   +        if spec.kind == "null_pct":
   +            return float(value) <= float(spec.params["max_pct"])
   +        if spec.kind == "row_count_min":
   +            return int(value) >= int(spec.params["min_count"])
   +        if spec.kind == "duplicate_pct":
   +            return float(value) <= float(spec.params["max_pct"])
   +        if spec.kind == "exact_value":
   +            return value == spec.params["expected"]
   +        if spec.kind == "numeric_between":
   +            return float(spec.params["min_val"]) <= float(value) <= float(spec.params["max_val"])
   +        raise ValueError(f"Unsupported aggregate validation kind: {spec.kind}")
   +
   +    def _validate_checks(self) -> None:
   +        if not self.checks:
   +            raise ValueError("checks must not be empty")
   
   
   diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
   @@
   -    LLMDataQualityOperator(
   -        prompts={
   -            "null_order_id": "Check the percentage of rows where order_id is NULL",
   -            "duplicate_orders": "Calculate the percentage of duplicate order_id values",
   -            "min_row_count": "Count the total number of rows in the orders table",
   -        },
   -        validators={
   -            "null_order_id": null_pct_check(max_pct=0.0),
   -            "duplicate_orders": duplicate_pct_check(max_pct=0.0),
   -            "min_row_count": row_count_check(min_count=10_000),
   -        },
   -    )
   +    LLMDataQualityOperator(
   +        checks=[
   +            {"name": "null_order_id", "description": "order_id should never be null"},
   +            {"name": "duplicate_orders", "description": "duplicate order_id percentage should be 0%"},
   +            {"name": "min_row_count", "description": "orders table should have at least 10,000 rows"},
   +        ],
   +    )
   
   
   ```


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3466217974


##########
providers/common/ai/src/airflow/providers/common/ai/utils/logging.py:
##########
@@ -36,7 +37,7 @@
 
 def log_run_summary(logger: Logger | logging.Logger, result: AgentRunResult[Any]) -> None:
     """Log model name, token usage, and tool call sequence from an agent run."""
-    usage = result.usage
+    usage = result.usage() if inspect.ismethod(result.usage) else result.usage
     model_name = getattr(result.response, "model_name", "unknown")

Review Comment:
   `result.usage` can be either a method or an attribute; using `inspect.ismethod(...)` is too narrow and will *not* call `usage()` for common callable types (e.g. MagicMock, functions), causing `usage.requests` etc. to be wrong (or fail). Prefer a simple callable check and call when appropriate.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dataquality/validation.py:
##########
@@ -0,0 +1,372 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Built-in and custom validator factories for data-quality checks.
+
+Each factory returns a ``Callable[[Any], bool]`` and can be passed as the
+``validator`` argument of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+to force a specific validator for that check, bypassing LLM selection.
+
+Custom validators registered with :func:`register_validator` are exposed to the LLM
+via the :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` catalog
+so the model can select them automatically.
+
+Usage::
+
+    from airflow.providers.common.ai.utils.dataquality import DQCheckInput
+    from airflow.providers.common.ai.utils.dataquality import null_pct_check, row_count_check
+
+    checks = [
+        DQCheckInput(
+            name="email_nulls",
+            description="Check for null emails",
+            validator=null_pct_check(max_pct=0.05),
+        ),
+        DQCheckInput(
+            name="min_customers",
+            description="Ensure at least 1000 rows",
+            validator=row_count_check(min_count=1000),
+        ),
+    ]
+
+
+    @register_validator(
+        "freshness_check",
+        llm_context=(
+            "Compute hours since the most recent row. "
+            "SQL pattern: EXTRACT(EPOCH FROM (NOW() - MAX(ts_col))) / 3600.0. "
+            "Returns a DOUBLE representing hours elapsed."
+        ),
+        check_category="freshness",
+    )
+    def freshness_check(*, max_hours: float):
+        def _check(value):
+            return float(value) <= max_hours
+
+        return _check
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+log = logging.getLogger(__name__)
+

Review Comment:
   `logging.getLogger(__name__)` is assigned to `log` but never used, which will fail ruff (unused import/variable) and add noise to the module. Either remove it or start using it for warnings/debug output.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,764 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for running data-quality checks from natural language using an LLM agent."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQPlan,
+    DQReport,
+)
+from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
+    from airflow.sdk import Context
+
+_DQ_SYSTEM_PROMPT = """\
+You are a data-quality expert. Evaluate the user's data-quality checks against a \
+live data source using the tools available to you.
+
+WORKFLOW:
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL CHECKS
+      procedure below instead of writing an aggregate query.
+   b. Otherwise, write a SELECT query that computes the relevant metric for this check.
+   c. Optionally call ``check_query`` to validate SQL syntax before executing.
+   d. Execute the query using the ``query`` tool.
+   e. You MUST call ``apply_validator`` for EVERY check — never skip this step, even for
+      fixed validators (``has_fixed_validator: true``).  Use ``validator_name: "fixed"`` for those.
+5. Return a ``DQReport`` with one ``DQCheckResult`` per check — every check must appear exactly once.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+  Example: -- check: null_order_id
+- For conditional counts, use CASE expressions (FILTER WHERE is not universally supported):
+    CORRECT:   COUNT(CASE WHEN col IS NULL THEN 1 END)
+    INCORRECT: COUNT(*) FILTER (WHERE col IS NULL)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- For duplicate percentages:
+    (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / NULLIF(COUNT(*), 0)
+- For float division, cast to avoid integer truncation:
+    CAST(numerator AS DOUBLE) / CAST(denominator AS DOUBLE)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+VALIDATOR SELECTION:
+- ``list_validators`` returns each validator's name, parameter signature, and description.
+  Read the parameter names and types before calling ``apply_validator``.
+- Pass the required ``validator_args`` as a JSON object matching the parameter signature exactly.
+- If a parameter has no default, it is REQUIRED — always include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — you must pass ``{"max_pct": <value>}``.
+  If you are unsure of a threshold, use a safe default (e.g. ``0.05`` for percentage checks).
+- NEVER pass empty ``{}`` args for a validator that has required parameters.
+- For checks marked ``has_fixed_validator: true`` in ``list_checks``, call ``apply_validator``
+  with ``validator_name: "fixed"`` — the pre-assigned validator runs automatically, no args needed.
+- If no suitable validator exists, pass ``validator_name: "none"`` — the check passes by default
+  but its metric value is still recorded.
+
+DQREPORT OUTPUT RULES:
+- ``check_name`` must exactly match the name returned by ``list_checks`` (no abbreviation).
+- Each check must appear in the report exactly once.
+- Set ``passed: false`` and provide a clear ``failure_reason`` for every failed check.
+- Set ``metric_key`` to the SQL column alias you used for this check's metric value.
+- Set ``sql_query`` to the exact SQL statement you executed for this check (including the leading comment).
+- Set ``validator_info`` to ``{"name": "<validator_name>", "args": {<validator_args>}}`` when a validator
+  was applied, or ``null`` when no validator was used (``validator_name: "none"``).
+
+CHECK CATEGORIES (use to guide SQL and validator selection):
+  null_check    — null / missing value counts or percentages
+  uniqueness    — duplicate detection, cardinality checks
+  validity      — regex / format / pattern matching on string columns
+  numeric_range — range, bounds, or statistical checks on numeric columns
+  row_count     — total row counts or existence checks
+  string_format — length, encoding, whitespace, or character-set checks
+  row_level     — per-row anomaly checks; validator receives a list, not an aggregate
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true`` for a check):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>`` that
+  returns one value per row — no GROUP BY, no COUNT, no CASE expressions.
+  Example: -- check: customer_email_format\n  SELECT email FROM customers
+- After executing, extract the column values into a Python list:
+    value = [row["<column>"] for row in result["rows"]]
+  where ``result`` is the JSON object returned by the ``query`` tool.
+- Call ``apply_validator`` with ``value`` set to that list and ``validator_name: "fixed"``.
+- The validator evaluates each item and returns a ``value`` payload containing a
+    ``RowLevelResult`` summary (total/invalid/invalid_pct/sample_violations/sample_size).
+- Set ``metric_key`` to the SQL column name (e.g. ``"email"``).
+- Set ``value`` in the DQCheckResult to the ``value`` object returned by ``apply_validator``.
+"""
+
+_DQ_PLAN_SYSTEM_PROMPT = """\
+You are a data-quality expert.  Your task is to plan *how* to evaluate each
+data-quality check: write the SQL query and choose the validator for each check.
+A human reviewer will inspect your plan before any SQL is executed.
+
+WORKFLOW (PLANNING MODE — no SQL execution, no apply_validator):
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to understand the data model.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL
+      procedure below.  Otherwise, write an aggregate SELECT query for the metric.
+   b. Optionally call ``check_query`` to validate your SQL syntax — but do NOT execute it.
+   c. Select the appropriate validator name and arguments:
+      - For ``has_fixed_validator: true`` checks: use ``validator_name: "fixed"``, empty args.
+      - If no validator fits: use ``validator_name: "none"``.
+5. Return a ``DQPlan`` with one ``DQCheckPlan`` per check — every check must appear exactly once.
+
+IMPORTANT: The ``query`` tool is NOT available.  Do NOT try to execute SQL.
+Do NOT call ``apply_validator`` — it is not available in planning mode.
+SQL will be executed after the human reviewer approves the plan.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+- For conditional counts, use CASE expressions:
+    COUNT(CASE WHEN col IS NULL THEN 1 END)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true``):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>``.
+- Set ``row_level: true`` in the DQCheckPlan output.
+- Set ``metric_key`` to the column name (e.g. ``"email"``).
+
+DQCheckPlan FIELDS:
+- ``check_name``: exact name from ``list_checks`` (no abbreviation).
+- ``sql_query``: the SQL statement to execute later (with leading ``-- check: <name>`` comment).
+- ``metric_key``: the SQL column alias / name to read as the primary metric value.
+- ``row_level``: ``true`` for row-level checks (SELECT returns one row per record), ``false`` for aggregate.
+- ``validator_name``: ``"fixed"``, a registered validator name, or ``"none"``.
+- ``validator_args``: kwargs for the validator factory, or ``{}`` for ``"fixed"``/``"none"``.
+
+VALIDATOR ARGS RULES (critical — missing args cause hard failures in Phase 2):
+- ``list_validators`` returns a ``parameters`` field for each validator showing its exact
+  parameter names, types, and defaults.  Read this carefully before setting ``validator_args``.
+- If a parameter has no default (required), you MUST include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — always set ``{"max_pct": <value>}``.
+- If you are unsure of the right threshold, pick a reasonable default (e.g. 0.05 for pct checks).
+- NEVER leave ``validator_args`` as ``{}`` for a validator that has required parameters.
+- Only use ``{}`` when ``validator_name`` is ``"fixed"`` or ``"none"``.
+"""
+
+_DIALECT_SQL_NOTES: dict[str, str] = {
+    "postgres": (
+        "  - Regex match: `col ~ 'pattern'`; not match: `col !~ 'pattern'` (case-sensitive).\n"
+        "    Case-insensitive variants: `~*` and `!~*`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)` (no CAST needed).\n"
+        "  - String functions: LENGTH(), LOWER(), TRIM(), SUBSTRING().\n"
+    ),
+    "mysql": (
+        "  - Regex match: `col REGEXP 'pattern'`; not match: `col NOT REGEXP 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - String length (multibyte-safe): CHAR_LENGTH(col); byte length: LENGTH(col).\n"
+    ),
+    "sqlite": (
+        "  - No native regex operator; use LIKE or GLOB:\n"
+        "    `col NOT LIKE '%@%'` or `col GLOB '*@*.*'`.\n"
+        "  - Float division: `CAST(numerator AS REAL) / NULLIF(denominator, 0)`.\n"
+    ),
+    "tsql": (
+        "  - No regex operator; use LIKE or PATINDEX:\n"
+        "    `PATINDEX('%@%.%', col) = 0` means no match.\n"
+        "  - Float division: `CAST(numerator AS FLOAT) / NULLIF(denominator, 0)`.\n"
+        "  - Use TOP n instead of LIMIT n; no OFFSET without ORDER BY.\n"
+        "  - String length: LEN(col) (not LENGTH).\n"
+    ),
+    "bigquery": (
+        "  - Regex: `REGEXP_CONTAINS(col, r'pattern')` — returns TRUE if match.\n"
+        "    Negate with `NOT REGEXP_CONTAINS(...)`.\n"
+        "  - Safe division: `SAFE_DIVIDE(numerator, denominator)` (NULL on zero denominator).\n"
+        "  - Quote reserved names with backticks.\n"
+    ),
+    "snowflake": (
+        "  - Regex: `REGEXP_LIKE(col, 'pattern')` (full-string match) or `col RLIKE 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - Unquoted identifiers are uppercased; use double-quotes when case matters.\n"
+    ),
+}
+
+
+def _extract_schema_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by build_schema_context."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Run data-quality checks described in natural language using an LLM agent.
+
+    The agent discovers the database schema, writes and executes SQL queries,
+    applies validators, and produces a
+    :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.  The task
+    fails when any check does not pass, gating downstream tasks on data quality.
+
+    Supply the data-source toolset and DQ toolset together in ``toolsets``::
+
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+        LLMDataQualityOperator(
+            task_id="quality_check",
+            checks=[
+                DQCheckInput(name="email_nulls", description="Check for null emails"),
+                DQCheckInput(
+                    name="row_count",
+                    description="At least 1000 rows",
+                    validator=row_count_check(min_count=1000),
+                ),
+            ],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers"]),
+                SQLDQToolset(),
+            ],
+        )
+
+    When ``toolsets`` is omitted, the operator auto-creates
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` from
+    ``db_conn_id`` and ``table_names``.
+
+    For config-generation backends (e.g. ``SodaDQToolset``), the operator
+    returns the generated config string as its XCom value instead of a report.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional
+        ``validator`` keys).  Names must be unique.
+    :param toolsets: Pydantic-AI toolsets for the agent.  Must include exactly
+        one :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`
+        subclass alongside a data-source toolset.  When ``None``, the operator
+        auto-creates toolsets from ``db_conn_id``.
+    :param db_conn_id: Connection ID for auto-creating an
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`.
+        Ignored when ``toolsets`` is provided.
+    :param table_names: Tables passed to the auto-created
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` as
+        ``allowed_tables``.  Ignored when ``toolsets`` is provided.
+    :param schema_context: Additional schema description injected into the
+        system prompt.  Useful when the data source cannot be introspected
+        at runtime.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        toolsets: list[AbstractToolset] | None = None,
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Pop operator-specific params that must not reach BaseOperator.__init__.
+        # Using kwargs.pop (rather than named params) is necessary because Airflow's
+        # apply_defaults metaclass captures the full **kwargs dict and uses it for
+        # task-map serialization; named params in the dict are NOT removed from the
+        # snapshot, so they would be re-injected on deserialization and reach
+        # BaseOperator as unknown kwargs.
+        durable: bool = kwargs.pop("durable", False)
+
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "Run the data-quality checks.")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            [DQCheckInput.coerce(c) for c in checks] if isinstance(checks, list) else checks  # type: ignore[assignment]
+        )
+        self.toolsets = toolsets
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.durable = durable
+
+        self._validate_checks()
+
+        if not toolsets and db_conn_id is None:
+            raise ValueError("Either toolsets or db_conn_id must be provided.")
+
+    def execute(self, context: Context) -> Any:
+        """
+        Run the LLM agent to execute all data-quality checks.
+
+        When ``require_approval=True``, the task runs in two phases:
+
+        1. **Phase 1** — the LLM discovers schema and selects validators *without
+           executing any SQL queries*.  The resulting plan (check names,
+           descriptions, and validator assignments) is shown to a human reviewer.
+        2. **Phase 2** (in :meth:`execute_complete`) — after approval, SQL queries
+           are executed and validators are applied in pure Python to produce the
+           :class:`~...DQReport`.
+
+        :returns: :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`
+            as a dict when the DQ toolset is in ``"execute"`` mode, or a config
+            string when in ``"generate"`` mode.
+        :raises DQCheckFailedError: When any check fails in ``"execute"`` mode.
+        """
+        if self.require_approval:
+            plan, _ = self._run_plan_phase(context)
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                subject=f"Review DQ plan for task `{self.task_id}`",
+                body=self._build_plan_approval_body(plan),
+            )
+        return self._run_and_report(context)
+
+    def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any:
+        """
+        Phase 2: execute SQL and apply validators after human approval.
+
+        Called automatically by Airflow when the HITL trigger fires.  The base
+        class validates the approval decision (raises on reject or timeout).
+        Then each SQL query from the approved :class:`~...DQPlan` is executed in
+        pure Python, the metric values are extracted, and validators are applied
+        — no LLM calls.
+
+        :param context: Airflow task context.
+        :param generated_output: The :class:`~...DQPlan` JSON that was deferred.
+        :param event: Trigger event payload.
+        :returns: Same as :meth:`execute`.
+        :raises HITLRejectException: If the reviewer rejected the plan.
+        """
+        approved_output = super().execute_complete(context, generated_output, event)  # type: ignore[misc]
+        plan = DQPlan.model_validate_json(approved_output)
+
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        report = self._execute_plan_validators(plan, dq_toolset, toolsets)
+        context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+        if not report.passed:
+            raise DQCheckFailedError(report.failure_summary)
+        return report.model_dump()
+
+    # ------------------------------------------------------------------
+    # Core execution helpers
+    # ------------------------------------------------------------------
+
+    def _run_plan_phase(self, context: Context) -> tuple[DQPlan, BaseDQToolset]:
+        """
+        Phase 1 of the two-phase approval flow.
+
+        Runs the LLM agent in planning mode: schema-discovery tools
+        (``list_tables``, ``get_schema``, ``check_query``) are available, but
+        the ``query`` and ``apply_validator`` tools are hidden.  The agent
+        writes SQL strings and selects validators for each check, then outputs
+        a :class:`~...DQPlan` — without executing any SQL.
+
+        :returns: ``(plan, dq_toolset)`` — the plan for the approval body and the
+            toolset (already configured with checks) for Phase 2.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset._planning_mode = True  # omit apply_validator from tool list
+        dq_toolset.set_checks(self.checks)
+
+        # Hide the SQL 'query' tool so the LLM cannot execute DQ queries.
+        for ts in toolsets:
+            if ts is not dq_toolset and hasattr(ts, "_query"):
+                ts._planning_mode = True  # type: ignore[attr-defined]
+
+        instructions = self._build_plan_system_prompt(toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _ = self._build_durable_agent(DQPlan, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=DQPlan,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay (plan phase): model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        return result.output, dq_toolset
+
+    def _run_and_report(self, context: Context) -> Any:
+        """
+        Resolve toolsets, run the LLM agent, push XCom, and handle failures.
+
+        This is the single place where the agent executes regardless of whether
+        the task took the direct path or the approval-gated path.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        output_type: type = DQReport if dq_toolset.output_mode == "execute" else str
+        instructions = self._build_system_prompt(dq_toolset, toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _storage = self._build_durable_agent(output_type, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=output_type,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay: model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        output = result.output
+
+        if dq_toolset.output_mode == "execute":
+            report: DQReport = output
+            context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+            if not report.passed:
+                raise DQCheckFailedError(report.failure_summary)
+            return report.model_dump()
+
+        return output
+
+    def _build_durable_agent(
+        self,
+        output_type: type,
+        instructions: str,
+        toolsets: list[AbstractToolset],
+    ) -> tuple[Any, Any, Any]:
+        """Build a pydantic-ai Agent with CachingModel and CachingToolset wrappers."""
+        from pydantic_ai import Agent
+
+        from airflow.providers.common.ai.durable.caching_model import CachingModel
+        from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
+        from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+        from airflow.providers.common.ai.durable.storage import DurableStorage
+
+        plan_hash = self._compute_plan_hash(instructions)
+        storage = DurableStorage(cache_id=plan_hash)
+        counter = DurableStepCounter()
+
+        wrapped_model = CachingModel(self.llm_hook.get_conn(), storage=storage, counter=counter)
+        cached_toolsets = [CachingToolset(ts, storage=storage, counter=counter) for ts in toolsets]
+
+        agent = Agent(
+            wrapped_model,
+            output_type=output_type,
+            instructions=instructions,
+            toolsets=cached_toolsets,
+            **self.agent_params,
+        )
+        return agent, counter, storage
+
+    def _compute_plan_hash(self, instructions: str) -> str:
+        """Return a stable short hash of the check plan for cross-run cache keying."""
+        checks_data = sorted(
+            [{"name": c.name, "description": c.description} for c in self.checks],
+            key=lambda x: x["name"],
+        )
+        payload = json.dumps(
+            {
+                "checks": checks_data,
+                "schema_context": self.schema_context or "",
+                "instructions": instructions,
+            },
+            sort_keys=True,
+        )
+        digest = hashlib.sha256(payload.encode()).hexdigest()[:16]
+        return f"dq_{digest}"
+
+    def _build_plan_approval_body(self, plan: DQPlan) -> str:
+        """
+        Build a Markdown body for the HITL approval form.
+
+        Shows the LLM-generated SQL queries and validator selections so the
+        reviewer can inspect the plan before any SQL is executed or any
+        validator is applied.
+        """
+        lines = [
+            "## Data Quality Check Plan \u2014 Awaiting Approval",
+            "",
+            f"**Total checks:** {len(plan.checks)}",
+            "",
+            "Review the SQL queries and validator selections below.  ",
+            "**Approve** to execute the SQL, apply validators, and produce the quality report.  ",
+            "**Reject** to cancel without executing anything.",
+            "",
+            "| # | Check | Metric Key | Row Level | Validator | Args |",
+            "|---|-------|------------|-----------|-----------|------|",
+        ]
+        for i, cp in enumerate(plan.checks, 1):
+            args_display = json.dumps(cp.validator_args) if cp.validator_args else "{}"
+            lines.append(
+                f"| {i} | `{cp.check_name}` | `{cp.metric_key}` "
+                f"| {'Yes' if cp.row_level else 'No'} "
+                f"| `{cp.validator_name}` | `{args_display}` |"
+            )
+
+        lines += ["", "---", "", "### SQL Queries", ""]
+        for cp in plan.checks:
+            lines += [f"**{cp.check_name}**", "", "```sql", cp.sql_query.strip(), "```", ""]
+
+        return "\n".join(lines)
+
+    def _execute_plan_validators(
+        self,
+        plan: DQPlan,
+        dq_toolset: BaseDQToolset,
+        toolsets: list[AbstractToolset],
+    ) -> DQReport:
+        """
+        Phase 2 of the two-phase approval flow.
+
+        For each check in the approved :class:`~...DQPlan`:
+
+        1. Executes the SQL query via the data-source toolset.
+        2. Extracts the metric value from the result.
+        3. Applies the chosen validator in pure Python (no LLM calls).
+        4. Builds the final :class:`~...DQReport`.
+
+        :param plan: The approved :class:`~...DQPlan` from Phase 1.
+        :param dq_toolset: The DQ toolset (must implement ``_apply_validator``).
+        :param toolsets: All toolsets (used to find the SQL toolset for execution).
+        :raises ValueError: If no SQL-capable toolset or no validator executor found.
+        """
+        from airflow.providers.common.ai.utils.dataquality.models import DQCheckResult
+
+        sql_toolset = next((ts for ts in toolsets if hasattr(ts, "_query")), None)
+        if sql_toolset is None:
+            raise ValueError(
+                "require_approval Phase 2 requires a toolset with a _query() method "
+                "(e.g. SQLToolset). Add a data-source toolset to the toolsets list."
+            )
+
+        apply_fn = getattr(dq_toolset, "_apply_validator", None)
+        if apply_fn is None:
+            raise ValueError(
+                "require_approval two-phase execution requires a DQ toolset that "
+                "implements _apply_validator (e.g. SQLDQToolset)."
+            )
+
+        results: list[DQCheckResult] = []
+        for cp in plan.checks:
+            try:
+                raw_result = sql_toolset._query(cp.sql_query)  # type: ignore[union-attr]
+                result_data = json.loads(raw_result)
+                rows: list[dict[str, Any]] = result_data.get("rows", [])
+            except Exception as exc:
+                results.append(
+                    DQCheckResult(
+                        check_name=cp.check_name,
+                        passed=False,
+                        failure_reason=f"SQL execution failed: {exc}",
+                        sql_query=cp.sql_query,
+                        validator_info={"name": cp.validator_name, "args": cp.validator_args},
+                    )
+                )
+                continue
+
+            if cp.row_level:
+                metric_value: Any = [row.get(cp.metric_key) for row in rows]
+            else:
+                metric_value = rows[0].get(cp.metric_key) if rows else None
+
+            raw = apply_fn(cp.check_name, metric_value, cp.validator_name, cp.validator_args)
+            parsed = json.loads(raw)
+            passed = bool(parsed.get("passed", False))
+            result_value = parsed.get("value", metric_value)
+            results.append(
+                DQCheckResult(
+                    check_name=cp.check_name,
+                    passed=passed,
+                    value=result_value,
+                    failure_reason=parsed.get("reason") if not passed else None,
+                    metric_key=cp.metric_key,
+                    sql_query=cp.sql_query,
+                    validator_info={"name": cp.validator_name, "args": cp.validator_args},
+                )
+            )
+
+        return DQReport.build(results)
+
+    # ------------------------------------------------------------------
+    # Toolset / prompt helpers
+    # ------------------------------------------------------------------
+
+    def _resolve_toolsets(self) -> list[AbstractToolset]:
+        """
+        Return explicit toolsets or auto-create from ``db_conn_id``.
+
+        When ``toolsets`` is provided but contains no
+        :class:`BaseDQToolset`, a default
+        :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset`
+        is appended automatically.
+        """
+        if self.toolsets:
+            has_dq = any(isinstance(ts, BaseDQToolset) for ts in self.toolsets)
+            if not has_dq:
+                from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+                self.toolsets.append(SQLDQToolset())
+            return self.toolsets

Review Comment:
   `_resolve_toolsets()` mutates the user-provided `toolsets` list in-place (`self.toolsets.append(...)`). This can have surprising side effects if the same list instance is reused elsewhere (e.g. in multiple tasks) and makes the operator stateful across calls. Prefer returning a new list instead of mutating the input.



##########
providers/common/ai/tests/unit/common/ai/toolsets/dataquality/test_base.py:
##########
@@ -0,0 +1,117 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import json
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import DQCheckInput
+
+
+class _ConcreteDQToolset(BaseDQToolset):
+    """Minimal concrete subclass for testing BaseDQToolset."""
+
+    @property
+    def id(self) -> str:
+        return "test-dq"
+
+    @property
+    def output_mode(self):
+        return "execute"
+
+
+class TestBaseDQToolsetSetChecks:
+    def test_set_checks_stores_checks(self):
+        ts = _ConcreteDQToolset()
+        checks = [DQCheckInput(name="c1", description="desc1")]
+        ts.set_checks(checks)
+        assert ts._checks == checks
+
+    def test_set_checks_replaces_previous(self):
+        ts = _ConcreteDQToolset()
+        ts.set_checks([DQCheckInput(name="old", description="old desc")])
+        new = [DQCheckInput(name="new", description="new desc")]
+        ts.set_checks(new)
+        assert ts._checks == new
+
+
+class TestBaseDQToolsetListChecks:
+    def test_list_checks_returns_all_checks(self):
+        ts = _ConcreteDQToolset()
+        ts.set_checks(
+            [
+                DQCheckInput(name="null_ids", description="Check for null IDs"),
+                DQCheckInput(name="row_count", description="Check minimum row count"),
+            ]
+        )
+        result = json.loads(ts._list_checks())
+        assert len(result) == 2
+        names = [c["name"] for c in result]
+        assert "null_ids" in names
+        assert "row_count" in names
+
+    def test_list_checks_empty_when_no_checks(self):
+        ts = _ConcreteDQToolset()
+        result = json.loads(ts._list_checks())
+        assert result == []
+
+    def test_list_checks_has_fixed_validator_field(self):
+        ts = _ConcreteDQToolset()
+        ts.set_checks(
+            [
+                DQCheckInput(name="c1", description="d1", validator=lambda v: True),
+                DQCheckInput(name="c2", description="d2"),
+            ]
+        )
+        result = {c["name"]: c for c in json.loads(ts._list_checks())}
+        assert result["c1"]["has_fixed_validator"] is True
+        assert result["c2"]["has_fixed_validator"] is False
+
+
+class TestBaseDQToolsetGetTools:
+    def test_get_tools_includes_list_checks(self):
+        ts = _ConcreteDQToolset()
+        tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+        assert "list_checks" in tools
+
+    def test_call_tool_list_checks(self):
+        ts = _ConcreteDQToolset()
+        ts.set_checks([DQCheckInput(name="x", description="y")])
+        result = asyncio.run(ts.call_tool("list_checks", {}, ctx=MagicMock(), tool=MagicMock()))
+        data = json.loads(result)
+        assert data[0]["name"] == "x"

Review Comment:
   This test suite uses `MagicMock()` without a spec for `ctx`/`tool` even though the implementation doesn’t inspect them. Prefer passing `object()` (or a minimal stub) to avoid unspec'd mocks, per the project’s testing guidelines.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,764 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for running data-quality checks from natural language using an LLM agent."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQPlan,
+    DQReport,
+)
+from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
+    from airflow.sdk import Context
+
+_DQ_SYSTEM_PROMPT = """\
+You are a data-quality expert. Evaluate the user's data-quality checks against a \
+live data source using the tools available to you.
+
+WORKFLOW:
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL CHECKS
+      procedure below instead of writing an aggregate query.
+   b. Otherwise, write a SELECT query that computes the relevant metric for this check.
+   c. Optionally call ``check_query`` to validate SQL syntax before executing.
+   d. Execute the query using the ``query`` tool.
+   e. You MUST call ``apply_validator`` for EVERY check — never skip this step, even for
+      fixed validators (``has_fixed_validator: true``).  Use ``validator_name: "fixed"`` for those.
+5. Return a ``DQReport`` with one ``DQCheckResult`` per check — every check must appear exactly once.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+  Example: -- check: null_order_id
+- For conditional counts, use CASE expressions (FILTER WHERE is not universally supported):
+    CORRECT:   COUNT(CASE WHEN col IS NULL THEN 1 END)
+    INCORRECT: COUNT(*) FILTER (WHERE col IS NULL)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- For duplicate percentages:
+    (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / NULLIF(COUNT(*), 0)
+- For float division, cast to avoid integer truncation:
+    CAST(numerator AS DOUBLE) / CAST(denominator AS DOUBLE)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+VALIDATOR SELECTION:
+- ``list_validators`` returns each validator's name, parameter signature, and description.
+  Read the parameter names and types before calling ``apply_validator``.
+- Pass the required ``validator_args`` as a JSON object matching the parameter signature exactly.
+- If a parameter has no default, it is REQUIRED — always include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — you must pass ``{"max_pct": <value>}``.
+  If you are unsure of a threshold, use a safe default (e.g. ``0.05`` for percentage checks).
+- NEVER pass empty ``{}`` args for a validator that has required parameters.
+- For checks marked ``has_fixed_validator: true`` in ``list_checks``, call ``apply_validator``
+  with ``validator_name: "fixed"`` — the pre-assigned validator runs automatically, no args needed.
+- If no suitable validator exists, pass ``validator_name: "none"`` — the check passes by default
+  but its metric value is still recorded.
+
+DQREPORT OUTPUT RULES:
+- ``check_name`` must exactly match the name returned by ``list_checks`` (no abbreviation).
+- Each check must appear in the report exactly once.
+- Set ``passed: false`` and provide a clear ``failure_reason`` for every failed check.
+- Set ``metric_key`` to the SQL column alias you used for this check's metric value.
+- Set ``sql_query`` to the exact SQL statement you executed for this check (including the leading comment).
+- Set ``validator_info`` to ``{"name": "<validator_name>", "args": {<validator_args>}}`` when a validator
+  was applied, or ``null`` when no validator was used (``validator_name: "none"``).
+
+CHECK CATEGORIES (use to guide SQL and validator selection):
+  null_check    — null / missing value counts or percentages
+  uniqueness    — duplicate detection, cardinality checks
+  validity      — regex / format / pattern matching on string columns
+  numeric_range — range, bounds, or statistical checks on numeric columns
+  row_count     — total row counts or existence checks
+  string_format — length, encoding, whitespace, or character-set checks
+  row_level     — per-row anomaly checks; validator receives a list, not an aggregate
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true`` for a check):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>`` that
+  returns one value per row — no GROUP BY, no COUNT, no CASE expressions.
+  Example: -- check: customer_email_format\n  SELECT email FROM customers
+- After executing, extract the column values into a Python list:
+    value = [row["<column>"] for row in result["rows"]]
+  where ``result`` is the JSON object returned by the ``query`` tool.
+- Call ``apply_validator`` with ``value`` set to that list and ``validator_name: "fixed"``.
+- The validator evaluates each item and returns a ``value`` payload containing a
+    ``RowLevelResult`` summary (total/invalid/invalid_pct/sample_violations/sample_size).
+- Set ``metric_key`` to the SQL column name (e.g. ``"email"``).
+- Set ``value`` in the DQCheckResult to the ``value`` object returned by ``apply_validator``.
+"""
+
+_DQ_PLAN_SYSTEM_PROMPT = """\
+You are a data-quality expert.  Your task is to plan *how* to evaluate each
+data-quality check: write the SQL query and choose the validator for each check.
+A human reviewer will inspect your plan before any SQL is executed.
+
+WORKFLOW (PLANNING MODE — no SQL execution, no apply_validator):
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to understand the data model.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL
+      procedure below.  Otherwise, write an aggregate SELECT query for the metric.
+   b. Optionally call ``check_query`` to validate your SQL syntax — but do NOT execute it.
+   c. Select the appropriate validator name and arguments:
+      - For ``has_fixed_validator: true`` checks: use ``validator_name: "fixed"``, empty args.
+      - If no validator fits: use ``validator_name: "none"``.
+5. Return a ``DQPlan`` with one ``DQCheckPlan`` per check — every check must appear exactly once.
+
+IMPORTANT: The ``query`` tool is NOT available.  Do NOT try to execute SQL.
+Do NOT call ``apply_validator`` — it is not available in planning mode.
+SQL will be executed after the human reviewer approves the plan.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+- For conditional counts, use CASE expressions:
+    COUNT(CASE WHEN col IS NULL THEN 1 END)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true``):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>``.
+- Set ``row_level: true`` in the DQCheckPlan output.
+- Set ``metric_key`` to the column name (e.g. ``"email"``).
+
+DQCheckPlan FIELDS:
+- ``check_name``: exact name from ``list_checks`` (no abbreviation).
+- ``sql_query``: the SQL statement to execute later (with leading ``-- check: <name>`` comment).
+- ``metric_key``: the SQL column alias / name to read as the primary metric value.
+- ``row_level``: ``true`` for row-level checks (SELECT returns one row per record), ``false`` for aggregate.
+- ``validator_name``: ``"fixed"``, a registered validator name, or ``"none"``.
+- ``validator_args``: kwargs for the validator factory, or ``{}`` for ``"fixed"``/``"none"``.
+
+VALIDATOR ARGS RULES (critical — missing args cause hard failures in Phase 2):
+- ``list_validators`` returns a ``parameters`` field for each validator showing its exact
+  parameter names, types, and defaults.  Read this carefully before setting ``validator_args``.
+- If a parameter has no default (required), you MUST include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — always set ``{"max_pct": <value>}``.
+- If you are unsure of the right threshold, pick a reasonable default (e.g. 0.05 for pct checks).
+- NEVER leave ``validator_args`` as ``{}`` for a validator that has required parameters.
+- Only use ``{}`` when ``validator_name`` is ``"fixed"`` or ``"none"``.
+"""
+
+_DIALECT_SQL_NOTES: dict[str, str] = {
+    "postgres": (
+        "  - Regex match: `col ~ 'pattern'`; not match: `col !~ 'pattern'` (case-sensitive).\n"
+        "    Case-insensitive variants: `~*` and `!~*`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)` (no CAST needed).\n"
+        "  - String functions: LENGTH(), LOWER(), TRIM(), SUBSTRING().\n"
+    ),
+    "mysql": (
+        "  - Regex match: `col REGEXP 'pattern'`; not match: `col NOT REGEXP 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - String length (multibyte-safe): CHAR_LENGTH(col); byte length: LENGTH(col).\n"
+    ),
+    "sqlite": (
+        "  - No native regex operator; use LIKE or GLOB:\n"
+        "    `col NOT LIKE '%@%'` or `col GLOB '*@*.*'`.\n"
+        "  - Float division: `CAST(numerator AS REAL) / NULLIF(denominator, 0)`.\n"
+    ),
+    "tsql": (
+        "  - No regex operator; use LIKE or PATINDEX:\n"
+        "    `PATINDEX('%@%.%', col) = 0` means no match.\n"
+        "  - Float division: `CAST(numerator AS FLOAT) / NULLIF(denominator, 0)`.\n"
+        "  - Use TOP n instead of LIMIT n; no OFFSET without ORDER BY.\n"
+        "  - String length: LEN(col) (not LENGTH).\n"
+    ),
+    "bigquery": (
+        "  - Regex: `REGEXP_CONTAINS(col, r'pattern')` — returns TRUE if match.\n"
+        "    Negate with `NOT REGEXP_CONTAINS(...)`.\n"
+        "  - Safe division: `SAFE_DIVIDE(numerator, denominator)` (NULL on zero denominator).\n"
+        "  - Quote reserved names with backticks.\n"
+    ),
+    "snowflake": (
+        "  - Regex: `REGEXP_LIKE(col, 'pattern')` (full-string match) or `col RLIKE 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - Unquoted identifiers are uppercased; use double-quotes when case matters.\n"
+    ),
+}
+
+
+def _extract_schema_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by build_schema_context."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Run data-quality checks described in natural language using an LLM agent.
+
+    The agent discovers the database schema, writes and executes SQL queries,
+    applies validators, and produces a
+    :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.  The task
+    fails when any check does not pass, gating downstream tasks on data quality.
+
+    Supply the data-source toolset and DQ toolset together in ``toolsets``::
+
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+        LLMDataQualityOperator(
+            task_id="quality_check",
+            checks=[
+                DQCheckInput(name="email_nulls", description="Check for null emails"),
+                DQCheckInput(
+                    name="row_count",
+                    description="At least 1000 rows",
+                    validator=row_count_check(min_count=1000),
+                ),
+            ],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers"]),
+                SQLDQToolset(),
+            ],
+        )
+
+    When ``toolsets`` is omitted, the operator auto-creates
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` from
+    ``db_conn_id`` and ``table_names``.
+
+    For config-generation backends (e.g. ``SodaDQToolset``), the operator
+    returns the generated config string as its XCom value instead of a report.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional
+        ``validator`` keys).  Names must be unique.
+    :param toolsets: Pydantic-AI toolsets for the agent.  Must include exactly
+        one :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`
+        subclass alongside a data-source toolset.  When ``None``, the operator
+        auto-creates toolsets from ``db_conn_id``.
+    :param db_conn_id: Connection ID for auto-creating an
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`.
+        Ignored when ``toolsets`` is provided.
+    :param table_names: Tables passed to the auto-created
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` as
+        ``allowed_tables``.  Ignored when ``toolsets`` is provided.
+    :param schema_context: Additional schema description injected into the
+        system prompt.  Useful when the data source cannot be introspected
+        at runtime.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        toolsets: list[AbstractToolset] | None = None,
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Pop operator-specific params that must not reach BaseOperator.__init__.
+        # Using kwargs.pop (rather than named params) is necessary because Airflow's
+        # apply_defaults metaclass captures the full **kwargs dict and uses it for
+        # task-map serialization; named params in the dict are NOT removed from the
+        # snapshot, so they would be re-injected on deserialization and reach
+        # BaseOperator as unknown kwargs.
+        durable: bool = kwargs.pop("durable", False)
+
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "Run the data-quality checks.")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            [DQCheckInput.coerce(c) for c in checks] if isinstance(checks, list) else checks  # type: ignore[assignment]
+        )
+        self.toolsets = toolsets
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.durable = durable
+
+        self._validate_checks()
+
+        if not toolsets and db_conn_id is None:
+            raise ValueError("Either toolsets or db_conn_id must be provided.")
+
+    def execute(self, context: Context) -> Any:
+        """
+        Run the LLM agent to execute all data-quality checks.
+
+        When ``require_approval=True``, the task runs in two phases:
+
+        1. **Phase 1** — the LLM discovers schema and selects validators *without
+           executing any SQL queries*.  The resulting plan (check names,
+           descriptions, and validator assignments) is shown to a human reviewer.
+        2. **Phase 2** (in :meth:`execute_complete`) — after approval, SQL queries
+           are executed and validators are applied in pure Python to produce the
+           :class:`~...DQReport`.
+
+        :returns: :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`
+            as a dict when the DQ toolset is in ``"execute"`` mode, or a config
+            string when in ``"generate"`` mode.
+        :raises DQCheckFailedError: When any check fails in ``"execute"`` mode.
+        """
+        if self.require_approval:
+            plan, _ = self._run_plan_phase(context)
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                subject=f"Review DQ plan for task `{self.task_id}`",
+                body=self._build_plan_approval_body(plan),
+            )
+        return self._run_and_report(context)
+
+    def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any:
+        """
+        Phase 2: execute SQL and apply validators after human approval.
+
+        Called automatically by Airflow when the HITL trigger fires.  The base
+        class validates the approval decision (raises on reject or timeout).
+        Then each SQL query from the approved :class:`~...DQPlan` is executed in
+        pure Python, the metric values are extracted, and validators are applied
+        — no LLM calls.
+
+        :param context: Airflow task context.
+        :param generated_output: The :class:`~...DQPlan` JSON that was deferred.
+        :param event: Trigger event payload.
+        :returns: Same as :meth:`execute`.
+        :raises HITLRejectException: If the reviewer rejected the plan.
+        """
+        approved_output = super().execute_complete(context, generated_output, event)  # type: ignore[misc]
+        plan = DQPlan.model_validate_json(approved_output)
+
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        report = self._execute_plan_validators(plan, dq_toolset, toolsets)
+        context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+        if not report.passed:
+            raise DQCheckFailedError(report.failure_summary)
+        return report.model_dump()
+
+    # ------------------------------------------------------------------
+    # Core execution helpers
+    # ------------------------------------------------------------------
+
+    def _run_plan_phase(self, context: Context) -> tuple[DQPlan, BaseDQToolset]:
+        """
+        Phase 1 of the two-phase approval flow.
+
+        Runs the LLM agent in planning mode: schema-discovery tools
+        (``list_tables``, ``get_schema``, ``check_query``) are available, but
+        the ``query`` and ``apply_validator`` tools are hidden.  The agent
+        writes SQL strings and selects validators for each check, then outputs
+        a :class:`~...DQPlan` — without executing any SQL.
+
+        :returns: ``(plan, dq_toolset)`` — the plan for the approval body and the
+            toolset (already configured with checks) for Phase 2.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset._planning_mode = True  # omit apply_validator from tool list
+        dq_toolset.set_checks(self.checks)
+
+        # Hide the SQL 'query' tool so the LLM cannot execute DQ queries.
+        for ts in toolsets:
+            if ts is not dq_toolset and hasattr(ts, "_query"):
+                ts._planning_mode = True  # type: ignore[attr-defined]
+
+        instructions = self._build_plan_system_prompt(toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _ = self._build_durable_agent(DQPlan, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=DQPlan,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay (plan phase): model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        return result.output, dq_toolset
+
+    def _run_and_report(self, context: Context) -> Any:
+        """
+        Resolve toolsets, run the LLM agent, push XCom, and handle failures.
+
+        This is the single place where the agent executes regardless of whether
+        the task took the direct path or the approval-gated path.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        output_type: type = DQReport if dq_toolset.output_mode == "execute" else str
+        instructions = self._build_system_prompt(dq_toolset, toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _storage = self._build_durable_agent(output_type, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=output_type,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay: model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        output = result.output
+
+        if dq_toolset.output_mode == "execute":
+            report: DQReport = output
+            context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+            if not report.passed:
+                raise DQCheckFailedError(report.failure_summary)
+            return report.model_dump()
+
+        return output
+
+    def _build_durable_agent(
+        self,
+        output_type: type,
+        instructions: str,
+        toolsets: list[AbstractToolset],
+    ) -> tuple[Any, Any, Any]:
+        """Build a pydantic-ai Agent with CachingModel and CachingToolset wrappers."""
+        from pydantic_ai import Agent
+
+        from airflow.providers.common.ai.durable.caching_model import CachingModel
+        from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
+        from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+        from airflow.providers.common.ai.durable.storage import DurableStorage
+
+        plan_hash = self._compute_plan_hash(instructions)
+        storage = DurableStorage(cache_id=plan_hash)
+        counter = DurableStepCounter()
+
+        wrapped_model = CachingModel(self.llm_hook.get_conn(), storage=storage, counter=counter)
+        cached_toolsets = [CachingToolset(ts, storage=storage, counter=counter) for ts in toolsets]

Review Comment:
   `durable=True` currently builds `DurableStorage(cache_id=plan_hash)` and never calls `cleanup()`. Because this operator executes live `query` tool calls, a stable cross-run cache_id can replay *old query results* (stale DQ outcomes) and the missing cleanup will also accumulate cache files indefinitely. Consider aligning with existing durable semantics in `operators/agent.py` (task-instance-scoped cache id + cleanup on success) or otherwise scoping the cache so only safe parts (e.g. plan generation) are replayed.



##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py:
##########
@@ -0,0 +1,329 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Example DAGs demonstrating LLMDataQualityOperator usage."""
+
+from __future__ import annotations
+
+import logging
+import re
+
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
+from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+from airflow.providers.common.ai.toolsets.sql import SQLToolset
+from airflow.providers.common.ai.utils.dataquality import (
+    DQCheckInput,
+    exact_check,
+    null_pct_check,
+    register_validator,
+    row_count_check,
+)
+from airflow.providers.common.compat.sdk import dag, task
+from airflow.providers.common.sql.config import DataSourceConfig
+
+# ------------------------------------------------------------------
+# Module-level custom validator used across multiple DAGs
+# ------------------------------------------------------------------

Review Comment:
   Repo prose convention is to write “Dag/Dags” (title case) rather than “DAG/DAGs” (see `AGENTS.md`). Update this comment to match that convention.



##########
providers/common/ai/src/airflow/providers/common/ai/decorators/llm_data_quality.py:
##########
@@ -0,0 +1,166 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+TaskFlow decorator for LLM-driven data-quality checks.
+
+The user writes a function that **returns a list of
+:class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`** objects.
+The decorator runs them through
+:class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Collection, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.compat.sdk import (
+    DecoratedOperator,
+    TaskDecorator,
+    context_merge,
+    determine_kwargs,
+    task_decorator_factory,
+)
+from airflow.sdk.definitions._internal.types import SET_DURING_EXECUTION
+
+if TYPE_CHECKING:
+    from airflow.sdk import Context
+
+
+class _LLMDQDecoratedOperator(DecoratedOperator, LLMDataQualityOperator):
+    """
+    Wraps a callable that returns a list of :class:`DQCheckInput` for LLM data-quality checks.
+
+    The user function is called at execution time to produce the checks list.
+    All other parameters (``llm_conn_id``, ``toolsets``, ``db_conn_id``, etc.)
+    are passed through to
+    :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`.
+
+    :param python_callable: A callable that returns a ``list[DQCheckInput]``.
+    :param op_args: Positional arguments for the callable.
+    :param op_kwargs: Keyword arguments for the callable.
+    """
+
+    template_fields: Sequence[str] = (
+        *DecoratedOperator.template_fields,
+        *LLMDataQualityOperator.template_fields,
+    )
+    template_fields_renderers: ClassVar[dict[str, str]] = {
+        **DecoratedOperator.template_fields_renderers,
+    }
+
+    custom_operator_name: str = "@task.llm_dq"
+
+    def __init__(
+        self,
+        *,
+        python_callable: Callable,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            python_callable=python_callable,
+            op_args=op_args,
+            op_kwargs=op_kwargs,
+            checks=SET_DURING_EXECUTION,
+            **kwargs,
+        )
+
+    def execute(self, context: Context) -> Any:
+        context_merge(context, self.op_kwargs)
+        kwargs = determine_kwargs(self.python_callable, self.op_args, context)
+
+        checks = self.python_callable(*self.op_args, **kwargs)
+
+        if not isinstance(checks, list) or not checks:
+            raise TypeError(
+                "The returned value from the @task.llm_dq callable must be a non-empty list[DQCheckInput]."
+            )
+
+        from airflow.providers.common.ai.utils.dataquality.models import DQCheckInput
+
+        self.checks = [DQCheckInput.coerce(c) for c in checks]
+        self._validate_checks()
+        self.render_template_fields(context)
+        return LLMDataQualityOperator.execute(self, context)
+
+
+def llm_dq_task(
+    python_callable: Callable | None = None,
+    **kwargs,
+) -> TaskDecorator:
+    """
+    Wrap a function that returns a checks list into an LLM data-quality task.
+
+    The function body builds a ``list[DQCheckInput]`` (can use Airflow context,
+    XCom, etc.).  The decorator handles: LLM plan generation, plan caching,
+    SQL execution against the target database, and metric validation.
+
+    Usage::
+
+        from airflow.providers.common.ai.utils.dataquality.models import DQCheckInput
+        from airflow.providers.common.ai.utils.dataquality.validation import (
+            null_pct_check,
+            row_count_check,
+        )
+
+
+        @task.llm_dq(
+            llm_conn_id="openai_default",
+            db_conn_id="postgres_default",
+            table_names=["orders", "customers"],
+        )
+        def dq_checks(ds=None):
+            return [
+                DQCheckInput(
+                    name="row_count",
+                    description=f"The orders table must have at least 1000 rows as of {ds}.",
+                    validator=row_count_check(min_count=1000),
+                ),
+                DQCheckInput(
+                    name="email_nulls",
+                    description="No more than 5% of customer emails should be null.",
+                    validator=null_pct_check(max_pct=0.05),
+                ),
+            ]
+
+    With ``dry_run=True`` to preview the generated plan before execution::
+
+        @task.llm_dq(
+            llm_conn_id="openai_default",
+            db_conn_id="postgres_default",
+            table_names=["orders"],
+            dry_run=True,
+        )
+        def preview_plan():

Review Comment:
   The decorator docstring advertises `dry_run=True`, but neither `LLMDataQualityOperator` nor `LLMOperator` defines a `dry_run` parameter. This will mislead users and likely results in a `TypeError` if passed. Either remove this section or implement/support the parameter.



##########
providers/common/ai/src/airflow/providers/common/ai/durable/storage.py:
##########
@@ -68,13 +68,17 @@ class DurableStorage:
     def __init__(
         self,
         *,
-        dag_id: str,
-        task_id: str,
-        run_id: str,
+        dag_id: str = "",
+        task_id: str = "",
+        run_id: str = "",
         map_index: int = -1,
+        cache_id: str | None = None,
     ) -> None:
-        suffix = f"_{map_index}" if map_index >= 0 else ""
-        self._cache_id = f"{dag_id}_{task_id}_{run_id}{suffix}"
+        if cache_id is not None:
+            self._cache_id = cache_id
+        else:
+            suffix = f"_{map_index}" if map_index >= 0 else ""
+            self._cache_id = f"{dag_id}_{task_id}_{run_id}{suffix}"
         self._cache: dict[str, Any] | None = None

Review Comment:
   `DurableStorage.__init__` now allows empty `dag_id/task_id/run_id` when `cache_id` is omitted, which can silently create a shared `"__.json"` cache file and collide across tasks/runs. It would be safer to validate the identifiers when `cache_id` is not provided (and document the new `cache_id` parameter in the class docstring).



##########
providers/common/ai/tests/unit/common/ai/toolsets/dataquality/test_sql.py:
##########
@@ -0,0 +1,154 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import json
+from unittest.mock import MagicMock
+
+from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import DQCheckInput
+from airflow.providers.common.ai.utils.dataquality.validation import ValidatorRegistry, null_pct_check
+
+
+class TestSQLDQToolsetInit:
+    def test_id(self):
+        ts = SQLDQToolset()
+        assert ts.id == "dq-sql"
+
+    def test_output_mode_is_execute(self):
+        ts = SQLDQToolset()
+        assert ts.output_mode == "execute"
+
+    def test_custom_registry_is_stored(self):
+        registry = ValidatorRegistry()
+        ts = SQLDQToolset(validator_registry=registry)
+        assert ts._registry is registry
+
+
+class TestSQLDQToolsetGetTools:
+    def test_get_tools_includes_list_validators_and_apply_validator(self):
+        ts = SQLDQToolset()
+        tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+        assert "list_checks" in tools
+        assert "list_validators" in tools
+        assert "apply_validator" in tools
+
+    def test_planning_mode_hides_apply_validator(self):
+        ts = SQLDQToolset()
+        ts._planning_mode = True
+        tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+        assert "apply_validator" not in tools
+        assert "list_validators" in tools

Review Comment:
   Several tests pass `MagicMock()` without a spec for `ctx`/`tool`. Since these objects are not inspected by the implementation, prefer `object()` (or a small stub type) to avoid unspec'd mocks (which can hide interface mistakes and is discouraged in this repo’s testing guidelines).



##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py:
##########
@@ -0,0 +1,329 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Example DAGs demonstrating LLMDataQualityOperator usage."""

Review Comment:
   Repo prose convention is to write “Dag/Dags” (title case) rather than “DAG/DAGs” (see `AGENTS.md`). Update this docstring to match that convention.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4184289676

   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 😄 ) it performs nicely llm generates rules and can be executed directly on the object stores.. 
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3230402293


##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,763 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for running data-quality checks from natural language using an LLM agent."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQPlan,
+    DQReport,
+)
+from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
+    from airflow.sdk import Context
+
+_DQ_SYSTEM_PROMPT = """\
+You are a data-quality expert. Evaluate the user's data-quality checks against a \
+live data source using the tools available to you.
+
+WORKFLOW:
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL CHECKS
+      procedure below instead of writing an aggregate query.
+   b. Otherwise, write a SELECT query that computes the relevant metric for this check.
+   c. Optionally call ``check_query`` to validate SQL syntax before executing.
+   d. Execute the query using the ``query`` tool.
+   e. You MUST call ``apply_validator`` for EVERY check — never skip this step, even for
+      fixed validators (``has_fixed_validator: true``).  Use ``validator_name: "fixed"`` for those.
+5. Return a ``DQReport`` with one ``DQCheckResult`` per check — every check must appear exactly once.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+  Example: -- check: null_order_id
+- For conditional counts, use CASE expressions (FILTER WHERE is not universally supported):
+    CORRECT:   COUNT(CASE WHEN col IS NULL THEN 1 END)
+    INCORRECT: COUNT(*) FILTER (WHERE col IS NULL)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- For duplicate percentages:
+    (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / NULLIF(COUNT(*), 0)
+- For float division, cast to avoid integer truncation:
+    CAST(numerator AS DOUBLE) / CAST(denominator AS DOUBLE)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+VALIDATOR SELECTION:
+- ``list_validators`` returns each validator's name, parameter signature, and description.
+  Read the parameter names and types before calling ``apply_validator``.
+- Pass the required ``validator_args`` as a JSON object matching the parameter signature exactly.
+- If a parameter has no default, it is REQUIRED — always include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — you must pass ``{"max_pct": <value>}``.
+  If you are unsure of a threshold, use a safe default (e.g. ``0.05`` for percentage checks).
+- NEVER pass empty ``{}`` args for a validator that has required parameters.
+- For checks marked ``has_fixed_validator: true`` in ``list_checks``, call ``apply_validator``
+  with ``validator_name: "fixed"`` — the pre-assigned validator runs automatically, no args needed.
+- If no suitable validator exists, pass ``validator_name: "none"`` — the check passes by default
+  but its metric value is still recorded.
+
+DQREPORT OUTPUT RULES:
+- ``check_name`` must exactly match the name returned by ``list_checks`` (no abbreviation).
+- Each check must appear in the report exactly once.
+- Set ``passed: false`` and provide a clear ``failure_reason`` for every failed check.
+- Set ``metric_key`` to the SQL column alias you used for this check's metric value.
+- Set ``sql_query`` to the exact SQL statement you executed for this check (including the leading comment).
+- Set ``validator_info`` to ``{"name": "<validator_name>", "args": {<validator_args>}}`` when a validator
+  was applied, or ``null`` when no validator was used (``validator_name: "none"``).
+
+CHECK CATEGORIES (use to guide SQL and validator selection):
+  null_check    — null / missing value counts or percentages
+  uniqueness    — duplicate detection, cardinality checks
+  validity      — regex / format / pattern matching on string columns
+  numeric_range — range, bounds, or statistical checks on numeric columns
+  row_count     — total row counts or existence checks
+  string_format — length, encoding, whitespace, or character-set checks
+  row_level     — per-row anomaly checks; validator receives a list, not an aggregate
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true`` for a check):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>`` that
+  returns one value per row — no GROUP BY, no COUNT, no CASE expressions.
+  Example: -- check: customer_email_format\n  SELECT email FROM customers
+- After executing, extract the column values into a Python list:
+    value = [row["<column>"] for row in result["rows"]]
+  where ``result`` is the JSON object returned by the ``query`` tool.
+- Call ``apply_validator`` with ``value`` set to that list and ``validator_name: "fixed"``.
+- The validator evaluates each item and returns a ``value`` payload containing a
+    ``RowLevelResult`` summary (total/invalid/invalid_pct/sample_violations/sample_size).
+- Set ``metric_key`` to the SQL column name (e.g. ``"email"``).
+- Set ``value`` in the DQCheckResult to the ``value`` object returned by ``apply_validator``.
+"""
+
+_DQ_PLAN_SYSTEM_PROMPT = """\
+You are a data-quality expert.  Your task is to plan *how* to evaluate each
+data-quality check: write the SQL query and choose the validator for each check.
+A human reviewer will inspect your plan before any SQL is executed.
+
+WORKFLOW (PLANNING MODE — no SQL execution, no apply_validator):
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to understand the data model.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL
+      procedure below.  Otherwise, write an aggregate SELECT query for the metric.
+   b. Optionally call ``check_query`` to validate your SQL syntax — but do NOT execute it.
+   c. Select the appropriate validator name and arguments:
+      - For ``has_fixed_validator: true`` checks: use ``validator_name: "fixed"``, empty args.
+      - If no validator fits: use ``validator_name: "none"``.
+5. Return a ``DQPlan`` with one ``DQCheckPlan`` per check — every check must appear exactly once.
+
+IMPORTANT: The ``query`` tool is NOT available.  Do NOT try to execute SQL.
+Do NOT call ``apply_validator`` — it is not available in planning mode.
+SQL will be executed after the human reviewer approves the plan.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+- For conditional counts, use CASE expressions:
+    COUNT(CASE WHEN col IS NULL THEN 1 END)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true``):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>``.
+- Set ``row_level: true`` in the DQCheckPlan output.
+- Set ``metric_key`` to the column name (e.g. ``"email"``).
+
+DQCheckPlan FIELDS:
+- ``check_name``: exact name from ``list_checks`` (no abbreviation).
+- ``sql_query``: the SQL statement to execute later (with leading ``-- check: <name>`` comment).
+- ``metric_key``: the SQL column alias / name to read as the primary metric value.
+- ``row_level``: ``true`` for row-level checks (SELECT returns one row per record), ``false`` for aggregate.
+- ``validator_name``: ``"fixed"``, a registered validator name, or ``"none"``.
+- ``validator_args``: kwargs for the validator factory, or ``{}`` for ``"fixed"``/``"none"``.
+
+VALIDATOR ARGS RULES (critical — missing args cause hard failures in Phase 2):
+- ``list_validators`` returns a ``parameters`` field for each validator showing its exact
+  parameter names, types, and defaults.  Read this carefully before setting ``validator_args``.
+- If a parameter has no default (required), you MUST include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — always set ``{"max_pct": <value>}``.
+- If you are unsure of the right threshold, pick a reasonable default (e.g. 0.05 for pct checks).
+- NEVER leave ``validator_args`` as ``{}`` for a validator that has required parameters.
+- Only use ``{}`` when ``validator_name`` is ``"fixed"`` or ``"none"``.
+"""
+
+_DIALECT_SQL_NOTES: dict[str, str] = {
+    "postgres": (
+        "  - Regex match: `col ~ 'pattern'`; not match: `col !~ 'pattern'` (case-sensitive).\n"
+        "    Case-insensitive variants: `~*` and `!~*`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)` (no CAST needed).\n"
+        "  - String functions: LENGTH(), LOWER(), TRIM(), SUBSTRING().\n"
+    ),
+    "mysql": (
+        "  - Regex match: `col REGEXP 'pattern'`; not match: `col NOT REGEXP 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - String length (multibyte-safe): CHAR_LENGTH(col); byte length: LENGTH(col).\n"
+    ),
+    "sqlite": (
+        "  - No native regex operator; use LIKE or GLOB:\n"
+        "    `col NOT LIKE '%@%'` or `col GLOB '*@*.*'`.\n"
+        "  - Float division: `CAST(numerator AS REAL) / NULLIF(denominator, 0)`.\n"
+    ),
+    "tsql": (
+        "  - No regex operator; use LIKE or PATINDEX:\n"
+        "    `PATINDEX('%@%.%', col) = 0` means no match.\n"
+        "  - Float division: `CAST(numerator AS FLOAT) / NULLIF(denominator, 0)`.\n"
+        "  - Use TOP n instead of LIMIT n; no OFFSET without ORDER BY.\n"
+        "  - String length: LEN(col) (not LENGTH).\n"
+    ),
+    "bigquery": (
+        "  - Regex: `REGEXP_CONTAINS(col, r'pattern')` — returns TRUE if match.\n"
+        "    Negate with `NOT REGEXP_CONTAINS(...)`.\n"
+        "  - Safe division: `SAFE_DIVIDE(numerator, denominator)` (NULL on zero denominator).\n"
+        "  - Quote reserved names with backticks.\n"
+    ),
+    "snowflake": (
+        "  - Regex: `REGEXP_LIKE(col, 'pattern')` (full-string match) or `col RLIKE 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - Unquoted identifiers are uppercased; use double-quotes when case matters.\n"
+    ),
+}
+
+
+def _extract_schema_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by build_schema_context."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Run data-quality checks described in natural language using an LLM agent.
+
+    The agent discovers the database schema, writes and executes SQL queries,
+    applies validators, and produces a
+    :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.  The task
+    fails when any check does not pass, gating downstream tasks on data quality.
+
+    Supply the data-source toolset and DQ toolset together in ``toolsets``::
+
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+        LLMDataQualityOperator(
+            task_id="quality_check",
+            checks=[
+                DQCheckInput(name="email_nulls", description="Check for null emails"),
+                DQCheckInput(
+                    name="row_count",
+                    description="At least 1000 rows",
+                    validator=row_count_check(min_count=1000),
+                ),
+            ],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers"]),
+                SQLDQToolset(),
+            ],
+        )
+
+    When ``toolsets`` is omitted, the operator auto-creates
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` from
+    ``db_conn_id`` and ``table_names``.
+
+    For config-generation backends (e.g. ``SodaDQToolset``), the operator
+    returns the generated config string as its XCom value instead of a report.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional
+        ``validator`` keys).  Names must be unique.
+    :param toolsets: Pydantic-AI toolsets for the agent.  Must include exactly
+        one :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`
+        subclass alongside a data-source toolset.  When ``None``, the operator
+        auto-creates toolsets from ``db_conn_id``.
+    :param db_conn_id: Connection ID for auto-creating an
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`.
+        Ignored when ``toolsets`` is provided.
+    :param table_names: Tables passed to the auto-created
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` as
+        ``allowed_tables``.  Ignored when ``toolsets`` is provided.
+    :param schema_context: Additional schema description injected into the
+        system prompt.  Useful when the data source cannot be introspected
+        at runtime.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        toolsets: list[AbstractToolset] | None = None,
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Pop operator-specific params that must not reach BaseOperator.__init__.
+        # Using kwargs.pop (rather than named params) is necessary because Airflow's
+        # apply_defaults metaclass captures the full **kwargs dict and uses it for
+        # task-map serialization; named params in the dict are NOT removed from the
+        # snapshot, so they would be re-injected on deserialization and reach
+        # BaseOperator as unknown kwargs.
+        durable: bool = kwargs.pop("durable", False)
+
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "Run the data-quality checks.")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            [DQCheckInput.coerce(c) for c in checks] if isinstance(checks, list) else checks  # type: ignore[assignment]
+        )
+        self.toolsets = toolsets
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.durable = durable
+
+        self._validate_checks()
+
+        if not toolsets and db_conn_id is None:
+            raise ValueError("Either toolsets or db_conn_id must be provided.")
+
+    def execute(self, context: Context) -> Any:
+        """
+        Run the LLM agent to execute all data-quality checks.
+
+        When ``require_approval=True``, the task runs in two phases:
+
+        1. **Phase 1** — the LLM discovers schema, executes SQL queries, and selects
+           validators (but does *not* call ``apply_validator``).  The resulting plan
+           (SQL queries + validator choices) is shown to a human reviewer.
+        2. **Phase 2** (in :meth:`execute_complete`) — after approval, validators are
+           applied in pure Python and the :class:`~...DQReport` is produced.
+
+        :returns: :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`
+            as a dict when the DQ toolset is in ``"execute"`` mode, or a config
+            string when in ``"generate"`` mode.
+        :raises DQCheckFailedError: When any check fails in ``"execute"`` mode.
+        """
+        if self.require_approval:
+            plan, _ = self._run_plan_phase(context)
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                subject=f"Review DQ plan for task `{self.task_id}`",
+                body=self._build_plan_approval_body(plan),
+            )
+        return self._run_and_report(context)
+
+    def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any:
+        """
+        Phase 2: execute SQL and apply validators after human approval.
+
+        Called automatically by Airflow when the HITL trigger fires.  The base
+        class validates the approval decision (raises on reject or timeout).
+        Then each SQL query from the approved :class:`~...DQPlan` is executed in
+        pure Python, the metric values are extracted, and validators are applied
+        — no LLM calls.
+
+        :param context: Airflow task context.
+        :param generated_output: The :class:`~...DQPlan` JSON that was deferred.
+        :param event: Trigger event payload.
+        :returns: Same as :meth:`execute`.
+        :raises HITLRejectException: If the reviewer rejected the plan.
+        """
+        approved_output = super().execute_complete(context, generated_output, event)  # type: ignore[misc]
+        plan = DQPlan.model_validate_json(approved_output)
+
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        report = self._execute_plan_validators(plan, dq_toolset, toolsets)
+        context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+        if not report.passed:
+            raise DQCheckFailedError(report.failure_summary)
+        return report.model_dump()
+
+    # ------------------------------------------------------------------
+    # Core execution helpers
+    # ------------------------------------------------------------------
+
+    def _run_plan_phase(self, context: Context) -> tuple[DQPlan, BaseDQToolset]:
+        """
+        Phase 1 of the two-phase approval flow.
+
+        Runs the LLM agent in planning mode: schema-discovery tools
+        (``list_tables``, ``get_schema``, ``check_query``) are available, but
+        the ``query`` and ``apply_validator`` tools are hidden.  The agent
+        writes SQL strings and selects validators for each check, then outputs
+        a :class:`~...DQPlan` — without executing any SQL.
+
+        :returns: ``(plan, dq_toolset)`` — the plan for the approval body and the
+            toolset (already configured with checks) for Phase 2.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset._planning_mode = True  # omit apply_validator from tool list
+        dq_toolset.set_checks(self.checks)
+
+        # Hide the SQL 'query' tool so the LLM cannot execute DQ queries.
+        for ts in toolsets:
+            if ts is not dq_toolset and hasattr(ts, "_query"):
+                ts._planning_mode = True  # type: ignore[attr-defined]
+
+        instructions = self._build_plan_system_prompt(toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _ = self._build_durable_agent(DQPlan, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=DQPlan,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay (plan phase): model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        return result.output, dq_toolset
+
+    def _run_and_report(self, context: Context) -> Any:
+        """
+        Resolve toolsets, run the LLM agent, push XCom, and handle failures.
+
+        This is the single place where the agent executes regardless of whether
+        the task took the direct path or the approval-gated path.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        output_type: type = DQReport if dq_toolset.output_mode == "execute" else str
+        instructions = self._build_system_prompt(dq_toolset, toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+
+        if self.durable:
+            agent, counter, _storage = self._build_durable_agent(output_type, instructions, logged_toolsets)
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=output_type,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay: model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        output = result.output
+
+        if dq_toolset.output_mode == "execute":
+            report: DQReport = output
+            context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+            if not report.passed:
+                raise DQCheckFailedError(report.failure_summary)
+            return report.model_dump()
+
+        return output
+
+    def _build_durable_agent(
+        self,
+        output_type: type,
+        instructions: str,
+        toolsets: list[AbstractToolset],
+    ) -> tuple[Any, Any, Any]:
+        """Build a pydantic-ai Agent with CachingModel and CachingToolset wrappers."""
+        from pydantic_ai import Agent
+
+        from airflow.providers.common.ai.durable.caching_model import CachingModel
+        from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
+        from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+        from airflow.providers.common.ai.durable.storage import DurableStorage
+
+        plan_hash = self._compute_plan_hash(instructions)
+        storage = DurableStorage(cache_id=plan_hash)
+        counter = DurableStepCounter()
+
+        wrapped_model = CachingModel(self.llm_hook.get_conn(), storage=storage, counter=counter)
+        cached_toolsets = [CachingToolset(ts, storage=storage, counter=counter) for ts in toolsets]
+
+        agent = Agent(
+            wrapped_model,
+            output_type=output_type,
+            instructions=instructions,
+            toolsets=cached_toolsets,
+            **self.agent_params,
+        )
+        return agent, counter, storage
+
+    def _compute_plan_hash(self, instructions: str) -> str:
+        """Return a stable short hash of the check plan for cross-run cache keying."""
+        checks_data = sorted(
+            [{"name": c.name, "description": c.description} for c in self.checks],
+            key=lambda x: x["name"],
+        )
+        payload = json.dumps(
+            {
+                "checks": checks_data,
+                "schema_context": self.schema_context or "",
+                "instructions": instructions,
+            },
+            sort_keys=True,
+        )
+        digest = hashlib.sha256(payload.encode()).hexdigest()[:16]
+        return f"dq_{digest}"

Review Comment:
   The PR description mentions caching generated DQ plans in Airflow Variables (e.g. `dq_plan_<version>_<sha256...>`), but this implementation appears to rely on the durable ObjectStorage cache (`DurableStorage(cache_id=...)`) and does not persist a DQPlan in Variables. Please either update the PR description to match the implemented caching approach, or add the Variable-based plan caching described in the PR text.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,763 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for running data-quality checks from natural language using an LLM agent."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQPlan,
+    DQReport,
+)
+from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
+    from airflow.sdk import Context
+
+_DQ_SYSTEM_PROMPT = """\
+You are a data-quality expert. Evaluate the user's data-quality checks against a \
+live data source using the tools available to you.
+
+WORKFLOW:
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL CHECKS
+      procedure below instead of writing an aggregate query.
+   b. Otherwise, write a SELECT query that computes the relevant metric for this check.
+   c. Optionally call ``check_query`` to validate SQL syntax before executing.
+   d. Execute the query using the ``query`` tool.
+   e. You MUST call ``apply_validator`` for EVERY check — never skip this step, even for
+      fixed validators (``has_fixed_validator: true``).  Use ``validator_name: "fixed"`` for those.
+5. Return a ``DQReport`` with one ``DQCheckResult`` per check — every check must appear exactly once.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+  Example: -- check: null_order_id
+- For conditional counts, use CASE expressions (FILTER WHERE is not universally supported):
+    CORRECT:   COUNT(CASE WHEN col IS NULL THEN 1 END)
+    INCORRECT: COUNT(*) FILTER (WHERE col IS NULL)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- For duplicate percentages:
+    (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / NULLIF(COUNT(*), 0)
+- For float division, cast to avoid integer truncation:
+    CAST(numerator AS DOUBLE) / CAST(denominator AS DOUBLE)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+VALIDATOR SELECTION:
+- ``list_validators`` returns each validator's name, parameter signature, and description.
+  Read the parameter names and types before calling ``apply_validator``.
+- Pass the required ``validator_args`` as a JSON object matching the parameter signature exactly.
+- If a parameter has no default, it is REQUIRED — always include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — you must pass ``{"max_pct": <value>}``.
+  If you are unsure of a threshold, use a safe default (e.g. ``0.05`` for percentage checks).
+- NEVER pass empty ``{}`` args for a validator that has required parameters.
+- For checks marked ``has_fixed_validator: true`` in ``list_checks``, call ``apply_validator``
+  with ``validator_name: "fixed"`` — the pre-assigned validator runs automatically, no args needed.
+- If no suitable validator exists, pass ``validator_name: "none"`` — the check passes by default
+  but its metric value is still recorded.
+
+DQREPORT OUTPUT RULES:
+- ``check_name`` must exactly match the name returned by ``list_checks`` (no abbreviation).
+- Each check must appear in the report exactly once.
+- Set ``passed: false`` and provide a clear ``failure_reason`` for every failed check.
+- Set ``metric_key`` to the SQL column alias you used for this check's metric value.
+- Set ``sql_query`` to the exact SQL statement you executed for this check (including the leading comment).
+- Set ``validator_info`` to ``{"name": "<validator_name>", "args": {<validator_args>}}`` when a validator
+  was applied, or ``null`` when no validator was used (``validator_name: "none"``).
+
+CHECK CATEGORIES (use to guide SQL and validator selection):
+  null_check    — null / missing value counts or percentages
+  uniqueness    — duplicate detection, cardinality checks
+  validity      — regex / format / pattern matching on string columns
+  numeric_range — range, bounds, or statistical checks on numeric columns
+  row_count     — total row counts or existence checks
+  string_format — length, encoding, whitespace, or character-set checks
+  row_level     — per-row anomaly checks; validator receives a list, not an aggregate
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true`` for a check):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>`` that
+  returns one value per row — no GROUP BY, no COUNT, no CASE expressions.
+  Example: -- check: customer_email_format\n  SELECT email FROM customers
+- After executing, extract the column values into a Python list:
+    value = [row["<column>"] for row in result["rows"]]
+  where ``result`` is the JSON object returned by the ``query`` tool.
+- Call ``apply_validator`` with ``value`` set to that list and ``validator_name: "fixed"``.
+- The validator evaluates each item and returns a ``value`` payload containing a
+    ``RowLevelResult`` summary (total/invalid/invalid_pct/sample_violations/sample_size).
+- Set ``metric_key`` to the SQL column name (e.g. ``"email"``).
+- Set ``value`` in the DQCheckResult to the ``value`` object returned by ``apply_validator``.
+"""
+
+_DQ_PLAN_SYSTEM_PROMPT = """\
+You are a data-quality expert.  Your task is to plan *how* to evaluate each
+data-quality check: write the SQL query and choose the validator for each check.
+A human reviewer will inspect your plan before any SQL is executed.
+
+WORKFLOW (PLANNING MODE — no SQL execution, no apply_validator):
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to understand the data model.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL
+      procedure below.  Otherwise, write an aggregate SELECT query for the metric.
+   b. Optionally call ``check_query`` to validate your SQL syntax — but do NOT execute it.
+   c. Select the appropriate validator name and arguments:
+      - For ``has_fixed_validator: true`` checks: use ``validator_name: "fixed"``, empty args.
+      - If no validator fits: use ``validator_name: "none"``.
+5. Return a ``DQPlan`` with one ``DQCheckPlan`` per check — every check must appear exactly once.
+
+IMPORTANT: The ``query`` tool is NOT available.  Do NOT try to execute SQL.
+Do NOT call ``apply_validator`` — it is not available in planning mode.
+SQL will be executed after the human reviewer approves the plan.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+- For conditional counts, use CASE expressions:
+    COUNT(CASE WHEN col IS NULL THEN 1 END)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true``):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>``.
+- Set ``row_level: true`` in the DQCheckPlan output.
+- Set ``metric_key`` to the column name (e.g. ``"email"``).
+
+DQCheckPlan FIELDS:
+- ``check_name``: exact name from ``list_checks`` (no abbreviation).
+- ``sql_query``: the SQL statement to execute later (with leading ``-- check: <name>`` comment).
+- ``metric_key``: the SQL column alias / name to read as the primary metric value.
+- ``row_level``: ``true`` for row-level checks (SELECT returns one row per record), ``false`` for aggregate.
+- ``validator_name``: ``"fixed"``, a registered validator name, or ``"none"``.
+- ``validator_args``: kwargs for the validator factory, or ``{}`` for ``"fixed"``/``"none"``.
+
+VALIDATOR ARGS RULES (critical — missing args cause hard failures in Phase 2):
+- ``list_validators`` returns a ``parameters`` field for each validator showing its exact
+  parameter names, types, and defaults.  Read this carefully before setting ``validator_args``.
+- If a parameter has no default (required), you MUST include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — always set ``{"max_pct": <value>}``.
+- If you are unsure of the right threshold, pick a reasonable default (e.g. 0.05 for pct checks).
+- NEVER leave ``validator_args`` as ``{}`` for a validator that has required parameters.
+- Only use ``{}`` when ``validator_name`` is ``"fixed"`` or ``"none"``.
+"""
+
+_DIALECT_SQL_NOTES: dict[str, str] = {
+    "postgres": (
+        "  - Regex match: `col ~ 'pattern'`; not match: `col !~ 'pattern'` (case-sensitive).\n"
+        "    Case-insensitive variants: `~*` and `!~*`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)` (no CAST needed).\n"
+        "  - String functions: LENGTH(), LOWER(), TRIM(), SUBSTRING().\n"
+    ),
+    "mysql": (
+        "  - Regex match: `col REGEXP 'pattern'`; not match: `col NOT REGEXP 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - String length (multibyte-safe): CHAR_LENGTH(col); byte length: LENGTH(col).\n"
+    ),
+    "sqlite": (
+        "  - No native regex operator; use LIKE or GLOB:\n"
+        "    `col NOT LIKE '%@%'` or `col GLOB '*@*.*'`.\n"
+        "  - Float division: `CAST(numerator AS REAL) / NULLIF(denominator, 0)`.\n"
+    ),
+    "tsql": (
+        "  - No regex operator; use LIKE or PATINDEX:\n"
+        "    `PATINDEX('%@%.%', col) = 0` means no match.\n"
+        "  - Float division: `CAST(numerator AS FLOAT) / NULLIF(denominator, 0)`.\n"
+        "  - Use TOP n instead of LIMIT n; no OFFSET without ORDER BY.\n"
+        "  - String length: LEN(col) (not LENGTH).\n"
+    ),
+    "bigquery": (
+        "  - Regex: `REGEXP_CONTAINS(col, r'pattern')` — returns TRUE if match.\n"
+        "    Negate with `NOT REGEXP_CONTAINS(...)`.\n"
+        "  - Safe division: `SAFE_DIVIDE(numerator, denominator)` (NULL on zero denominator).\n"
+        "  - Quote reserved names with backticks.\n"
+    ),
+    "snowflake": (
+        "  - Regex: `REGEXP_LIKE(col, 'pattern')` (full-string match) or `col RLIKE 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - Unquoted identifiers are uppercased; use double-quotes when case matters.\n"
+    ),
+}
+
+
+def _extract_schema_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by build_schema_context."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Run data-quality checks described in natural language using an LLM agent.
+
+    The agent discovers the database schema, writes and executes SQL queries,
+    applies validators, and produces a
+    :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.  The task
+    fails when any check does not pass, gating downstream tasks on data quality.
+
+    Supply the data-source toolset and DQ toolset together in ``toolsets``::
+
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+        LLMDataQualityOperator(
+            task_id="quality_check",
+            checks=[
+                DQCheckInput(name="email_nulls", description="Check for null emails"),
+                DQCheckInput(
+                    name="row_count",
+                    description="At least 1000 rows",
+                    validator=row_count_check(min_count=1000),
+                ),
+            ],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers"]),
+                SQLDQToolset(),
+            ],
+        )
+
+    When ``toolsets`` is omitted, the operator auto-creates
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` from
+    ``db_conn_id`` and ``table_names``.
+
+    For config-generation backends (e.g. ``SodaDQToolset``), the operator
+    returns the generated config string as its XCom value instead of a report.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional
+        ``validator`` keys).  Names must be unique.
+    :param toolsets: Pydantic-AI toolsets for the agent.  Must include exactly
+        one :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`
+        subclass alongside a data-source toolset.  When ``None``, the operator
+        auto-creates toolsets from ``db_conn_id``.
+    :param db_conn_id: Connection ID for auto-creating an
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`.
+        Ignored when ``toolsets`` is provided.
+    :param table_names: Tables passed to the auto-created
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` as
+        ``allowed_tables``.  Ignored when ``toolsets`` is provided.
+    :param schema_context: Additional schema description injected into the
+        system prompt.  Useful when the data source cannot be introspected
+        at runtime.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        toolsets: list[AbstractToolset] | None = None,
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Pop operator-specific params that must not reach BaseOperator.__init__.
+        # Using kwargs.pop (rather than named params) is necessary because Airflow's
+        # apply_defaults metaclass captures the full **kwargs dict and uses it for
+        # task-map serialization; named params in the dict are NOT removed from the
+        # snapshot, so they would be re-injected on deserialization and reach
+        # BaseOperator as unknown kwargs.
+        durable: bool = kwargs.pop("durable", False)
+
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "Run the data-quality checks.")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            [DQCheckInput.coerce(c) for c in checks] if isinstance(checks, list) else checks  # type: ignore[assignment]
+        )
+        self.toolsets = toolsets
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.durable = durable
+
+        self._validate_checks()
+
+        if not toolsets and db_conn_id is None:
+            raise ValueError("Either toolsets or db_conn_id must be provided.")
+
+    def execute(self, context: Context) -> Any:
+        """
+        Run the LLM agent to execute all data-quality checks.
+
+        When ``require_approval=True``, the task runs in two phases:
+
+        1. **Phase 1** — the LLM discovers schema, executes SQL queries, and selects
+           validators (but does *not* call ``apply_validator``).  The resulting plan
+           (SQL queries + validator choices) is shown to a human reviewer.
+        2. **Phase 2** (in :meth:`execute_complete`) — after approval, validators are
+           applied in pure Python and the :class:`~...DQReport` is produced.

Review Comment:
   The `execute()` docstring says Phase 1 "executes SQL queries" when `require_approval=True`, but `_run_plan_phase()` explicitly hides the `query` tool and states "without executing any SQL". Please update the docstring to reflect the actual two-phase behavior (plan only in Phase 1; SQL execution happens in `execute_complete`).
   



##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py:
##########
@@ -0,0 +1,328 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Example DAGs demonstrating LLMDataQualityOperator usage."""
+
+from __future__ import annotations
+
+import logging
+
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
+from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+from airflow.providers.common.ai.toolsets.sql import SQLToolset
+from airflow.providers.common.ai.utils.dataquality import (
+    DQCheckInput,
+    exact_check,
+    null_pct_check,
+    register_validator,
+    row_count_check,
+)
+from airflow.providers.common.compat.sdk import dag, task
+from airflow.providers.common.sql.config import DataSourceConfig
+
+# ------------------------------------------------------------------
+# Module-level custom validator used across multiple DAGs
+# ------------------------------------------------------------------
+
+
+# [START howto_operator_llm_dq_email_format_validator]
+@register_validator(
+    "email_format_check",
+    llm_context=(
+        "Row-level validator. Run `SELECT email FROM <table>` (no aggregation), "
+        "collect all email values into a list, and pass that list to apply_validator. "
+        "The validator checks each value against a basic email regex and passes when "
+        "the fraction of invalid emails is at or below max_invalid_pct."
+    ),
+    check_category="validity",
+    row_level=True,
+)
+def email_format_check(*, max_invalid_pct: float = 0.0):
+    """Return a per-row predicate used by row-level aggregation in SQLDQToolset."""
+    import re as _re
+
+    _pattern = _re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")

Review Comment:
   There is a function-local import (`import re as _re`) inside `email_format_check()`. In this codebase, imports are expected at module scope unless there’s a clear need (circular import avoidance / optional dependency / lazy loading). Please move this import to the top of the module (and keep the compiled pattern at module scope if possible).



##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py:
##########
@@ -0,0 +1,328 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Example DAGs demonstrating LLMDataQualityOperator usage."""
+
+from __future__ import annotations
+
+import logging
+
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
+from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+from airflow.providers.common.ai.toolsets.sql import SQLToolset
+from airflow.providers.common.ai.utils.dataquality import (
+    DQCheckInput,
+    exact_check,
+    null_pct_check,
+    register_validator,
+    row_count_check,
+)
+from airflow.providers.common.compat.sdk import dag, task
+from airflow.providers.common.sql.config import DataSourceConfig
+
+# ------------------------------------------------------------------
+# Module-level custom validator used across multiple DAGs
+# ------------------------------------------------------------------
+
+
+# [START howto_operator_llm_dq_email_format_validator]
+@register_validator(
+    "email_format_check",
+    llm_context=(
+        "Row-level validator. Run `SELECT email FROM <table>` (no aggregation), "
+        "collect all email values into a list, and pass that list to apply_validator. "
+        "The validator checks each value against a basic email regex and passes when "
+        "the fraction of invalid emails is at or below max_invalid_pct."
+    ),
+    check_category="validity",
+    row_level=True,
+)
+def email_format_check(*, max_invalid_pct: float = 0.0):
+    """Return a per-row predicate used by row-level aggregation in SQLDQToolset."""
+    import re as _re
+
+    _pattern = _re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")
+
+    def _check(value: object) -> bool:
+        if value in (None, ""):
+            return False
+        return bool(_pattern.match(str(value)))
+
+    return _check
+
+
+# [END howto_operator_llm_dq_email_format_validator]
+
+
+# [START howto_operator_llm_dq_postgres_basic]
+@dag
+def example_llm_dq_postgres_basic():
+    """
+    Run data-quality checks on PostgreSQL tables using explicit toolsets.
+
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` handles schema
+    discovery and SQL execution, while
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` provides the
+    ``list_validators`` and ``apply_validator`` tools to the LLM agent.
+
+    Connections required:
+
+    - ``pydanticai_default``: Pydantic AI connection.
+    - ``postgres_default``: PostgreSQL connection pointing at the seeded database.
+    """
+    LLMDataQualityOperator(
+        task_id="validate_orders_and_customers",
+        llm_conn_id="pydanticai_default",
+        toolsets=[
+            SQLToolset(db_conn_id="postgres_default", allowed_tables=["orders", "customers"]),
+            SQLDQToolset(),
+        ],
+        checks=[
+            DQCheckInput(
+                name="null_order_id",
+                description="Check the percentage of rows in orders where order_id is NULL",
+            ),
+            DQCheckInput(
+                name="duplicate_orders",
+                description="Calculate the percentage of duplicate order_id values in orders",
+            ),
+            DQCheckInput(
+                name="null_customer_email",
+                description="Check the percentage of rows in customers where email is NULL",
+            ),
+            DQCheckInput(
+                name="min_order_count",
+                description="Count the total number of rows in the orders table",
+                validator=row_count_check(min_count=1),
+            ),
+            DQCheckInput(
+                name="negative_amount",
+                description="Count rows in orders where order_amount is negative or NULL",
+                validator=exact_check(expected=0),
+            ),
+            DQCheckInput(
+                name="customer_email_format",
+                description="Validate that each email in the customers table matches a valid email format",
+                validator=email_format_check(max_invalid_pct=0.01),
+            ),
+        ],
+    )
+
+
+# [END howto_operator_llm_dq_postgres_basic]
+
+example_llm_dq_postgres_basic()
+
+
+# [START howto_operator_llm_dq_explicit_toolsets]
+@dag
+def example_llm_dq_explicit_toolsets():
+    """
+    Run data-quality checks with explicitly composed toolsets.
+
+    Passing ``toolsets`` directly gives full control over which tables the
+    agent can access and which validators are available.  Use this pattern
+    when you need fine-grained permissions or want to add custom toolsets.
+
+    Connections required:
+
+    - ``pydanticai_default``: Pydantic AI connection.
+    - ``postgres_default``: PostgreSQL connection.
+    """
+    LLMDataQualityOperator(
+        task_id="validate_products",
+        llm_conn_id="pydanticai_default",
+        toolsets=[
+            SQLToolset(db_conn_id="postgres_default", allowed_tables=["products"]),
+            SQLDQToolset(),
+        ],
+        checks=[
+            DQCheckInput(
+                name="null_sku",
+                description="Check the percentage of rows where sku is NULL",
+                validator=null_pct_check(max_pct=0.0),
+            ),
+            DQCheckInput(
+                name="negative_price",
+                description="Count rows where price is negative or NULL",
+                validator=exact_check(expected=0),
+            ),
+            DQCheckInput(
+                name="min_product_count",
+                description="Ensure the products table has at least 100 rows",
+                validator=row_count_check(min_count=100),
+            ),
+        ],
+    )
+
+
+# [END howto_operator_llm_dq_explicit_toolsets]
+
+example_llm_dq_explicit_toolsets()
+
+
+# [START howto_operator_llm_dq_task_decorator]
+@dag
+def example_llm_dq_task_decorator():
+    """
+    Combine a ``@task``-decorated step with a ``@task.llm_dq``-decorated DQ task.
+
+    ``notify_start`` is a plain ``@task`` that runs first.
+    ``validate_users`` is decorated with ``@task.llm_dq``: the function body
+    returns the checks list, and the decorator handles the LLM agent run,
+    SQL execution, and metric validation.
+
+    Connections required:
+
+    - ``pydanticai_default``: Pydantic AI connection.
+    - ``postgres_default``: PostgreSQL connection for the ``users`` table.
+    """
+
+    @task
+    def notify_start() -> dict:
+        log = logging.getLogger(__name__)
+        log.info("Starting DQ validation for the users table.")
+        return {"status": "started"}
+
+    @task.llm_dq(
+        llm_conn_id="pydanticai_default",
+        toolsets=[
+            SQLToolset(db_conn_id="postgres_default", allowed_tables=["users"]),
+            SQLDQToolset(),
+        ],
+    )
+    def validate_users():
+        return [
+            DQCheckInput(
+                name="null_username",
+                description="Check the percentage of rows where username is NULL",
+            ),
+            DQCheckInput(
+                name="null_email",
+                description="Check the percentage of rows where email is NULL",
+                validator=null_pct_check(max_pct=0.02),
+            ),
+        ]
+
+    notify_start() >> validate_users()
+
+
+# [END howto_operator_llm_dq_task_decorator]
+
+example_llm_dq_task_decorator()
+
+
+# [START howto_operator_llm_dq_require_approval]
+@dag
+def example_llm_dq_require_approval():
+    """
+    Gate check execution on human review using the built-in HITL mechanism.
+
+    When ``require_approval=True`` the task defers *before* the LLM agent runs
+    and surfaces the check plan (names, descriptions, and pre-assigned validators)
+    to a human reviewer in the Airflow UI.  The agent only executes after the
+    reviewer approves.  If rejected, the task raises
+    :class:`~airflow.providers.standard.exceptions.HITLRejectException`.

Review Comment:
   This docstring claims that with `require_approval=True` the task defers *before* the LLM agent runs, but `LLMDataQualityOperator.execute()` runs an LLM planning phase (`_run_plan_phase`) before deferring. Please adjust the wording to clarify that it defers before **SQL execution / validator application**, not before plan generation.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.

kaxil closed pull request #62963: Aip 99 llmdataqualityoperator
URL: https://github.com/apache/airflow/pull/62963


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3071668952


##########
providers/common/ai/tests/unit/common/ai/utils/test_dq_planner.py:
##########
@@ -0,0 +1,1444 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.utils.dq_models import DQCheck, DQCheckGroup, DQPlan
+from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+from airflow.providers.common.ai.utils.sql_validation import SQLSafetyError
+from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+
+def _make_plan(*check_names: str) -> DQPlan:
+    """Helper: build a minimal DQPlan with one group per check."""
+    groups = [
+        DQCheckGroup(
+            group_id="numeric_aggregate",
+            query=f"SELECT COUNT(*) AS {name}_count FROM t",
+            checks=[DQCheck(check_name=name, metric_key=f"{name}_count", group_id="numeric_aggregate")],
+        )
+        for name in check_names
+    ]
+    return DQPlan(groups=groups)
+
+
+def _make_llm_hook(plan: DQPlan) -> MagicMock:
+    """Helper: mock PydanticAIHook that returns *plan* from agent.run_sync."""
+    mock_usage = MagicMock(requests=1, tool_calls=0, input_tokens=100, output_tokens=50, total_tokens=150)
+    mock_result = MagicMock(spec=["output", "all_messages", "usage", "response"])
+    mock_result.output = plan
+    mock_result.all_messages.return_value = []
+    mock_result.usage.return_value = mock_usage
+    mock_result.response.model_name = "test-model"
+    mock_agent = MagicMock(spec=["run_sync"])
+    mock_agent.run_sync.return_value = mock_result
+    mock_hook = MagicMock(spec=PydanticAIHook)
+    mock_hook.create_agent.return_value = mock_agent
+    return mock_hook
+
+
+class TestSQLDQPlannerBuildSchema:
+    def test_returns_manual_schema_context_verbatim(self):
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=None)
+        result = planner.build_schema_context(
+            table_names=None,
+            schema_context="Table: t\nColumns: id INT",
+        )
+        assert result == "Table: t\nColumns: id INT"
+
+    def test_introspects_via_db_hook_when_no_manual_context(self):
+        mock_db_hook = MagicMock()
+        mock_db_hook.get_table_schema.return_value = [{"name": "id", "type": "INT"}]
+
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=mock_db_hook)
+        result = planner.build_schema_context(
+            table_names=["customers"],
+            schema_context=None,
+        )
+
+        mock_db_hook.get_table_schema.assert_called_once_with("customers")
+        assert "customers" in result
+        assert "id INT" in result
+
+    def test_manual_context_takes_priority_over_db_hook(self):
+        mock_db_hook = MagicMock()
+
+        planner = SQLDQPlanner(llm_hook=MagicMock(spec=PydanticAIHook), db_hook=mock_db_hook)
+        result = planner.build_schema_context(
+            table_names=["t"],

Review Comment:
   Same issue here: `mock_db_hook = MagicMock()` is unspec'd, so `assert_not_called()` etc may pass even if the production API changes. Use `MagicMock(spec=DbApiHook)` (or at least `spec=["get_table_schema"]`) to keep the test aligned with the real hook interface.



##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,908 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        self._cached_datafusion_engine: DataFusionEngine | None = None
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+

Review Comment:
   `generate_plan()` logs the full system prompt and user message at INFO. The system prompt is very large and may include schema context and custom validator hints, which can bloat logs and potentially expose sensitive schema details. Consider logging these at DEBUG (or behind an opt-in flag) and keeping INFO to a short summary (e.g. counts + plan_hash).



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,703 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for generating and executing data-quality checks from natural language using LLMs."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.db_schema import get_db_hook
+from airflow.providers.common.ai.utils.dq_models import (
+    DQCheckFailedError,
+    DQCheckGroup,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+    RowLevelResult,
+    UnexpectedResult,
+)
+from airflow.providers.common.ai.utils.dq_validation import default_registry
+from airflow.providers.common.compat.sdk import Variable
+
+try:
+    from airflow.providers.common.ai.utils.dq_planner import SQLDQPlanner
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.sdk import Context
+
+_PLAN_VARIABLE_PREFIX = "dq_plan_"
+_PLAN_VARIABLE_KEY_MAX_LEN = 200  # stay well under Airflow Variable key length limit
+
+
+def _describe_validator(validator: Callable[[Any], bool]) -> str:
+    """Return a human-readable validator label for failure messages."""
+    display = getattr(validator, "_validator_display", None)
+    if isinstance(display, str) and display:
+        return display
+    validator_name = getattr(validator, "_validator_name", None)
+    if isinstance(validator_name, str) and validator_name:
+        return validator_name
+    validator_name = getattr(validator, "__name__", None)
+    if isinstance(validator_name, str) and validator_name:
+        return validator_name
+    return repr(validator)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Generate and execute data-quality checks from natural language descriptions.
+
+    Each entry in ``prompts`` describes **one** data-quality expectation.
+    The LLM groups related checks into optimised SQL queries, executes them
+    against the target database, and validates each metric against the
+    corresponding entry in ``validators``.  The task fails if any check
+    does not pass, gating downstream tasks on data quality.
+
+    Generated SQL plans are cached in Airflow
+    :class:`~airflow.models.variable.Variable` to avoid repeat LLM calls.
+    Set ``dry_run=True`` to preview the plan without executing it — the
+    serialised plan dict is returned without running any SQL.
+    Set ``require_approval=True`` to gate execution on human review via the
+    HITL interface: the plan is presented to the reviewer first, and SQL
+    checks run only after approval.  ``dry_run`` and ``require_approval``
+    are independent — enabling both returns the plan dict without any
+    approval prompt.
+
+    :param prompts: Mapping of ``{check_name: natural_language_description}``.
+        Each key must be unique.  Use one check per key; the operator enforces
+        a strict one-key → one-check mapping.
+    :param llm_conn_id: Connection ID for the LLM provider.
+    :param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
+        Overrides the model stored in the connection's extra field.
+    :param system_prompt: Additional instructions appended to the planning prompt.
+    :param agent_params: Additional keyword arguments passed to the pydantic-ai
+        ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param db_conn_id: Connection ID for the database to run checks against.
+        Must resolve to a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
+    :param table_names: Tables to include in the LLM's schema context.
+    :param schema_context: Manual schema description; bypasses DB introspection.
+    :param validators: Mapping of ``{check_name: callable}`` where each callable
+        receives the raw metric value and returns ``True`` (pass) or ``False`` (fail).
+        Keys must be a subset of ``prompts.keys()``.
+        Use built-in factories from
+        :mod:`~airflow.providers.common.ai.utils.dq_validation` or plain lambdas::
+
+            from airflow.providers.common.ai.utils.dq_validation import null_pct_check
+
+            validators = {
+                "email_nulls": null_pct_check(max_pct=0.05),
+                "row_check": lambda v: v >= 1000,
+            }
+
+    :param dialect: SQL dialect override (``postgres``, ``mysql``, etc.).
+        Auto-detected from *db_conn_id* when not set.
+    :param datasource_config: DataFusion datasource for object-storage schema.
+    :param dry_run: When ``True``, generate and cache the plan but skip execution.
+        Returns the serialised plan dict instead of a :class:`~airflow.providers.common.ai.utils.dq_models.DQReport`.
+    :param prompt_version: Optional version tag included in the plan cache key.
+        Bump this to invalidate cached plans when prompts change semantically
+        without changing their text.
+    :param collect_unexpected: When ``True``, the LLM generates an
+        ``unexpected_query`` for validity / string-format checks.
+        If any of those checks fail, the unexpected query is executed and
+        the resulting sample rows are included in the report.
+    :param unexpected_sample_size: Maximum number of violating rows to return
+        per failed check.  Default ``100``.
+    :param row_level_sample_size: Maximum number of rows to fetch per row-level
+        check.  ``None`` (default) performs a full table scan — every row is
+        fetched and validated.  A positive integer is passed to the LLM as a
+        ``LIMIT`` clause on the generated SELECT, bounding execution time and
+        memory usage at the cost of sampling coverage.
+    :param require_approval: When ``True``, the operator defers after generating
+        and caching the DQ plan.  The plan SQL is surfaced in the HITL interface
+        for human review; checks run only after the reviewer approves.  Inherited
+        from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+        ``dry_run=True`` takes precedence — combining both flags returns the plan
+        dict immediately without requesting approval.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "prompts",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+        "prompt_version",
+        "collect_unexpected",
+        "unexpected_sample_size",
+        "row_level_sample_size",
+    )
+
+    def __init__(
+        self,
+        *,
+        prompts: dict[str, str],
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        validators: dict[str, Callable[[Any], bool]] | None = None,
+        dialect: str | None = None,
+        datasource_config: DataSourceConfig | None = None,
+        prompt_version: str | None = None,
+        dry_run: bool = False,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        row_level_sample_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "LLMDataQualityOperator")
+        super().__init__(**kwargs)
+
+        self.prompts = prompts
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.validators = validators or {}
+        self.dialect = dialect
+        self.datasource_config = datasource_config
+        self.prompt_version = prompt_version
+        self.dq_dry_run = dry_run
+        self.collect_unexpected = collect_unexpected
+        self.unexpected_sample_size = unexpected_sample_size
+        self.row_level_sample_size = row_level_sample_size
+
+        self._validate_prompts()
+        self._validate_validator_keys()
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """
+        Generate the DQ plan (or load from cache), then execute or defer for approval.
+
+        When ``dry_run=True`` the serialised plan dict is returned immediately —
+        no SQL is executed and no approval is requested.
+        When ``require_approval=True`` the task defers, presenting the plan to a
+        human reviewer; data-quality checks run only after the reviewer approves.
+
+        :returns: Dict with keys ``plan``, ``passed``, and ``results``.  On success
+            ``passed=True`` and ``results`` is a list of per-check result dicts.
+            For row-level checks the ``value`` entry in each result dict is itself
+            a dict with keys ``total``, ``invalid``, ``invalid_pct``, and
+            ``sample_violations`` rather than a raw scalar.
+            When ``dry_run=True`` ``passed=None`` and ``results=None`` — no SQL
+            is executed.  The ``plan`` key is always present in all modes.
+        :raises DQCheckFailedError: If any data-quality check fails threshold validation.
+        :raises TaskDeferred: When ``require_approval=True``, defers for human review
+            before executing the checks.
+        """
+        planner = self._build_planner()
+
+        schema_ctx = planner.build_schema_context(
+            table_names=self.table_names, schema_context=self.schema_context
+        )
+
+        self.log.info("Using schema context:\n%s", schema_ctx)
+
+        plan = self._load_or_generate_plan(planner, schema_ctx)
+
+        if self.dq_dry_run:
+            self.log.info(
+                "dry_run=True — skipping execution. Plan contains %d group(s), %d check(s).",
+                len(plan.groups),
+                len(plan.check_names),
+            )
+            for group in plan.groups:
+                self.log.info(
+                    "Group: %s\nChecks: %s\nSQL Query:\n%s\n",
+                    group.group_id,
+                    ", ".join(c.check_name for c in group.checks),
+                    group.query,
+                )
+            return {"plan": plan.model_dump(), "passed": None, "results": None}
+
+        if self.require_approval:
+            # Defer BEFORE execution — approval gates the SQL checks.
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                body=self._build_dry_run_markdown(plan),
+            )
+            return {}  # type: ignore[return-value]  # pragma: no cover
+
+        return self._run_checks_and_report(context, planner, plan)
+
+    def _build_planner(self) -> SQLDQPlanner:
+        """Construct a :class:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner` from operator config."""
+        return SQLDQPlanner(
+            llm_hook=self.llm_hook,
+            db_hook=self.db_hook,
+            dialect=self.dialect,
+            datasource_config=self.datasource_config,
+            system_prompt=self.system_prompt,
+            agent_params=self.agent_params,
+            collect_unexpected=self.collect_unexpected,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=self.validator_contexts,
+            row_validators=self._collect_row_validators(),
+            row_level_sample_size=self.row_level_sample_size,
+        )
+
+    @cached_property
+    def validator_contexts(self) -> str:
+        """Return validator-specific LLM context rendered from configured validators."""
+        return default_registry.build_llm_context(self.validators)
+
+    def _run_checks_and_report(
+        self,
+        context: Context,
+        planner: SQLDQPlanner,
+        plan: DQPlan,
+    ) -> dict[str, Any]:
+        """
+        Execute *plan* against the database, apply validators, and return the serialised report.
+
+        :raises DQCheckFailedError: If any data-quality check fails.
+        """
+        results_map = planner.execute_plan(plan)
+        check_results = self._validate_results(results_map, plan)
+
+        # Collect unexpected rows for failed validity/format checks.
+        if self.collect_unexpected:
+            failed_names = {r.check_name for r in check_results if not r.passed}
+            if failed_names:
+                unexpected_map = planner.execute_unexpected_queries(plan, failed_names)
+                self._attach_unexpected(check_results, unexpected_map)
+
+        report = DQReport.build(check_results)
+
+        output: dict[str, Any] = {
+            "plan": plan.model_dump(),
+            "passed": report.passed,
+            "results": [
+                {
+                    "check_name": r.check_name,
+                    "metric_key": r.metric_key,
+                    # RowLevelResult is not JSON-serialisable; convert to a plain dict.
+                    "value": (
+                        {
+                            "total": r.value.total,
+                            "invalid": r.value.invalid,
+                            "invalid_pct": r.value.invalid_pct,
+                            "sample_violations": r.value.sample_violations,
+                            "sample_size": r.value.sample_size,
+                        }
+                        if isinstance(r.value, RowLevelResult)
+                        else r.value
+                    ),
+                    "passed": r.passed,
+                    "failure_reason": r.failure_reason,
+                    **(
+                        {
+                            "unexpected_records": r.unexpected.unexpected_records,
+                            "unexpected_sample_size": r.unexpected.sample_size,
+                        }
+                        if r.unexpected
+                        else {}
+                    ),
+                }
+                for r in report.results
+            ],
+        }
+
+        if not report.passed:
+            # Push results to XCom before failing so downstream tasks
+            # (e.g. with trigger_rule=all_done) can still inspect them.
+            context["ti"].xcom_push(key="return_value", value=output)
+            raise DQCheckFailedError(report.failure_summary)
+
+        self.log.info("All %d data-quality check(s) passed.", len(report.results))
+        return output
+
+    def _build_dry_run_markdown(self, plan: DQPlan) -> str:
+        """
+        Build a structured markdown summary of the DQ plan for the HITL review body.
+
+        Aggregate groups and row-level groups are rendered in separate sections so
+        reviewers can immediately distinguish SQL-aggregate checks from per-row
+        validation logic.
+        """
+        aggregate_groups = [g for g in plan.groups if not any(c.row_level for c in g.checks)]
+        row_level_groups = [g for g in plan.groups if any(c.row_level for c in g.checks)]
+
+        total_checks = len(plan.check_names)
+        agg_count = sum(len(g.checks) for g in aggregate_groups)
+        row_count = sum(len(g.checks) for g in row_level_groups)
+
+        lines: list[str] = [
+            "# LLM Data Quality Plan",
+            "",
+            "| | |",
+            "|---|---|",
+            f"| **Plan hash** | `{plan.plan_hash or 'N/A'}` |",
+            f"| **Total checks** | {total_checks} |",
+            f"| **Aggregate checks** | {agg_count} ({len(aggregate_groups)} group{'s' if len(aggregate_groups) != 1 else ''}) |",
+            f"| **Row-level checks** | {row_count} ({len(row_level_groups)} group{'s' if len(row_level_groups) != 1 else ''}) |",
+            "",
+        ]
+
+        if aggregate_groups:
+            lines += [
+                "---",
+                "",
+                "## Aggregate Checks",
+                "",
+                "> Each group runs as a **single SQL query**. "
+                "Result columns are matched to check names by metric key.",
+                "",
+            ]
+            for group in aggregate_groups:
+                lines += self._render_aggregate_group(group)
+
+        if row_level_groups:
+            lines += [
+                "---",
+                "",
+                "## Row-Level Checks",
+                "",
+                "> Row-level checks fetch **raw column values** and apply Python-side "
+                "validation per row. The threshold controls the maximum allowed fraction "
+                "of invalid rows before the check fails.",
+                "",
+            ]
+            for group in row_level_groups:
+                lines += self._render_row_level_group(group)
+
+        return "\n".join(lines).rstrip()
+
+    def _render_aggregate_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one aggregate SQL group as a markdown subsection."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Category |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            category = check.check_category or "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {category} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+
+        # Unexpected queries — only show when present.
+        unexpected = [(c.check_name, c.unexpected_query) for c in group.checks if c.unexpected_query]
+        if unexpected:
+            lines += ["<details><summary>Unexpected-row queries</summary>", ""]
+            for check_name, uq in unexpected:
+                lines += [
+                    f"**`{check_name}`**",
+                    "",
+                    "```sql",
+                    (uq or "").strip(),
+                    "```",
+                    "",
+                ]
+            lines += ["</details>", ""]
+
+        return lines
+
+    def _render_row_level_group(self, group: DQCheckGroup) -> list[str]:
+        """Render one row-level group as a markdown subsection with threshold info."""
+        lines: list[str] = [
+            f"### `{group.group_id}`",
+            "",
+            "| Check name | Metric key | Max invalid % |",
+            "|---|---|---|",
+        ]
+        for check in group.checks:
+            validator = self.validators.get(check.check_name)
+            max_pct = getattr(validator, "_max_invalid_pct", None)
+            threshold_str = f"{max_pct:.2%}" if max_pct is not None else "—"
+            lines.append(f"| `{check.check_name}` | `{check.metric_key}` | {threshold_str} |")
+
+        lines += [
+            "",
+            "```sql",
+            group.query.strip(),
+            "```",
+            "",
+        ]
+        return lines
+
+    def _load_or_generate_plan(self, planner: SQLDQPlanner, schema_ctx: str) -> DQPlan:
+        """Return a cached plan when available, otherwise generate and cache a new one."""
+        if not isinstance(self.prompts, dict):
+            raise TypeError("prompts must be a dict[str, str] before generating a DQ plan.")
+
+        row_validator_thresholds = self._collect_row_validator_thresholds()
+
+        plan_hash = _compute_plan_hash(
+            self.prompts,
+            self.prompt_version,
+            self.collect_unexpected,
+            self.row_level_sample_size,
+            schema_context=schema_ctx,
+            unexpected_sample_size=self.unexpected_sample_size,
+            validator_contexts=self.validator_contexts,
+            row_validator_thresholds=row_validator_thresholds,
+        )
+        variable_key = f"{_PLAN_VARIABLE_PREFIX}{plan_hash}"
+
+        cached_json = Variable.get(variable_key, None)
+        if cached_json is not None:
+            self.log.info("DQ plan cache hit — key: %r", variable_key)
+            plan = DQPlan.model_validate_json(cached_json)
+            if not plan.plan_hash:
+                plan.plan_hash = plan_hash
+            return plan
+
+        self.log.info("DQ plan cache miss — generating via LLM (key: %r).", variable_key)
+        plan = planner.generate_plan(self.prompts, schema_ctx)
+        plan.plan_hash = plan_hash
+        Variable.set(variable_key, plan.model_dump_json())
+        return plan
+
+    def _validate_results(
+        self,
+        results_map: dict[str, Any],
+        plan: DQPlan,
+    ) -> list[DQCheckResult]:
+        """
+        Apply validators to each metric value and return per-check results.
+
+        For aggregate checks each validator callable receives the raw metric
+        value returned by the database.  For row-level checks, where *value* is
+        a :class:`~airflow.providers.common.ai.utils.dq_models.RowLevelResult`,
+        the pass/fail decision compares ``invalid_pct`` against the validator's
+        ``_max_invalid_pct`` attribute (defaulting to ``0.0`` when absent).
+        Aggregate checks without a registered validator are logged and marked
+        as passed; row-level checks require an explicit validator and fail
+        when none is provided.
+
+        :param results_map: ``{check_name: metric_value_or_RowLevelResult}`` as
+            returned by
+            :meth:`~airflow.providers.common.ai.utils.dq_planner.SQLDQPlanner.execute_plan`.
+        :param plan: The DQ plan whose groups and checks drive iteration order.
+        :returns: Per-check
+            :class:`~airflow.providers.common.ai.utils.dq_models.DQCheckResult`
+            list in plan-group order.
+        :raises ValueError: If *results_map* is missing a key for any check in *plan*.
+        """
+        check_results: list[DQCheckResult] = []
+
+        for group in plan.groups:
+            for check in group.checks:
+                if check.check_name not in results_map:
+                    raise ValueError(
+                        f"Planner did not return a result for check {check.check_name!r} "
+                        f"(group {group.group_id!r}). Available keys: {sorted(results_map)}"
+                    )
+                value = results_map[check.check_name]
+                validator = self.validators.get(check.check_name)
+
+                passed = True
+                failure_reason: str | None = None
+
+                if isinstance(value, RowLevelResult):
+                    if validator is None:
+                        self.log.error(
+                            "No validator found for row-level check %r (metric key: %r). "
+                            "Row-level checks require an explicit validator.",
+                            check.check_name,
+                            check.metric_key,
+                        )
+                        passed = False
+                        failure_reason = (
+                            "Row-level check requires a registered row-level validator, "
+                            "but none was provided."
+                        )
+                    else:
+                        if not hasattr(validator, "_max_invalid_pct"):
+                            self.log.warning(
+                                "Row-level validator for check %r has no '_max_invalid_pct' attribute — "
+                                "defaulting threshold to 0.0%%. Every invalid row will fail the check.",
+                                check.check_name,
+                            )
+                        max_pct = getattr(validator, "_max_invalid_pct", 0.0)
+                        passed = value.invalid_pct <= max_pct
+                        if not passed:
+                            failure_reason = (
+                                f"Row-level check failed: {value.invalid}/{value.total} rows invalid "
+                                f"({value.invalid_pct:.4%}), threshold {max_pct:.4%}"
+                            )

Review Comment:
   Row-level validation assumes `validator._max_invalid_pct` is numeric. If a user passes a non-numeric value (e.g. string via templating), `value.invalid_pct <= max_pct` will raise `TypeError` and fail the task with a low-signal exception. Consider validating/coercing `max_pct` to `float` and raising a clear `ValueError` that names the check when it's missing/invalid.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.

Copilot commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3467353983


##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_data_quality.py:
##########
@@ -0,0 +1,440 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for LLMDataQualityOperator."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQCheckPlan,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+)
+
+_CHECKS = [
+    DQCheckInput(name="null_emails", description="Check for null email addresses"),
+    DQCheckInput(name="dup_ids", description="Check for duplicate customer IDs"),
+]
+
+
+class _TaskInstanceLike:
+    def xcom_push(self, key: str, value: Any) -> None:
+        pass
+
+
+class _UsageLike:
+    requests = 1
+    tool_calls = 0
+    input_tokens = 1
+    output_tokens = 1
+    total_tokens = 2
+
+
+class _ResponseLike:
+    model_name = "test-model"
+
+
+class _AgentResultLike:
+    def __init__(self, output: Any) -> None:
+        self.output = output
+        self.response = _ResponseLike()
+
+    def usage(self) -> _UsageLike:
+        return _UsageLike()
+
+    def all_messages(self) -> list[Any]:
+        return []
+
+
+def _make_context() -> Any:
+    task_instance = MagicMock(spec=_TaskInstanceLike)
+    return {"task_instance": task_instance, "ti": task_instance}
+
+
+def _make_operator(**overrides: Any) -> LLMDataQualityOperator:
+    defaults: dict[str, Any] = {
+        "task_id": "test_dq",
+        "checks": _CHECKS,
+        "llm_conn_id": "pydantic_ai_default",
+        "db_conn_id": "postgres_default",
+    }
+    defaults.update(overrides)
+    op = LLMDataQualityOperator(**defaults)
+    op.llm_hook = MagicMock(spec=PydanticAIHook)
+    return op
+
+
+def _passing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=True),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=True,
+    )
+
+
+def _failing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=False, failure_reason="100 nulls found"),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=False,
+        failure_summary="null_emails: 100 nulls found",
+    )
+
+
+def _mock_agent_result(output: Any) -> _AgentResultLike:
+    return _AgentResultLike(output)
+
+
+class TestLLMDataQualityOperatorInit:
+    def test_requires_llm_conn_id(self):
+        with pytest.raises(TypeError):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                db_conn_id="postgres_default",
+            )
+
+    def test_raises_when_no_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+            )
+
+    def test_raises_when_empty_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+                toolsets=[],
+            )
+
+    def test_empty_checks_raises_value_error(self):
+        with pytest.raises(ValueError, match="checks must not be empty"):
+            _make_operator(checks=[])
+
+    def test_duplicate_check_names_raises(self):
+        with pytest.raises(ValueError, match="duplicate"):
+            _make_operator(
+                checks=[
+                    DQCheckInput(name="dup", description="first"),
+                    DQCheckInput(name="dup", description="second"),
+                ]
+            )
+
+    def test_dict_checks_auto_coerced(self):
+        op = _make_operator(
+            checks=[
+                {"name": "null_emails", "description": "Check nulls"},
+                {"name": "dup_ids", "description": "Check dups"},
+            ]
+        )
+        assert all(isinstance(c, DQCheckInput) for c in op.checks)
+
+    def test_template_fields_include_required_keys(self):
+        op = _make_operator()
+        tf = set(op.template_fields)
+        assert {"checks", "system_prompt", "agent_params", "db_conn_id", "table_names"} <= tf
+
+
+class TestResolveToolsets:
+    def test_explicit_toolsets_returned_unchanged(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        resolved = op._resolve_toolsets()
+        assert resolved is op.toolsets

Review Comment:
   This test uses `MagicMock(spec=BaseDQToolset)`, but `LLMDataQualityOperator._resolve_toolsets()` detects DQ toolsets via `isinstance(ts, BaseDQToolset)`, so a MagicMock will not be recognized. Also `_resolve_toolsets()` returns a new list (`list(self.toolsets)`), so asserting identity (`is`) will fail even with a real toolset instance.



##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/dataquality/sql.py:
##########
@@ -0,0 +1,399 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""SQL data-quality toolset for LLMDataQualityOperator."""
+
+from __future__ import annotations
+
+import inspect
+import json
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal
+
+try:
+    from airflow.providers.common.ai.utils.dataquality.validation import ValidatorRegistry, default_registry
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.toolsets.abstract import ToolsetTool
+
+from airflow.providers.common.ai.toolsets.dataquality.base import _PASSTHROUGH_VALIDATOR, BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import RowLevelResult
+
+if TYPE_CHECKING:
+    from pydantic_ai._run_context import RunContext
+
+_log = logging.getLogger(__name__)
+
+# JSON Schemas for the three SQL-DQ-specific tools.
+_LIST_VALIDATORS_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "properties": {},
+}
+
+_APPLY_VALIDATOR_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "properties": {
+        "check_name": {"type": "string", "description": "Name of the DQ check being evaluated."},
+        "value": {"description": "Metric value returned by the SQL query for this check."},
+        "validator_name": {
+            "type": "string",
+            "description": "Registered validator name (from list_validators). Pass 'none' to skip.",
+        },
+        "validator_args": {
+            "type": "object",
+            "description": 'Keyword arguments for the validator factory (e.g. {"max_pct": 0.05}).',
+        },
+    },
+    "required": ["check_name", "value", "validator_name", "validator_args"],
+}
+
+_ROW_LEVEL_SAMPLE_LIMIT = 20
+
+
+class SQLDQToolset(BaseDQToolset):
+    """
+    Data-quality toolset for SQL-based checks.
+
+    Provides three tools on top of
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`:
+
+    * ``list_validators`` — exposes the registered validator catalog so the LLM
+      can choose an appropriate validator for each check.
+    * ``apply_validator`` — instantiates a validator and applies it to a metric
+      value, returning ``{"passed": bool, "reason": ...}``.
+
+    Use this toolset alongside :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`
+    or :class:`~airflow.providers.common.ai.toolsets.datafusion.DataFusionToolset` in
+    :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`.
+    The data-source toolset handles schema discovery and query execution; this
+    toolset handles the DQ-specific logic::
+
+        LLMDataQualityOperator(
+            checks=[...],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default"),
+                SQLDQToolset(),
+            ],
+        )
+
+    :param validator_registry: Validator registry to expose to the LLM.
+        Defaults to :data:`~airflow.providers.common.ai.utils.dataquality.validation.default_registry`.
+    """
+
+    def __init__(self, *, validator_registry: ValidatorRegistry | None = None) -> None:
+        super().__init__()
+        self._registry = validator_registry if validator_registry is not None else default_registry
+
+    @property
+    def id(self) -> str:
+        return "dq-sql"
+
+    @property
+    def output_mode(self) -> Literal["execute", "generate"]:
+        return "execute"
+
+    # ------------------------------------------------------------------
+    # AbstractToolset interface
+    # ------------------------------------------------------------------
+
+    async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
+        tools = await super().get_tools(ctx)
+
+        for name, description, schema in (
+            (
+                "list_validators",
+                "List available validator names, parameters, and descriptions.",
+                _LIST_VALIDATORS_SCHEMA,
+            ),
+            (
+                "apply_validator",
+                "Apply a registered validator to a metric value and return pass/fail.",
+                _APPLY_VALIDATOR_SCHEMA,
+            ),
+        ):
+            # In planning mode (two-phase approval flow) the agent must NOT call
+            # apply_validator — it only selects a validator by name/args and records
+            # that choice in the DQPlan output.  apply_validator runs in Phase 2
+            # (pure Python) after the human reviewer approves.
+            if name == "apply_validator" and getattr(self, "_planning_mode", False):
+                continue
+            tool_def = ToolDefinition(
+                name=name,
+                description=description,
+                parameters_json_schema=schema,
+                sequential=True,
+            )
+            tools[name] = ToolsetTool(
+                toolset=self,
+                tool_def=tool_def,
+                max_retries=1,
+                args_validator=_PASSTHROUGH_VALIDATOR,
+            )
+        return tools
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[Any],
+        tool: ToolsetTool[Any],
+    ) -> Any:
+        if name == "list_validators":
+            result = self._list_validators()
+            _log.info("list_validators result: %s", result)
+            return result
+        if name == "apply_validator":
+            check_name = tool_args["check_name"]
+            validator_name = tool_args["validator_name"]
+            validator_args = tool_args.get("validator_args") or {}
+            value = tool_args["value"]
+
+            # Resolve the display name: for fixed validators, show the actual
+            # validator's name/signature followed by "(fixed)" so logs are
+            # meaningful instead of just showing "fixed".
+            if validator_name.lower() == "fixed" or self._has_fixed_validator(check_name):
+                fixed = self._get_fixed_validator(check_name)
+                display_name = getattr(fixed, "_validator_display", None) or getattr(
+                    fixed, "_validator_name", "fixed"
+                )
+                log_validator = f"{display_name}(fixed)"
+            else:
+                log_validator = validator_name
+
+            _log.info(
+                "apply_validator: check=%s, validator=%s, args=%s, value=%s",
+                check_name,
+                log_validator,
+                validator_args,
+                value,
+            )
+            result = self._apply_validator(check_name, value, validator_name, validator_args)
+            parsed = json.loads(result)
+            _log.info(
+                "apply_validator result: check=%s, passed=%s, reason=%s",
+                parsed.get("check_name"),
+                parsed.get("passed"),
+                parsed.get("reason"),
+            )
+            return result
+        return await super().call_tool(name, tool_args, ctx, tool)
+
+    # ------------------------------------------------------------------
+    # Tool implementations
+    # ------------------------------------------------------------------
+
+    def _list_validators(self) -> str:
+        entries = []
+        for name in self._registry.list_validators():
+            entry = self._registry.get(name)
+            entries.append(
+                {
+                    "name": name,
+                    "category": entry.check_category or None,
+                    "row_level": entry.row_level,
+                    "parameters": self._format_signature(entry.factory),
+                    "description": entry.llm_context or None,
+                }
+            )
+        return json.dumps(entries)
+
+    def _apply_validator(
+        self,
+        check_name: str,
+        value: Any,
+        validator_name: str,
+        validator_args: dict[str, Any],
+    ) -> str:
+        if validator_name.lower() == "fixed" or self._has_fixed_validator(check_name):
+            fixed = self._get_fixed_validator(check_name)
+            if fixed is None:
+                return json.dumps(
+                    {
+                        "check_name": check_name,
+                        "passed": False,
+                        "reason": f"No fixed validator configured for check {check_name!r}.",
+                    }
+                )
+            is_row_level = bool(getattr(fixed, "_row_level", False))
+            if is_row_level:
+                return self._apply_row_level_validator(check_name, value, fixed, validator_name="fixed")
+
+            _log.info("apply_validator scalar: check=%s, value=%s", check_name, value)
+            try:
+                passed = bool(fixed(value))
+            except Exception as exc:
+                return json.dumps({"check_name": check_name, "passed": False, "reason": str(exc)})
+            reason = None if passed else f"Fixed validator returned False for value {value!r}"
+            return json.dumps({"check_name": check_name, "passed": passed, "reason": reason})
+
+        if not validator_name or validator_name.lower() == "none":
+            return json.dumps({"check_name": check_name, "passed": True, "reason": "no validator"})
+
+        ok, err = self._validate_suggestion(validator_name, validator_args)
+        if not ok:
+            return json.dumps({"check_name": check_name, "passed": False, "reason": err})
+
+        try:
+            entry = self._registry.get(validator_name)
+            validator = entry.factory(**validator_args)
+        except Exception as exc:
+            return json.dumps({"check_name": check_name, "passed": False, "reason": str(exc)})
+
+        if entry.row_level:
+            return self._apply_row_level_validator(
+                check_name,
+                value,
+                validator,
+                validator_name=validator_name,
+            )
+
+        try:
+            passed = bool(validator(value))
+        except Exception as exc:
+            return json.dumps({"check_name": check_name, "passed": False, "reason": str(exc)})
+
+        reason = None if passed else f"{validator_name} returned False for value {value!r}"
+        return json.dumps({"check_name": check_name, "passed": passed, "reason": reason})
+
+    def _apply_row_level_validator(
+        self,
+        check_name: str,
+        value: Any,
+        validator: Callable[[Any], bool],
+        *,
+        validator_name: str,
+    ) -> str:
+        """Apply a row-level validator to each row value and return a RowLevelResult payload."""
+        values = self._coerce_row_values(check_name, value)
+        invalid_values: list[Any] = []
+
+        for row_value in values:
+            try:
+                is_valid = bool(validator(row_value))
+            except Exception as exc:
+                return json.dumps(
+                    {
+                        "check_name": check_name,
+                        "passed": False,
+                        "reason": (
+                            f"{validator_name} raised {exc!s} while evaluating row-level value {row_value!r}."
+                        ),
+                    }
+                )
+            if not is_valid:
+                invalid_values.append(row_value)
+
+        summary = RowLevelResult.build(invalid_values, len(values), sample_limit=_ROW_LEVEL_SAMPLE_LIMIT)
+        max_invalid_pct = self._get_max_invalid_pct(validator)

Review Comment:
   Row-level validation accumulates every invalid row value into `invalid_values` before sampling, which can become a large in-memory list when `values` is big (e.g. if `max_rows` is raised or a non-SQL toolset returns many rows). It’s enough to track an invalid count plus a bounded sample list, and build `RowLevelResult` from those fields directly.



##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_data_quality.py:
##########
@@ -0,0 +1,440 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for LLMDataQualityOperator."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQCheckPlan,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+)
+
+_CHECKS = [
+    DQCheckInput(name="null_emails", description="Check for null email addresses"),
+    DQCheckInput(name="dup_ids", description="Check for duplicate customer IDs"),
+]
+
+
+class _TaskInstanceLike:
+    def xcom_push(self, key: str, value: Any) -> None:
+        pass
+
+
+class _UsageLike:
+    requests = 1
+    tool_calls = 0
+    input_tokens = 1
+    output_tokens = 1
+    total_tokens = 2
+
+
+class _ResponseLike:
+    model_name = "test-model"
+
+
+class _AgentResultLike:
+    def __init__(self, output: Any) -> None:
+        self.output = output
+        self.response = _ResponseLike()
+
+    def usage(self) -> _UsageLike:
+        return _UsageLike()
+
+    def all_messages(self) -> list[Any]:
+        return []
+
+
+def _make_context() -> Any:
+    task_instance = MagicMock(spec=_TaskInstanceLike)
+    return {"task_instance": task_instance, "ti": task_instance}
+
+
+def _make_operator(**overrides: Any) -> LLMDataQualityOperator:
+    defaults: dict[str, Any] = {
+        "task_id": "test_dq",
+        "checks": _CHECKS,
+        "llm_conn_id": "pydantic_ai_default",
+        "db_conn_id": "postgres_default",
+    }
+    defaults.update(overrides)
+    op = LLMDataQualityOperator(**defaults)
+    op.llm_hook = MagicMock(spec=PydanticAIHook)
+    return op
+
+
+def _passing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=True),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=True,
+    )
+
+
+def _failing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=False, failure_reason="100 nulls found"),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=False,
+        failure_summary="null_emails: 100 nulls found",
+    )
+
+
+def _mock_agent_result(output: Any) -> _AgentResultLike:
+    return _AgentResultLike(output)
+
+
+class TestLLMDataQualityOperatorInit:
+    def test_requires_llm_conn_id(self):
+        with pytest.raises(TypeError):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                db_conn_id="postgres_default",
+            )
+
+    def test_raises_when_no_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+            )
+
+    def test_raises_when_empty_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+                toolsets=[],
+            )
+
+    def test_empty_checks_raises_value_error(self):
+        with pytest.raises(ValueError, match="checks must not be empty"):
+            _make_operator(checks=[])
+
+    def test_duplicate_check_names_raises(self):
+        with pytest.raises(ValueError, match="duplicate"):
+            _make_operator(
+                checks=[
+                    DQCheckInput(name="dup", description="first"),
+                    DQCheckInput(name="dup", description="second"),
+                ]
+            )
+
+    def test_dict_checks_auto_coerced(self):
+        op = _make_operator(
+            checks=[
+                {"name": "null_emails", "description": "Check nulls"},
+                {"name": "dup_ids", "description": "Check dups"},
+            ]
+        )
+        assert all(isinstance(c, DQCheckInput) for c in op.checks)
+
+    def test_template_fields_include_required_keys(self):
+        op = _make_operator()
+        tf = set(op.template_fields)
+        assert {"checks", "system_prompt", "agent_params", "db_conn_id", "table_names"} <= tf
+
+
+class TestResolveToolsets:
+    def test_explicit_toolsets_returned_unchanged(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        resolved = op._resolve_toolsets()
+        assert resolved is op.toolsets
+
+    def test_auto_creates_sql_toolsets_from_db_conn_id(self):
+        op = _make_operator()
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_dq,
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(db_conn_id="postgres_default", allowed_tables=None)
+            mock_dq.assert_called_once_with()
+
+    def test_auto_creates_sql_toolsets_with_table_names(self):
+        op = _make_operator(table_names=["customers", "orders"])
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset"),
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(
+                db_conn_id="postgres_default", allowed_tables=["customers", "orders"]
+            )
+
+    def test_auto_appends_sql_dq_toolset_when_missing(self):
+        mock_other = object()
+        op = _make_operator(toolsets=[mock_other])
+        with patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_cls:
+            mock_cls.return_value = MagicMock(spec=BaseDQToolset)
+            resolved = op._resolve_toolsets()
+            mock_cls.assert_called_once_with()
+            assert len(resolved) == 2
+
+
+class TestFindDqToolset:
+    def test_returns_first_base_dq_toolset(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_other = object()
+        result = LLMDataQualityOperator._find_dq_toolset([mock_other, mock_dq])
+        assert result is mock_dq
+
+    def test_raises_when_no_dq_toolset(self):
+        mock_other = object()
+        with pytest.raises(ValueError, match="No BaseDQToolset found"):
+            LLMDataQualityOperator._find_dq_toolset([mock_other])
+
+
+class TestExecuteMode:
+    def _run(self, report: DQReport, **overrides: Any) -> Any:
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq], **overrides)
+        op.llm_hook.create_agent.return_value.run_sync.return_value = _mock_agent_result(report)  # type: ignore[attr-defined]
+        return op.execute(context=_make_context())
+
+    def test_set_checks_called_on_dq_toolset(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        op.llm_hook.create_agent.return_value.run_sync.return_value = _mock_agent_result(_passing_report())
+        op.execute(context=_make_context())
+        mock_dq.set_checks.assert_called_once_with(op.checks)
+
+    def test_passing_report_returns_model_dump(self):
+        result = self._run(_passing_report())
+        assert result["passed"] is True
+        assert isinstance(result["results"], list)
+
+    def test_failing_report_raises_dq_check_failed_error(self):
+        with pytest.raises(DQCheckFailedError, match="null_emails"):
+            self._run(_failing_report())
+
+    def test_output_type_is_dq_report(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        op.llm_hook.create_agent.return_value.run_sync.return_value = _mock_agent_result(_passing_report())
+        op.execute(context=_make_context())
+        create_agent_kwargs = op.llm_hook.create_agent.call_args.kwargs
+        assert create_agent_kwargs["output_type"] is DQReport
+
+
+class TestGenerateMode:
+    def _run_generate(self, config_str: str = "checks:\n  - name: foo") -> Any:
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "generate"
+        op = _make_operator(toolsets=[mock_dq])
+        op.llm_hook.create_agent.return_value.run_sync.return_value = _mock_agent_result(config_str)  # type: ignore[attr-defined]
+        return op.execute(context=_make_context())

Review Comment:
   Like the execute-mode tests, using `MagicMock(spec=BaseDQToolset)` here won’t satisfy `isinstance(..., BaseDQToolset)` checks in the operator, so the generate-mode path may not be exercised as intended. Use a minimal concrete `BaseDQToolset` implementation with `output_mode="generate"`.



##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_data_quality.py:
##########
@@ -0,0 +1,440 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for LLMDataQualityOperator."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQCheckPlan,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+)
+
+_CHECKS = [
+    DQCheckInput(name="null_emails", description="Check for null email addresses"),
+    DQCheckInput(name="dup_ids", description="Check for duplicate customer IDs"),
+]
+
+
+class _TaskInstanceLike:
+    def xcom_push(self, key: str, value: Any) -> None:
+        pass
+
+
+class _UsageLike:
+    requests = 1
+    tool_calls = 0
+    input_tokens = 1
+    output_tokens = 1
+    total_tokens = 2
+
+
+class _ResponseLike:
+    model_name = "test-model"
+
+
+class _AgentResultLike:
+    def __init__(self, output: Any) -> None:
+        self.output = output
+        self.response = _ResponseLike()
+
+    def usage(self) -> _UsageLike:
+        return _UsageLike()
+
+    def all_messages(self) -> list[Any]:
+        return []
+
+
+def _make_context() -> Any:
+    task_instance = MagicMock(spec=_TaskInstanceLike)
+    return {"task_instance": task_instance, "ti": task_instance}
+
+
+def _make_operator(**overrides: Any) -> LLMDataQualityOperator:
+    defaults: dict[str, Any] = {
+        "task_id": "test_dq",
+        "checks": _CHECKS,
+        "llm_conn_id": "pydantic_ai_default",
+        "db_conn_id": "postgres_default",
+    }
+    defaults.update(overrides)
+    op = LLMDataQualityOperator(**defaults)
+    op.llm_hook = MagicMock(spec=PydanticAIHook)
+    return op
+
+
+def _passing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=True),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=True,
+    )
+
+
+def _failing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=False, failure_reason="100 nulls found"),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=False,
+        failure_summary="null_emails: 100 nulls found",
+    )
+
+
+def _mock_agent_result(output: Any) -> _AgentResultLike:
+    return _AgentResultLike(output)
+
+
+class TestLLMDataQualityOperatorInit:
+    def test_requires_llm_conn_id(self):
+        with pytest.raises(TypeError):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                db_conn_id="postgres_default",
+            )
+
+    def test_raises_when_no_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+            )
+
+    def test_raises_when_empty_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+                toolsets=[],
+            )
+
+    def test_empty_checks_raises_value_error(self):
+        with pytest.raises(ValueError, match="checks must not be empty"):
+            _make_operator(checks=[])
+
+    def test_duplicate_check_names_raises(self):
+        with pytest.raises(ValueError, match="duplicate"):
+            _make_operator(
+                checks=[
+                    DQCheckInput(name="dup", description="first"),
+                    DQCheckInput(name="dup", description="second"),
+                ]
+            )
+
+    def test_dict_checks_auto_coerced(self):
+        op = _make_operator(
+            checks=[
+                {"name": "null_emails", "description": "Check nulls"},
+                {"name": "dup_ids", "description": "Check dups"},
+            ]
+        )
+        assert all(isinstance(c, DQCheckInput) for c in op.checks)
+
+    def test_template_fields_include_required_keys(self):
+        op = _make_operator()
+        tf = set(op.template_fields)
+        assert {"checks", "system_prompt", "agent_params", "db_conn_id", "table_names"} <= tf
+
+
+class TestResolveToolsets:
+    def test_explicit_toolsets_returned_unchanged(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        resolved = op._resolve_toolsets()
+        assert resolved is op.toolsets
+
+    def test_auto_creates_sql_toolsets_from_db_conn_id(self):
+        op = _make_operator()
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_dq,
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(db_conn_id="postgres_default", allowed_tables=None)
+            mock_dq.assert_called_once_with()
+
+    def test_auto_creates_sql_toolsets_with_table_names(self):
+        op = _make_operator(table_names=["customers", "orders"])
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset"),
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(
+                db_conn_id="postgres_default", allowed_tables=["customers", "orders"]
+            )
+
+    def test_auto_appends_sql_dq_toolset_when_missing(self):
+        mock_other = object()
+        op = _make_operator(toolsets=[mock_other])
+        with patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_cls:
+            mock_cls.return_value = MagicMock(spec=BaseDQToolset)
+            resolved = op._resolve_toolsets()
+            mock_cls.assert_called_once_with()
+            assert len(resolved) == 2
+
+
+class TestFindDqToolset:
+    def test_returns_first_base_dq_toolset(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_other = object()
+        result = LLMDataQualityOperator._find_dq_toolset([mock_other, mock_dq])
+        assert result is mock_dq
+
+    def test_raises_when_no_dq_toolset(self):
+        mock_other = object()
+        with pytest.raises(ValueError, match="No BaseDQToolset found"):
+            LLMDataQualityOperator._find_dq_toolset([mock_other])
+
+
+class TestExecuteMode:
+    def _run(self, report: DQReport, **overrides: Any) -> Any:
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq], **overrides)
+        op.llm_hook.create_agent.return_value.run_sync.return_value = _mock_agent_result(report)  # type: ignore[attr-defined]
+        return op.execute(context=_make_context())

Review Comment:
   Here and in several other tests in this file, `MagicMock(spec=BaseDQToolset)` is passed as a toolset, but production code relies on `isinstance(..., BaseDQToolset)` for discovery. That means the operator will treat the toolsets list as missing a DQ toolset and may append a real `SQLDQToolset` or raise, so these tests won’t exercise the intended path.



##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -213,6 +229,20 @@ def _is_ref_allowed(self, catalog: str | None, schema: str | None, table: str) -
     def id(self) -> str:
         return f"sql-{self._db_conn_id}"
 
+    @property
+    def sqlglot_dialect(self) -> str | None:
+        """Sqlglot dialect name for this connection, used for SQL validation and LLM prompting."""
+        if self._explicit_dialect is not None:

Review Comment:
   `sqlglot_dialect` docstring says it’s used for “SQL validation and LLM prompting”, but the toolset’s validation path currently uses `_dialect_for_validation()` (hook.dialect_name) and doesn’t consult this property. Either wire this into the validation path, or adjust the docstring to avoid misleading API consumers.



##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_data_quality.py:
##########
@@ -0,0 +1,440 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for LLMDataQualityOperator."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQCheckPlan,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+)
+
+_CHECKS = [
+    DQCheckInput(name="null_emails", description="Check for null email addresses"),
+    DQCheckInput(name="dup_ids", description="Check for duplicate customer IDs"),
+]
+
+
+class _TaskInstanceLike:
+    def xcom_push(self, key: str, value: Any) -> None:
+        pass
+
+
+class _UsageLike:
+    requests = 1
+    tool_calls = 0
+    input_tokens = 1
+    output_tokens = 1
+    total_tokens = 2
+
+
+class _ResponseLike:
+    model_name = "test-model"
+
+
+class _AgentResultLike:
+    def __init__(self, output: Any) -> None:
+        self.output = output
+        self.response = _ResponseLike()
+
+    def usage(self) -> _UsageLike:
+        return _UsageLike()
+
+    def all_messages(self) -> list[Any]:
+        return []
+
+
+def _make_context() -> Any:
+    task_instance = MagicMock(spec=_TaskInstanceLike)
+    return {"task_instance": task_instance, "ti": task_instance}
+
+
+def _make_operator(**overrides: Any) -> LLMDataQualityOperator:
+    defaults: dict[str, Any] = {
+        "task_id": "test_dq",
+        "checks": _CHECKS,
+        "llm_conn_id": "pydantic_ai_default",
+        "db_conn_id": "postgres_default",
+    }
+    defaults.update(overrides)
+    op = LLMDataQualityOperator(**defaults)
+    op.llm_hook = MagicMock(spec=PydanticAIHook)
+    return op
+
+
+def _passing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=True),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=True,
+    )
+
+
+def _failing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=False, failure_reason="100 nulls found"),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=False,
+        failure_summary="null_emails: 100 nulls found",
+    )
+
+
+def _mock_agent_result(output: Any) -> _AgentResultLike:
+    return _AgentResultLike(output)
+
+
+class TestLLMDataQualityOperatorInit:
+    def test_requires_llm_conn_id(self):
+        with pytest.raises(TypeError):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                db_conn_id="postgres_default",
+            )
+
+    def test_raises_when_no_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+            )
+
+    def test_raises_when_empty_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+                toolsets=[],
+            )
+
+    def test_empty_checks_raises_value_error(self):
+        with pytest.raises(ValueError, match="checks must not be empty"):
+            _make_operator(checks=[])
+
+    def test_duplicate_check_names_raises(self):
+        with pytest.raises(ValueError, match="duplicate"):
+            _make_operator(
+                checks=[
+                    DQCheckInput(name="dup", description="first"),
+                    DQCheckInput(name="dup", description="second"),
+                ]
+            )
+
+    def test_dict_checks_auto_coerced(self):
+        op = _make_operator(
+            checks=[
+                {"name": "null_emails", "description": "Check nulls"},
+                {"name": "dup_ids", "description": "Check dups"},
+            ]
+        )
+        assert all(isinstance(c, DQCheckInput) for c in op.checks)
+
+    def test_template_fields_include_required_keys(self):
+        op = _make_operator()
+        tf = set(op.template_fields)
+        assert {"checks", "system_prompt", "agent_params", "db_conn_id", "table_names"} <= tf
+
+
+class TestResolveToolsets:
+    def test_explicit_toolsets_returned_unchanged(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        resolved = op._resolve_toolsets()
+        assert resolved is op.toolsets
+
+    def test_auto_creates_sql_toolsets_from_db_conn_id(self):
+        op = _make_operator()
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_dq,
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(db_conn_id="postgres_default", allowed_tables=None)
+            mock_dq.assert_called_once_with()
+
+    def test_auto_creates_sql_toolsets_with_table_names(self):
+        op = _make_operator(table_names=["customers", "orders"])
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset"),
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(
+                db_conn_id="postgres_default", allowed_tables=["customers", "orders"]
+            )
+
+    def test_auto_appends_sql_dq_toolset_when_missing(self):
+        mock_other = object()
+        op = _make_operator(toolsets=[mock_other])
+        with patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_cls:
+            mock_cls.return_value = MagicMock(spec=BaseDQToolset)
+            resolved = op._resolve_toolsets()
+            mock_cls.assert_called_once_with()
+            assert len(resolved) == 2
+
+
+class TestFindDqToolset:
+    def test_returns_first_base_dq_toolset(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_other = object()
+        result = LLMDataQualityOperator._find_dq_toolset([mock_other, mock_dq])
+        assert result is mock_dq
+
+    def test_raises_when_no_dq_toolset(self):
+        mock_other = object()
+        with pytest.raises(ValueError, match="No BaseDQToolset found"):
+            LLMDataQualityOperator._find_dq_toolset([mock_other])
+
+
+class TestExecuteMode:
+    def _run(self, report: DQReport, **overrides: Any) -> Any:
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq], **overrides)
+        op.llm_hook.create_agent.return_value.run_sync.return_value = _mock_agent_result(report)  # type: ignore[attr-defined]
+        return op.execute(context=_make_context())
+
+    def test_set_checks_called_on_dq_toolset(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        op.llm_hook.create_agent.return_value.run_sync.return_value = _mock_agent_result(_passing_report())
+        op.execute(context=_make_context())
+        mock_dq.set_checks.assert_called_once_with(op.checks)
+

Review Comment:
   `set_checks` call assertion relies on a `MagicMock(spec=BaseDQToolset)` toolset, but the operator discovers DQ toolsets via `isinstance`, so this mock won’t be used as the DQ toolset. Use a real `BaseDQToolset` subclass instance and wrap `set_checks` with a spy mock if you need call assertions.



##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_data_quality.py:
##########
@@ -0,0 +1,440 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for LLMDataQualityOperator."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.operators.llm_data_quality import LLMDataQualityOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQCheckPlan,
+    DQCheckResult,
+    DQPlan,
+    DQReport,
+)
+
+_CHECKS = [
+    DQCheckInput(name="null_emails", description="Check for null email addresses"),
+    DQCheckInput(name="dup_ids", description="Check for duplicate customer IDs"),
+]
+
+
+class _TaskInstanceLike:
+    def xcom_push(self, key: str, value: Any) -> None:
+        pass
+
+
+class _UsageLike:
+    requests = 1
+    tool_calls = 0
+    input_tokens = 1
+    output_tokens = 1
+    total_tokens = 2
+
+
+class _ResponseLike:
+    model_name = "test-model"
+
+
+class _AgentResultLike:
+    def __init__(self, output: Any) -> None:
+        self.output = output
+        self.response = _ResponseLike()
+
+    def usage(self) -> _UsageLike:
+        return _UsageLike()
+
+    def all_messages(self) -> list[Any]:
+        return []
+
+
+def _make_context() -> Any:
+    task_instance = MagicMock(spec=_TaskInstanceLike)
+    return {"task_instance": task_instance, "ti": task_instance}
+
+
+def _make_operator(**overrides: Any) -> LLMDataQualityOperator:
+    defaults: dict[str, Any] = {
+        "task_id": "test_dq",
+        "checks": _CHECKS,
+        "llm_conn_id": "pydantic_ai_default",
+        "db_conn_id": "postgres_default",
+    }
+    defaults.update(overrides)
+    op = LLMDataQualityOperator(**defaults)
+    op.llm_hook = MagicMock(spec=PydanticAIHook)
+    return op
+
+
+def _passing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=True),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=True,
+    )
+
+
+def _failing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=False, failure_reason="100 nulls found"),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=False,
+        failure_summary="null_emails: 100 nulls found",
+    )
+
+
+def _mock_agent_result(output: Any) -> _AgentResultLike:
+    return _AgentResultLike(output)
+
+
+class TestLLMDataQualityOperatorInit:
+    def test_requires_llm_conn_id(self):
+        with pytest.raises(TypeError):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                db_conn_id="postgres_default",
+            )
+
+    def test_raises_when_no_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+            )
+
+    def test_raises_when_empty_toolsets_and_no_db_conn_id(self):
+        with pytest.raises(ValueError, match="Either toolsets or db_conn_id"):
+            LLMDataQualityOperator(
+                task_id="test_dq",
+                checks=_CHECKS,
+                llm_conn_id="pydantic_ai_default",
+                toolsets=[],
+            )
+
+    def test_empty_checks_raises_value_error(self):
+        with pytest.raises(ValueError, match="checks must not be empty"):
+            _make_operator(checks=[])
+
+    def test_duplicate_check_names_raises(self):
+        with pytest.raises(ValueError, match="duplicate"):
+            _make_operator(
+                checks=[
+                    DQCheckInput(name="dup", description="first"),
+                    DQCheckInput(name="dup", description="second"),
+                ]
+            )
+
+    def test_dict_checks_auto_coerced(self):
+        op = _make_operator(
+            checks=[
+                {"name": "null_emails", "description": "Check nulls"},
+                {"name": "dup_ids", "description": "Check dups"},
+            ]
+        )
+        assert all(isinstance(c, DQCheckInput) for c in op.checks)
+
+    def test_template_fields_include_required_keys(self):
+        op = _make_operator()
+        tf = set(op.template_fields)
+        assert {"checks", "system_prompt", "agent_params", "db_conn_id", "table_names"} <= tf
+
+
+class TestResolveToolsets:
+    def test_explicit_toolsets_returned_unchanged(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_dq.output_mode = "execute"
+        op = _make_operator(toolsets=[mock_dq])
+        resolved = op._resolve_toolsets()
+        assert resolved is op.toolsets
+
+    def test_auto_creates_sql_toolsets_from_db_conn_id(self):
+        op = _make_operator()
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_dq,
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(db_conn_id="postgres_default", allowed_tables=None)
+            mock_dq.assert_called_once_with()
+
+    def test_auto_creates_sql_toolsets_with_table_names(self):
+        op = _make_operator(table_names=["customers", "orders"])
+        with (
+            patch("airflow.providers.common.ai.toolsets.sql.SQLToolset") as mock_sql,
+            patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset"),
+        ):
+            op._resolve_toolsets()
+            mock_sql.assert_called_once_with(
+                db_conn_id="postgres_default", allowed_tables=["customers", "orders"]
+            )
+
+    def test_auto_appends_sql_dq_toolset_when_missing(self):
+        mock_other = object()
+        op = _make_operator(toolsets=[mock_other])
+        with patch("airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset") as mock_cls:
+            mock_cls.return_value = MagicMock(spec=BaseDQToolset)
+            resolved = op._resolve_toolsets()
+            mock_cls.assert_called_once_with()
+            assert len(resolved) == 2
+
+
+class TestFindDqToolset:
+    def test_returns_first_base_dq_toolset(self):
+        mock_dq = MagicMock(spec=BaseDQToolset)
+        mock_other = object()
+        result = LLMDataQualityOperator._find_dq_toolset([mock_other, mock_dq])
+        assert result is mock_dq

Review Comment:
   `_find_dq_toolset()` uses `isinstance(toolset, BaseDQToolset)`, so `MagicMock(spec=BaseDQToolset)` won’t match and this test will fail. Use a minimal concrete `BaseDQToolset` subclass instance instead of a mock for the DQ toolset slot.



##########
providers/common/ai/tests/unit/common/ai/decorators/test_llm_data_quality.py:
##########
@@ -0,0 +1,141 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.providers.common.ai.decorators.llm_data_quality import (
+    _LLMDQDecoratedOperator,
+)
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckInput,
+    DQCheckResult,
+    DQReport,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+_CHECKS = [
+    DQCheckInput(name="null_emails", description="Check for null email addresses"),
+    DQCheckInput(name="dup_ids", description="Check for duplicate customer IDs"),
+]
+
+
+def _passing_report() -> DQReport:
+    return DQReport(
+        results=[
+            DQCheckResult(check_name="null_emails", passed=True),
+            DQCheckResult(check_name="dup_ids", passed=True),
+        ],
+        passed=True,
+    )
+
+
+def _make_op(callable_fn=None, **kwargs) -> _LLMDQDecoratedOperator:
+    if callable_fn is None:
+
+        def callable_fn():
+            return _CHECKS
+
+    return _LLMDQDecoratedOperator(
+        task_id="test_dq",
+        python_callable=callable_fn,
+        llm_conn_id="pydantic_ai_default",
+        db_conn_id="postgres_default",
+        **kwargs,
+    )
+
+
+def _mock_agent_result(output) -> MagicMock:
+    result = MagicMock()
+    result.output = output
+    result.usage.return_value = MagicMock(
+        requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0
+    )
+    result.response = MagicMock(model_name="test-model")
+    result.all_messages.return_value = []
+    return result

Review Comment:
   This helper constructs several `MagicMock()` objects without a `spec`/`autospec`, which can hide interface drift (e.g. `usage` vs `usage()` changes) and make refactors harder to catch. Prefer `MagicMock(spec=[...])` for the agent result and response objects, consistent with other tests in this provider.



##########
providers/common/ai/docs/operators/llm_data_quality.rst:
##########
@@ -0,0 +1,502 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _howto/operator:llm_data_quality:
+
+``LLMDataQualityOperator``
+==========================
+
+Use :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator`
+to run natural-language data-quality checks against a live data source using an LLM agent.
+
+The agent:
+
+1. Calls ``list_checks`` to read the user's quality expectations.
+2. Calls ``list_validators`` to discover available validators.
+3. Uses schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. Writes and executes SQL queries, applies validators, and produces a
+   :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.
+5. Fails the task when any check does not pass.
+
+.. seealso::
+    :ref:`Connection configuration <howto/connection:pydanticai>`
+
+Architecture — Toolsets
+-----------------------
+
+The operator is composed of two toolsets that are passed together in ``toolsets``:
+
+- **Data-source toolset** — gives the agent access to the data (SQL queries, schema discovery).
+  Examples: :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` for relational
+  databases, :class:`~airflow.providers.common.ai.toolsets.datafusion.DataFusionToolset`
+  for object-storage formats (Parquet, CSV, Avro on S3, GCS, etc.).
+
+- **DQ toolset** (:class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset`) —
+  adds the ``list_validators`` and ``apply_validator`` tools so the agent can evaluate metrics
+  against registered validator thresholds.
+
+When ``toolsets`` is omitted, the operator auto-creates
+:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+:class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset`
+from ``db_conn_id`` and ``table_names``.
+
+PostgreSQL / Relational Databases
+----------------------------------
+
+Pass ``SQLToolset`` (schema-discovery + SQL execution) alongside ``SQLDQToolset``:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_postgres_basic]
+    :end-before: [END howto_operator_llm_dq_postgres_basic]
+
+Object Storage (S3 / Parquet)
+------------------------------
+
+Use :class:`~airflow.providers.common.ai.toolsets.datafusion.DataFusionToolset` to register
+S3 (or other object-store) Parquet/CSV data as a queryable DataFusion table, then pair it with
+``SQLDQToolset``:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_object_storage]
+    :end-before: [END howto_operator_llm_dq_object_storage]
+
+.. note::
+    The ``region`` key in the ``aws_default`` connection **Extra** field must match the
+    bucket's AWS region to avoid redirect errors, e.g. ``{"region": "eu-central-1"}``.
+
+Explicit Toolsets
+-----------------
+
+Passing ``toolsets`` directly gives full control over which tables the agent can access.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_explicit_toolsets]
+    :end-before: [END howto_operator_llm_dq_explicit_toolsets]
+
+Checks
+------
+
+Each entry in ``checks`` is a
+:class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+(or a plain ``dict`` with ``name``, ``description``, and optional ``validator`` keys).
+
+.. code-block:: python
+
+    from airflow.providers.common.ai.utils.dataquality import DQCheckInput, null_pct_check
+
+    checks = [
+        DQCheckInput(
+            name="null_email",
+            description="Check the percentage of rows where email is NULL",
+            validator=null_pct_check(max_pct=0.05),
+        ),
+        DQCheckInput(
+            name="row_count",
+            description="Ensure the table has at least 1000 rows",
+        ),
+    ]
+
+Check names must be unique within a single operator call.
+
+Validators
+----------
+
+Built-in validator factories are exported from
+:mod:`~airflow.providers.common.ai.utils.dataquality`:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Factory
+     - Description
+   * - ``null_pct_check(max_pct=...)``
+     - Passes when the fraction of NULL values is at or below ``max_pct``.
+   * - ``row_count_check(min_count=...)``
+     - Passes when the row count is at least ``min_count``.
+   * - ``duplicate_pct_check(max_pct=...)``
+     - Passes when the fraction of duplicate values is at or below ``max_pct``.
+   * - ``between_check(min_val=..., max_val=...)``
+     - Passes when the metric value falls within ``[min_val, max_val]``.
+   * - ``exact_check(expected=...)``
+     - Passes when the metric value equals ``expected`` exactly.
+
+Checks without a validator are measured and included in the report.
+If no validator is applied, the check is treated as passed.
+
+Custom Validators
+-----------------
+
+Use :func:`~airflow.providers.common.ai.utils.dataquality.register_validator`
+to register validator factories that the LLM can discover via ``list_validators``
+and select dynamically:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_email_format_validator]
+    :end-before: [END howto_operator_llm_dq_email_format_validator]
+
+Row-Level Checks
+----------------
+
+Row-level validators receive a list of values (one per row) instead of a single
+aggregate metric.  The agent issues a plain ``SELECT <column> FROM <table>`` query
+(no aggregation) and passes the column values to ``apply_validator``.
+
+The validator returns a
+:class:`~airflow.providers.common.ai.utils.dataquality.models.RowLevelResult` with
+these fields:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 25 75
+
+   * - Field
+     - Description
+   * - ``total``
+     - Number of evaluated rows.
+   * - ``invalid``
+     - Number of rows that failed the validator predicate.
+   * - ``invalid_pct``
+     - ``invalid / total`` (``0.0`` when ``total`` is zero).
+   * - ``sample_violations``
+     - Sample of failing values (string form).
+   * - ``sample_size``
+     - Length of ``sample_violations``.
+
+Register a row-level validator with ``row_level=True``:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_email_format_validator]
+    :end-before: [END howto_operator_llm_dq_email_format_validator]
+
+Schema Context
+--------------
+
+Pass ``schema_context`` to inject a manual schema description into the system prompt.
+Useful when the data source cannot be introspected at runtime, or to restrict
+the agent to a specific set of tables:
+
+.. code-block:: python
+
+    LLMDataQualityOperator(
+        task_id="validate_with_schema",
+        llm_conn_id="pydanticai_default",
+        db_conn_id="postgres_default",
+        schema_context=(
+            "Table: orders\n"
+            "Columns: id INT, customer_id INT, amount DECIMAL, created_at TIMESTAMP\n\n"
+            "Table: customers\n"
+            "Columns: id INT, email TEXT, country TEXT"
+        ),
+        checks=[
+            DQCheckInput(
+                name="null_amount",
+                description="Check the percentage of orders with a NULL amount",
+                validator=null_pct_check(max_pct=0.0),
+            ),
+        ],
+    )
+
+Human-in-the-Loop Approval
+---------------------------
+
+Set ``require_approval=True`` to gate execution on a human reviewer.
+The task runs in two phases:
+
+1. **Phase 1** — the LLM discovers schema, writes SQL queries, and selects validators
+   but does *not* execute any SQL or call ``apply_validator``.  The resulting plan
+   (SQL statements + validator choices) is surfaced in the Airflow UI.
+2. **Phase 2** (after approval) — the approved SQL queries run, validators are applied
+   in pure Python, and the final :class:`~...DQReport` is produced.  No LLM calls.
+
+Rejecting the plan raises
+:class:`~airflow.providers.standard.exceptions.HITLRejectException`.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_require_approval]
+    :end-before: [END howto_operator_llm_dq_require_approval]
+
+See :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` for shared HITL parameters
+(``approval_timeout``, ``allow_modifications``).
+
+Durable Execution
+-----------------
+
+Set ``durable=True`` to enable cross-run caching of LLM and tool-call results.
+When a task is retried or restarted, the operator replays cached model responses and
+tool calls instead of re-running them, saving tokens and wall-clock time.
+
+.. code-block:: python
+
+    LLMDataQualityOperator(
+        task_id="validate_products",
+        llm_conn_id="pydanticai_default",
+        toolsets=[
+            SQLToolset(db_conn_id="postgres_default", allowed_tables=["products"]),
+            SQLDQToolset(),
+        ],
+        durable=True,
+        checks=[...],
+    )
+
+TaskFlow Decorator
+------------------
+
+Use ``@task.llm_dq`` when checks are produced dynamically by a Python callable.
+The function body returns the checks list; the decorator handles the LLM agent run,
+SQL execution, and validator application automatically.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_data_quality.py
+    :language: python
+    :start-after: [START howto_operator_llm_dq_task_decorator]
+    :end-before: [END howto_operator_llm_dq_task_decorator]
+
+Output and XCom
+---------------
+
+The operator returns a ``dict`` representation of
+:class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport` as its XCom value
+(``output_mode="execute"``).
+
+Additionally, the full report is pushed to XCom under the key ``dq_report`` on the
+task instance:
+
+.. code-block:: python
+
+    report = context["ti"].xcom_pull(task_ids="validate_orders", key="dq_report")
+
+For config-generation backends (``output_mode="generate"``), the operator returns
+the generated configuration string instead of a report.
+
+Logging
+-------
+
+After each LLM call, the operator logs a summary with model name, token usage,
+and request count at INFO level:
+
+.. code-block:: text
+
+    [INFO] Agent run summary — model: gpt-4o, requests: 6, tool_calls: 14,
+           input_tokens: 3201, output_tokens: 412, total_tokens: 3613
+
+Parameters Reference
+--------------------
+
+.. list-table::
+   :header-rows: 1
+   :widths: 25 15 60
+
+   * - Parameter
+     - Type
+     - Description
+   * - ``checks``
+     - ``list``
+     - List of :class:`~...DQCheckInput` objects or plain dicts.  Names must be unique.
+   * - ``toolsets``
+     - ``list | None``
+     - Pydantic-AI toolsets.  Must include a data-source toolset and a
+       :class:`~...BaseDQToolset`.  Auto-created from ``db_conn_id`` when ``None``.
+   * - ``db_conn_id``
+     - ``str | None``
+     - Connection ID for auto-creating ``SQLToolset``.  Ignored when ``toolsets`` is set.
+   * - ``table_names``
+     - ``list[str] | None``
+     - Tables exposed to the auto-created ``SQLToolset``.  Ignored when ``toolsets`` is set.
+   * - ``schema_context``
+     - ``str | None``
+     - Manual schema description injected into the system prompt.
+   * - ``durable``
+     - ``bool``
+     - Enable cross-run LLM and tool-call caching.  Default ``False``.
+   * - ``require_approval``
+     - ``bool``
+     - Enable two-phase HITL approval before SQL execution.  Default ``False``.
+   * - ``llm_conn_id``
+     - ``str``
+     - Pydantic-AI connection ID (inherited from
+       :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`).
+
+
+Each entry in ``checks`` is a
+:class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`.
+The operator:
+
+1. Exposes checks and validators as tools to the model.
+2. Lets the model discover schema and generate SQL via your data-source toolset.
+3. Applies validators and produces a
+   :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.
+4. Fails the task when any check fails.
+
+.. seealso::
+    :ref:`Connection configuration <howto/connection:pydanticai>`
+
+Basic Usage
+-----------
+
+If you pass ``db_conn_id`` (without explicit toolsets), the operator auto-creates:
+
+- :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`
+- :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset`
+

Review Comment:
   This doc page appears to contain a duplicated second pass of major sections (e.g. “Basic Usage”, “Explicit Toolsets”, “Validators”, etc.) starting around here, repeating content already covered earlier in the page. It would be better to remove the duplicated block and keep a single coherent flow to avoid conflicting updates later.



##########
providers/common/ai/provider.yaml:
##########
@@ -457,6 +459,8 @@ task-decorators:
     name: llm_file_analysis
   - class-name: airflow.providers.common.ai.decorators.llm_branch.llm_branch_task
     name: llm_branch
+  - class-name: airflow.providers.common.ai.decorators.llm_data_quality.llm_dq_task
+    name: llm_dq

Review Comment:
   PR description mentions an `@task.llm_data_quality` decorator, but the implementation/registration uses `llm_dq` (e.g. task-decorator name `llm_dq` and `custom_operator_name = "@task.llm_dq"`). Please align naming across docs/description/API (either rename the decorator to `llm_data_quality` or update the PR description to match `llm_dq`).



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,770 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for running data-quality checks from natural language using an LLM agent."""
+
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQPlan,
+    DQReport,
+)
+from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
+    from airflow.sdk import Context
+
+_DQ_SYSTEM_PROMPT = """\
+You are a data-quality expert. Evaluate the user's data-quality checks against a \
+live data source using the tools available to you.
+
+WORKFLOW:
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL CHECKS
+      procedure below instead of writing an aggregate query.
+   b. Otherwise, write a SELECT query that computes the relevant metric for this check.
+   c. Optionally call ``check_query`` to validate SQL syntax before executing.
+   d. Execute the query using the ``query`` tool.
+   e. You MUST call ``apply_validator`` for EVERY check — never skip this step, even for
+      fixed validators (``has_fixed_validator: true``).  Use ``validator_name: "fixed"`` for those.
+5. Return a ``DQReport`` with one ``DQCheckResult`` per check — every check must appear exactly once.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+  Example: -- check: null_order_id
+- For conditional counts, use CASE expressions (FILTER WHERE is not universally supported):
+    CORRECT:   COUNT(CASE WHEN col IS NULL THEN 1 END)
+    INCORRECT: COUNT(*) FILTER (WHERE col IS NULL)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- For duplicate percentages:
+    (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / NULLIF(COUNT(*), 0)
+- For float division, cast to avoid integer truncation:
+    CAST(numerator AS DOUBLE) / CAST(denominator AS DOUBLE)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+VALIDATOR SELECTION:
+- ``list_validators`` returns each validator's name, parameter signature, and description.
+  Read the parameter names and types before calling ``apply_validator``.
+- Pass the required ``validator_args`` as a JSON object matching the parameter signature exactly.
+- If a parameter has no default, it is REQUIRED — always include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — you must pass ``{"max_pct": <value>}``.
+  If you are unsure of a threshold, use a safe default (e.g. ``0.05`` for percentage checks).
+- NEVER pass empty ``{}`` args for a validator that has required parameters.
+- For checks marked ``has_fixed_validator: true`` in ``list_checks``, call ``apply_validator``
+  with ``validator_name: "fixed"`` — the pre-assigned validator runs automatically, no args needed.
+- If no suitable validator exists, pass ``validator_name: "none"`` — the check passes by default
+  but its metric value is still recorded.
+
+DQREPORT OUTPUT RULES:
+- ``check_name`` must exactly match the name returned by ``list_checks`` (no abbreviation).
+- Each check must appear in the report exactly once.
+- Set ``passed: false`` and provide a clear ``failure_reason`` for every failed check.
+- Set ``metric_key`` to the SQL column alias you used for this check's metric value.
+- Set ``sql_query`` to the exact SQL statement you executed for this check (including the leading comment).
+- Set ``validator_info`` to ``{"name": "<validator_name>", "args": {<validator_args>}}`` when a validator
+  was applied, or ``null`` when no validator was used (``validator_name: "none"``).
+
+CHECK CATEGORIES (use to guide SQL and validator selection):
+  null_check    — null / missing value counts or percentages
+  uniqueness    — duplicate detection, cardinality checks
+  validity      — regex / format / pattern matching on string columns
+  numeric_range — range, bounds, or statistical checks on numeric columns
+  row_count     — total row counts or existence checks
+  string_format — length, encoding, whitespace, or character-set checks
+  row_level     — per-row anomaly checks; validator receives a list, not an aggregate
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true`` for a check):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>`` that
+  returns one value per row — no GROUP BY, no COUNT, no CASE expressions.
+  Example: -- check: customer_email_format\n  SELECT email FROM customers
+- After executing, extract the column values into a Python list:
+    value = [row["<column>"] for row in result["rows"]]
+  where ``result`` is the JSON object returned by the ``query`` tool.
+- Call ``apply_validator`` with ``value`` set to that list and ``validator_name: "fixed"``.
+- The validator evaluates each item and returns a ``value`` payload containing a
+    ``RowLevelResult`` summary (total/invalid/invalid_pct/sample_violations/sample_size).
+- Set ``metric_key`` to the SQL column name (e.g. ``"email"``).
+- Set ``value`` in the DQCheckResult to the ``value`` object returned by ``apply_validator``.
+"""
+
+_DQ_PLAN_SYSTEM_PROMPT = """\
+You are a data-quality expert.  Your task is to plan *how* to evaluate each
+data-quality check: write the SQL query and choose the validator for each check.
+A human reviewer will inspect your plan before any SQL is executed.
+
+WORKFLOW (PLANNING MODE — no SQL execution, no apply_validator):
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to understand the data model.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL
+      procedure below.  Otherwise, write an aggregate SELECT query for the metric.
+   b. Optionally call ``check_query`` to validate your SQL syntax — but do NOT execute it.
+   c. Select the appropriate validator name and arguments:
+      - For ``has_fixed_validator: true`` checks: use ``validator_name: "fixed"``, empty args.
+      - If no validator fits: use ``validator_name: "none"``.
+5. Return a ``DQPlan`` with one ``DQCheckPlan`` per check — every check must appear exactly once.
+
+IMPORTANT: The ``query`` tool is NOT available.  Do NOT try to execute SQL.
+Do NOT call ``apply_validator`` — it is not available in planning mode.
+SQL will be executed after the human reviewer approves the plan.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+- For conditional counts, use CASE expressions:
+    COUNT(CASE WHEN col IS NULL THEN 1 END)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true``):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>``.
+- Set ``row_level: true`` in the DQCheckPlan output.
+- Set ``metric_key`` to the column name (e.g. ``"email"``).
+
+DQCheckPlan FIELDS:
+- ``check_name``: exact name from ``list_checks`` (no abbreviation).
+- ``sql_query``: the SQL statement to execute later (with leading ``-- check: <name>`` comment).
+- ``metric_key``: the SQL column alias / name to read as the primary metric value.
+- ``row_level``: ``true`` for row-level checks (SELECT returns one row per record), ``false`` for aggregate.
+- ``validator_name``: ``"fixed"``, a registered validator name, or ``"none"``.
+- ``validator_args``: kwargs for the validator factory, or ``{}`` for ``"fixed"``/``"none"``.
+
+VALIDATOR ARGS RULES (critical — missing args cause hard failures in Phase 2):
+- ``list_validators`` returns a ``parameters`` field for each validator showing its exact
+  parameter names, types, and defaults.  Read this carefully before setting ``validator_args``.
+- If a parameter has no default (required), you MUST include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — always set ``{"max_pct": <value>}``.
+- If you are unsure of the right threshold, pick a reasonable default (e.g. 0.05 for pct checks).
+- NEVER leave ``validator_args`` as ``{}`` for a validator that has required parameters.
+- Only use ``{}`` when ``validator_name`` is ``"fixed"`` or ``"none"``.
+"""
+
+_DIALECT_SQL_NOTES: dict[str, str] = {
+    "postgres": (
+        "  - Regex match: `col ~ 'pattern'`; not match: `col !~ 'pattern'` (case-sensitive).\n"
+        "    Case-insensitive variants: `~*` and `!~*`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)` (no CAST needed).\n"
+        "  - String functions: LENGTH(), LOWER(), TRIM(), SUBSTRING().\n"
+    ),
+    "mysql": (
+        "  - Regex match: `col REGEXP 'pattern'`; not match: `col NOT REGEXP 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - String length (multibyte-safe): CHAR_LENGTH(col); byte length: LENGTH(col).\n"
+    ),
+    "sqlite": (
+        "  - No native regex operator; use LIKE or GLOB:\n"
+        "    `col NOT LIKE '%@%'` or `col GLOB '*@*.*'`.\n"
+        "  - Float division: `CAST(numerator AS REAL) / NULLIF(denominator, 0)`.\n"
+    ),
+    "tsql": (
+        "  - No regex operator; use LIKE or PATINDEX:\n"
+        "    `PATINDEX('%@%.%', col) = 0` means no match.\n"
+        "  - Float division: `CAST(numerator AS FLOAT) / NULLIF(denominator, 0)`.\n"
+        "  - Use TOP n instead of LIMIT n; no OFFSET without ORDER BY.\n"
+        "  - String length: LEN(col) (not LENGTH).\n"
+    ),
+    "bigquery": (
+        "  - Regex: `REGEXP_CONTAINS(col, r'pattern')` — returns TRUE if match.\n"
+        "    Negate with `NOT REGEXP_CONTAINS(...)`.\n"
+        "  - Safe division: `SAFE_DIVIDE(numerator, denominator)` (NULL on zero denominator).\n"
+        "  - Quote reserved names with backticks.\n"
+    ),
+    "snowflake": (
+        "  - Regex: `REGEXP_LIKE(col, 'pattern')` (full-string match) or `col RLIKE 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - Unquoted identifiers are uppercased; use double-quotes when case matters.\n"
+    ),
+}
+
+
+def _extract_schema_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by build_schema_context."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Run data-quality checks described in natural language using an LLM agent.
+
+    The agent discovers the database schema, writes and executes SQL queries,
+    applies validators, and produces a
+    :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.  The task
+    fails when any check does not pass, gating downstream tasks on data quality.
+
+    Supply the data-source toolset and DQ toolset together in ``toolsets``::
+
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+        LLMDataQualityOperator(
+            task_id="quality_check",
+            checks=[
+                DQCheckInput(name="email_nulls", description="Check for null emails"),
+                DQCheckInput(
+                    name="row_count",
+                    description="At least 1000 rows",
+                    validator=row_count_check(min_count=1000),
+                ),
+            ],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers"]),
+                SQLDQToolset(),
+            ],
+        )
+
+    When ``toolsets`` is omitted, the operator auto-creates
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` from
+    ``db_conn_id`` and ``table_names``.
+
+    For config-generation backends (e.g. ``SodaDQToolset``), the operator
+    returns the generated config string as its XCom value instead of a report.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional
+        ``validator`` keys).  Names must be unique.
+    :param toolsets: Pydantic-AI toolsets for the agent.  Must include exactly
+        one :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`
+        subclass alongside a data-source toolset.  When ``None``, the operator
+        auto-creates toolsets from ``db_conn_id``.
+    :param db_conn_id: Connection ID for auto-creating an
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`.
+        Ignored when ``toolsets`` is provided.
+    :param table_names: Tables passed to the auto-created
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` as
+        ``allowed_tables``.  Ignored when ``toolsets`` is provided.
+    :param schema_context: Additional schema description injected into the
+        system prompt.  Useful when the data source cannot be introspected
+        at runtime.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        toolsets: list[AbstractToolset] | None = None,
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Pop operator-specific params that must not reach BaseOperator.__init__.
+        # Using kwargs.pop (rather than named params) is necessary because Airflow's
+        # apply_defaults metaclass captures the full **kwargs dict and uses it for
+        # task-map serialization; named params in the dict are NOT removed from the
+        # snapshot, so they would be re-injected on deserialization and reach
+        # BaseOperator as unknown kwargs.
+        durable: bool = kwargs.pop("durable", False)
+
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "Run the data-quality checks.")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            [DQCheckInput.coerce(c) for c in checks] if isinstance(checks, list) else checks  # type: ignore[assignment]
+        )
+        self.toolsets = toolsets
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.durable = durable
+
+        self._validate_checks()
+
+        if not toolsets and db_conn_id is None:
+            raise ValueError("Either toolsets or db_conn_id must be provided.")
+
+    def execute(self, context: Context) -> Any:
+        """
+        Run the LLM agent to execute all data-quality checks.
+
+        When ``require_approval=True``, the task runs in two phases:
+
+        1. **Phase 1** — the LLM discovers schema and selects validators *without
+           executing any SQL queries*.  The resulting plan (check names,
+           descriptions, and validator assignments) is shown to a human reviewer.
+        2. **Phase 2** (in :meth:`execute_complete`) — after approval, SQL queries
+           are executed and validators are applied in pure Python to produce the
+           :class:`~...DQReport`.
+
+        :returns: :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`
+            as a dict when the DQ toolset is in ``"execute"`` mode, or a config
+            string when in ``"generate"`` mode.
+        :raises DQCheckFailedError: When any check fails in ``"execute"`` mode.
+        """
+        if self.require_approval:
+            plan, _ = self._run_plan_phase(context)
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                subject=f"Review DQ plan for task `{self.task_id}`",
+                body=self._build_plan_approval_body(plan),
+            )
+        return self._run_and_report(context)
+
+    def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any:
+        """
+        Phase 2: execute SQL and apply validators after human approval.
+
+        Called automatically by Airflow when the HITL trigger fires.  The base
+        class validates the approval decision (raises on reject or timeout).
+        Then each SQL query from the approved :class:`~...DQPlan` is executed in
+        pure Python, the metric values are extracted, and validators are applied
+        — no LLM calls.
+
+        :param context: Airflow task context.
+        :param generated_output: The :class:`~...DQPlan` JSON that was deferred.
+        :param event: Trigger event payload.
+        :returns: Same as :meth:`execute`.
+        :raises HITLRejectException: If the reviewer rejected the plan.
+        """
+        approved_output = super().execute_complete(context, generated_output, event)  # type: ignore[misc]
+        plan = DQPlan.model_validate_json(approved_output)
+
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        report = self._execute_plan_validators(plan, dq_toolset, toolsets)
+        context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+        if not report.passed:
+            raise DQCheckFailedError(report.failure_summary)
+        return report.model_dump()
+
+    # ------------------------------------------------------------------
+    # Core execution helpers
+    # ------------------------------------------------------------------
+
+    def _run_plan_phase(self, context: Context) -> tuple[DQPlan, BaseDQToolset]:
+        """
+        Phase 1 of the two-phase approval flow.
+
+        Runs the LLM agent in planning mode: schema-discovery tools
+        (``list_tables``, ``get_schema``, ``check_query``) are available, but
+        the ``query`` and ``apply_validator`` tools are hidden.  The agent
+        writes SQL strings and selects validators for each check, then outputs
+        a :class:`~...DQPlan` — without executing any SQL.
+
+        :returns: ``(plan, dq_toolset)`` — the plan for the approval body and the
+            toolset (already configured with checks) for Phase 2.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset._planning_mode = True  # omit apply_validator from tool list
+        dq_toolset.set_checks(self.checks)
+
+        # Hide the SQL 'query' tool so the LLM cannot execute DQ queries.
+        for ts in toolsets:
+            if ts is not dq_toolset and hasattr(ts, "_query"):
+                ts._planning_mode = True  # type: ignore[attr-defined]
+
+        instructions = self._build_plan_system_prompt(toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+        storage = None
+
+        if self.durable:
+            agent, counter, storage = self._build_durable_agent(
+                context,
+                DQPlan,
+                instructions,
+                logged_toolsets,
+            )
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=DQPlan,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay (plan phase): model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        if storage is not None:
+            storage.cleanup()
+
+        return result.output, dq_toolset
+
+    def _run_and_report(self, context: Context) -> Any:
+        """
+        Resolve toolsets, run the LLM agent, push XCom, and handle failures.
+
+        This is the single place where the agent executes regardless of whether
+        the task took the direct path or the approval-gated path.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        output_type: type = DQReport if dq_toolset.output_mode == "execute" else str
+        instructions = self._build_system_prompt(dq_toolset, toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)

Review Comment:
   When `dq_toolset.output_mode == "generate"`, the operator switches `output_type` to `str` but still uses the same execution-oriented system prompt (`_DQ_SYSTEM_PROMPT`) that instructs the agent to run `query`/`apply_validator` and return a `DQReport`. This prompt/output mismatch will likely cause models to return the wrong shape or attempt execution steps that a generate-mode toolset may not support.



##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_data_quality.py:
##########
@@ -0,0 +1,770 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Operator for running data-quality checks from natural language using an LLM agent."""
+
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.toolsets.dataquality.base import BaseDQToolset
+from airflow.providers.common.ai.utils.dataquality.models import (
+    DQCheckFailedError,
+    DQCheckInput,
+    DQPlan,
+    DQReport,
+)
+from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
+    from airflow.sdk import Context
+
+_DQ_SYSTEM_PROMPT = """\
+You are a data-quality expert. Evaluate the user's data-quality checks against a \
+live data source using the tools available to you.
+
+WORKFLOW:
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to explore the data source.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL CHECKS
+      procedure below instead of writing an aggregate query.
+   b. Otherwise, write a SELECT query that computes the relevant metric for this check.
+   c. Optionally call ``check_query`` to validate SQL syntax before executing.
+   d. Execute the query using the ``query`` tool.
+   e. You MUST call ``apply_validator`` for EVERY check — never skip this step, even for
+      fixed validators (``has_fixed_validator: true``).  Use ``validator_name: "fixed"`` for those.
+5. Return a ``DQReport`` with one ``DQCheckResult`` per check — every check must appear exactly once.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+  Example: -- check: null_order_id
+- For conditional counts, use CASE expressions (FILTER WHERE is not universally supported):
+    CORRECT:   COUNT(CASE WHEN col IS NULL THEN 1 END)
+    INCORRECT: COUNT(*) FILTER (WHERE col IS NULL)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- For duplicate percentages:
+    (COUNT(*) - COUNT(DISTINCT col)) * 1.0 / NULLIF(COUNT(*), 0)
+- For float division, cast to avoid integer truncation:
+    CAST(numerator AS DOUBLE) / CAST(denominator AS DOUBLE)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+VALIDATOR SELECTION:
+- ``list_validators`` returns each validator's name, parameter signature, and description.
+  Read the parameter names and types before calling ``apply_validator``.
+- Pass the required ``validator_args`` as a JSON object matching the parameter signature exactly.
+- If a parameter has no default, it is REQUIRED — always include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — you must pass ``{"max_pct": <value>}``.
+  If you are unsure of a threshold, use a safe default (e.g. ``0.05`` for percentage checks).
+- NEVER pass empty ``{}`` args for a validator that has required parameters.
+- For checks marked ``has_fixed_validator: true`` in ``list_checks``, call ``apply_validator``
+  with ``validator_name: "fixed"`` — the pre-assigned validator runs automatically, no args needed.
+- If no suitable validator exists, pass ``validator_name: "none"`` — the check passes by default
+  but its metric value is still recorded.
+
+DQREPORT OUTPUT RULES:
+- ``check_name`` must exactly match the name returned by ``list_checks`` (no abbreviation).
+- Each check must appear in the report exactly once.
+- Set ``passed: false`` and provide a clear ``failure_reason`` for every failed check.
+- Set ``metric_key`` to the SQL column alias you used for this check's metric value.
+- Set ``sql_query`` to the exact SQL statement you executed for this check (including the leading comment).
+- Set ``validator_info`` to ``{"name": "<validator_name>", "args": {<validator_args>}}`` when a validator
+  was applied, or ``null`` when no validator was used (``validator_name: "none"``).
+
+CHECK CATEGORIES (use to guide SQL and validator selection):
+  null_check    — null / missing value counts or percentages
+  uniqueness    — duplicate detection, cardinality checks
+  validity      — regex / format / pattern matching on string columns
+  numeric_range — range, bounds, or statistical checks on numeric columns
+  row_count     — total row counts or existence checks
+  string_format — length, encoding, whitespace, or character-set checks
+  row_level     — per-row anomaly checks; validator receives a list, not an aggregate
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true`` for a check):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>`` that
+  returns one value per row — no GROUP BY, no COUNT, no CASE expressions.
+  Example: -- check: customer_email_format\n  SELECT email FROM customers
+- After executing, extract the column values into a Python list:
+    value = [row["<column>"] for row in result["rows"]]
+  where ``result`` is the JSON object returned by the ``query`` tool.
+- Call ``apply_validator`` with ``value`` set to that list and ``validator_name: "fixed"``.
+- The validator evaluates each item and returns a ``value`` payload containing a
+    ``RowLevelResult`` summary (total/invalid/invalid_pct/sample_violations/sample_size).
+- Set ``metric_key`` to the SQL column name (e.g. ``"email"``).
+- Set ``value`` in the DQCheckResult to the ``value`` object returned by ``apply_validator``.
+"""
+
+_DQ_PLAN_SYSTEM_PROMPT = """\
+You are a data-quality expert.  Your task is to plan *how* to evaluate each
+data-quality check: write the SQL query and choose the validator for each check.
+A human reviewer will inspect your plan before any SQL is executed.
+
+WORKFLOW (PLANNING MODE — no SQL execution, no apply_validator):
+1. Call ``list_checks`` to read the user's quality expectations.
+2. Call ``list_validators`` to see available validators (names, parameters, descriptions).
+3. Use schema-discovery tools (``list_tables``, ``get_schema``) to understand the data model.
+4. For each check:
+   a. If the check has ``row_level: true`` (from ``list_checks``), follow the ROW-LEVEL
+      procedure below.  Otherwise, write an aggregate SELECT query for the metric.
+   b. Optionally call ``check_query`` to validate your SQL syntax — but do NOT execute it.
+   c. Select the appropriate validator name and arguments:
+      - For ``has_fixed_validator: true`` checks: use ``validator_name: "fixed"``, empty args.
+      - If no validator fits: use ``validator_name: "none"``.
+5. Return a ``DQPlan`` with one ``DQCheckPlan`` per check — every check must appear exactly once.
+
+IMPORTANT: The ``query`` tool is NOT available.  Do NOT try to execute SQL.
+Do NOT call ``apply_validator`` — it is not available in planning mode.
+SQL will be executed after the human reviewer approves the plan.
+
+SQL GENERATION RULES:
+- Generate ONLY SELECT statements.  NEVER use INSERT, UPDATE, DELETE, DROP, TRUNCATE, or DDL.
+- Begin each query with a SQL comment that names the check it serves:
+    -- check: <check_name>
+- For conditional counts, use CASE expressions:
+    COUNT(CASE WHEN col IS NULL THEN 1 END)
+- For null/invalid percentages:
+    COUNT(CASE WHEN condition THEN 1 END) * 1.0 / NULLIF(COUNT(*), 0)
+- Give each metric column a descriptive snake_case alias (e.g. ``null_email_pct``).
+
+ROW-LEVEL CHECKS (applies when ``list_checks`` returns ``row_level: true``):
+- Do NOT write an aggregate query.  Write a plain ``SELECT <column> FROM <table>``.
+- Set ``row_level: true`` in the DQCheckPlan output.
+- Set ``metric_key`` to the column name (e.g. ``"email"``).
+
+DQCheckPlan FIELDS:
+- ``check_name``: exact name from ``list_checks`` (no abbreviation).
+- ``sql_query``: the SQL statement to execute later (with leading ``-- check: <name>`` comment).
+- ``metric_key``: the SQL column alias / name to read as the primary metric value.
+- ``row_level``: ``true`` for row-level checks (SELECT returns one row per record), ``false`` for aggregate.
+- ``validator_name``: ``"fixed"``, a registered validator name, or ``"none"``.
+- ``validator_args``: kwargs for the validator factory, or ``{}`` for ``"fixed"``/``"none"``.
+
+VALIDATOR ARGS RULES (critical — missing args cause hard failures in Phase 2):
+- ``list_validators`` returns a ``parameters`` field for each validator showing its exact
+  parameter names, types, and defaults.  Read this carefully before setting ``validator_args``.
+- If a parameter has no default (required), you MUST include it in ``validator_args``.
+  Example: ``null_pct_check`` requires ``max_pct`` — always set ``{"max_pct": <value>}``.
+- If you are unsure of the right threshold, pick a reasonable default (e.g. 0.05 for pct checks).
+- NEVER leave ``validator_args`` as ``{}`` for a validator that has required parameters.
+- Only use ``{}`` when ``validator_name`` is ``"fixed"`` or ``"none"``.
+"""
+
+_DIALECT_SQL_NOTES: dict[str, str] = {
+    "postgres": (
+        "  - Regex match: `col ~ 'pattern'`; not match: `col !~ 'pattern'` (case-sensitive).\n"
+        "    Case-insensitive variants: `~*` and `!~*`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)` (no CAST needed).\n"
+        "  - String functions: LENGTH(), LOWER(), TRIM(), SUBSTRING().\n"
+    ),
+    "mysql": (
+        "  - Regex match: `col REGEXP 'pattern'`; not match: `col NOT REGEXP 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - String length (multibyte-safe): CHAR_LENGTH(col); byte length: LENGTH(col).\n"
+    ),
+    "sqlite": (
+        "  - No native regex operator; use LIKE or GLOB:\n"
+        "    `col NOT LIKE '%@%'` or `col GLOB '*@*.*'`.\n"
+        "  - Float division: `CAST(numerator AS REAL) / NULLIF(denominator, 0)`.\n"
+    ),
+    "tsql": (
+        "  - No regex operator; use LIKE or PATINDEX:\n"
+        "    `PATINDEX('%@%.%', col) = 0` means no match.\n"
+        "  - Float division: `CAST(numerator AS FLOAT) / NULLIF(denominator, 0)`.\n"
+        "  - Use TOP n instead of LIMIT n; no OFFSET without ORDER BY.\n"
+        "  - String length: LEN(col) (not LENGTH).\n"
+    ),
+    "bigquery": (
+        "  - Regex: `REGEXP_CONTAINS(col, r'pattern')` — returns TRUE if match.\n"
+        "    Negate with `NOT REGEXP_CONTAINS(...)`.\n"
+        "  - Safe division: `SAFE_DIVIDE(numerator, denominator)` (NULL on zero denominator).\n"
+        "  - Quote reserved names with backticks.\n"
+    ),
+    "snowflake": (
+        "  - Regex: `REGEXP_LIKE(col, 'pattern')` (full-string match) or `col RLIKE 'pattern'`.\n"
+        "  - Float division: `numerator * 1.0 / NULLIF(denominator, 0)`.\n"
+        "  - Unquoted identifiers are uppercased; use double-quotes when case matters.\n"
+    ),
+}
+
+
+def _extract_schema_table_names(schema_context: str) -> list[str]:
+    """Extract table names from a schema context string produced by build_schema_context."""
+    return re.findall(r"^Table:\s+(\S+)", schema_context, re.MULTILINE)
+
+
+class LLMDataQualityOperator(LLMOperator):
+    """
+    Run data-quality checks described in natural language using an LLM agent.
+
+    The agent discovers the database schema, writes and executes SQL queries,
+    applies validators, and produces a
+    :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`.  The task
+    fails when any check does not pass, gating downstream tasks on data quality.
+
+    Supply the data-source toolset and DQ toolset together in ``toolsets``::
+
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+        LLMDataQualityOperator(
+            task_id="quality_check",
+            checks=[
+                DQCheckInput(name="email_nulls", description="Check for null emails"),
+                DQCheckInput(
+                    name="row_count",
+                    description="At least 1000 rows",
+                    validator=row_count_check(min_count=1000),
+                ),
+            ],
+            llm_conn_id="pydanticai_default",
+            toolsets=[
+                SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers"]),
+                SQLDQToolset(),
+            ],
+        )
+
+    When ``toolsets`` is omitted, the operator auto-creates
+    :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
+    :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset` from
+    ``db_conn_id`` and ``table_names``.
+
+    For config-generation backends (e.g. ``SodaDQToolset``), the operator
+    returns the generated config string as its XCom value instead of a report.
+
+    :param checks: List of :class:`~airflow.providers.common.ai.utils.dataquality.models.DQCheckInput`
+        objects (or plain dicts with ``name``, ``description``, and optional
+        ``validator`` keys).  Names must be unique.
+    :param toolsets: Pydantic-AI toolsets for the agent.  Must include exactly
+        one :class:`~airflow.providers.common.ai.toolsets.dataquality.base.BaseDQToolset`
+        subclass alongside a data-source toolset.  When ``None``, the operator
+        auto-creates toolsets from ``db_conn_id``.
+    :param db_conn_id: Connection ID for auto-creating an
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`.
+        Ignored when ``toolsets`` is provided.
+    :param table_names: Tables passed to the auto-created
+        :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` as
+        ``allowed_tables``.  Ignored when ``toolsets`` is provided.
+    :param schema_context: Additional schema description injected into the
+        system prompt.  Useful when the data source cannot be introspected
+        at runtime.
+    """
+
+    template_fields: Sequence[str] = (
+        *LLMOperator.template_fields,
+        "checks",
+        "db_conn_id",
+        "table_names",
+        "schema_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        checks: list[DQCheckInput | dict[str, Any]],
+        toolsets: list[AbstractToolset] | None = None,
+        db_conn_id: str | None = None,
+        table_names: list[str] | None = None,
+        schema_context: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Pop operator-specific params that must not reach BaseOperator.__init__.
+        # Using kwargs.pop (rather than named params) is necessary because Airflow's
+        # apply_defaults metaclass captures the full **kwargs dict and uses it for
+        # task-map serialization; named params in the dict are NOT removed from the
+        # snapshot, so they would be re-injected on deserialization and reach
+        # BaseOperator as unknown kwargs.
+        durable: bool = kwargs.pop("durable", False)
+
+        kwargs.pop("output_type", None)
+        kwargs.setdefault("prompt", "Run the data-quality checks.")
+        super().__init__(**kwargs)
+
+        self.checks: list[DQCheckInput] = (
+            [DQCheckInput.coerce(c) for c in checks] if isinstance(checks, list) else checks  # type: ignore[assignment]
+        )
+        self.toolsets = toolsets
+        self.db_conn_id = db_conn_id
+        self.table_names = table_names
+        self.schema_context = schema_context
+        self.durable = durable
+
+        self._validate_checks()
+
+        if not toolsets and db_conn_id is None:
+            raise ValueError("Either toolsets or db_conn_id must be provided.")
+
+    def execute(self, context: Context) -> Any:
+        """
+        Run the LLM agent to execute all data-quality checks.
+
+        When ``require_approval=True``, the task runs in two phases:
+
+        1. **Phase 1** — the LLM discovers schema and selects validators *without
+           executing any SQL queries*.  The resulting plan (check names,
+           descriptions, and validator assignments) is shown to a human reviewer.
+        2. **Phase 2** (in :meth:`execute_complete`) — after approval, SQL queries
+           are executed and validators are applied in pure Python to produce the
+           :class:`~...DQReport`.
+
+        :returns: :class:`~airflow.providers.common.ai.utils.dataquality.models.DQReport`
+            as a dict when the DQ toolset is in ``"execute"`` mode, or a config
+            string when in ``"generate"`` mode.
+        :raises DQCheckFailedError: When any check fails in ``"execute"`` mode.
+        """
+        if self.require_approval:
+            plan, _ = self._run_plan_phase(context)
+            self.defer_for_approval(  # type: ignore[misc]
+                context,
+                plan.model_dump_json(),
+                subject=f"Review DQ plan for task `{self.task_id}`",
+                body=self._build_plan_approval_body(plan),
+            )
+        return self._run_and_report(context)
+
+    def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any:
+        """
+        Phase 2: execute SQL and apply validators after human approval.
+
+        Called automatically by Airflow when the HITL trigger fires.  The base
+        class validates the approval decision (raises on reject or timeout).
+        Then each SQL query from the approved :class:`~...DQPlan` is executed in
+        pure Python, the metric values are extracted, and validators are applied
+        — no LLM calls.
+
+        :param context: Airflow task context.
+        :param generated_output: The :class:`~...DQPlan` JSON that was deferred.
+        :param event: Trigger event payload.
+        :returns: Same as :meth:`execute`.
+        :raises HITLRejectException: If the reviewer rejected the plan.
+        """
+        approved_output = super().execute_complete(context, generated_output, event)  # type: ignore[misc]
+        plan = DQPlan.model_validate_json(approved_output)
+
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        report = self._execute_plan_validators(plan, dq_toolset, toolsets)
+        context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+        if not report.passed:
+            raise DQCheckFailedError(report.failure_summary)
+        return report.model_dump()
+
+    # ------------------------------------------------------------------
+    # Core execution helpers
+    # ------------------------------------------------------------------
+
+    def _run_plan_phase(self, context: Context) -> tuple[DQPlan, BaseDQToolset]:
+        """
+        Phase 1 of the two-phase approval flow.
+
+        Runs the LLM agent in planning mode: schema-discovery tools
+        (``list_tables``, ``get_schema``, ``check_query``) are available, but
+        the ``query`` and ``apply_validator`` tools are hidden.  The agent
+        writes SQL strings and selects validators for each check, then outputs
+        a :class:`~...DQPlan` — without executing any SQL.
+
+        :returns: ``(plan, dq_toolset)`` — the plan for the approval body and the
+            toolset (already configured with checks) for Phase 2.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset._planning_mode = True  # omit apply_validator from tool list
+        dq_toolset.set_checks(self.checks)
+
+        # Hide the SQL 'query' tool so the LLM cannot execute DQ queries.
+        for ts in toolsets:
+            if ts is not dq_toolset and hasattr(ts, "_query"):
+                ts._planning_mode = True  # type: ignore[attr-defined]
+
+        instructions = self._build_plan_system_prompt(toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+        storage = None
+
+        if self.durable:
+            agent, counter, storage = self._build_durable_agent(
+                context,
+                DQPlan,
+                instructions,
+                logged_toolsets,
+            )
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=DQPlan,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay (plan phase): model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        if storage is not None:
+            storage.cleanup()
+
+        return result.output, dq_toolset
+
+    def _run_and_report(self, context: Context) -> Any:
+        """
+        Resolve toolsets, run the LLM agent, push XCom, and handle failures.
+
+        This is the single place where the agent executes regardless of whether
+        the task took the direct path or the approval-gated path.
+        """
+        toolsets = self._resolve_toolsets()
+        dq_toolset = self._find_dq_toolset(toolsets)
+        dq_toolset.set_checks(self.checks)
+
+        output_type: type = DQReport if dq_toolset.output_mode == "execute" else str
+        instructions = self._build_system_prompt(dq_toolset, toolsets)
+        logged_toolsets = wrap_toolsets_for_logging(toolsets, self.log)
+        storage = None
+
+        if self.durable:
+            agent, counter, storage = self._build_durable_agent(
+                context,
+                output_type,
+                instructions,
+                logged_toolsets,
+            )
+        else:
+            agent = self.llm_hook.create_agent(
+                output_type=output_type,
+                instructions=instructions,
+                toolsets=logged_toolsets,
+                **self.agent_params,
+            )
+            counter = None
+
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        log_run_summary(self.log, result)
+
+        if counter is not None and (counter.replayed_model > 0 or counter.replayed_tool > 0):
+            self.log.info(
+                "Durable cache replay: model_steps=%d/%d, tool_steps=%d/%d",
+                counter.replayed_model,
+                counter.replayed_model + counter.cached_model,
+                counter.replayed_tool,
+                counter.replayed_tool + counter.cached_tool,
+            )
+
+        if storage is not None:
+            storage.cleanup()
+
+        output = result.output
+
+        if dq_toolset.output_mode == "execute":
+            report: DQReport = output
+            context["task_instance"].xcom_push(key="dq_report", value=report.model_dump())
+            if not report.passed:
+                raise DQCheckFailedError(report.failure_summary)
+            return report.model_dump()
+
+        return output
+
+    def _build_durable_agent(
+        self,
+        context: Context,
+        output_type: type,
+        instructions: str,
+        toolsets: list[AbstractToolset],
+    ) -> tuple[Any, Any, Any]:
+        """Build a pydantic-ai Agent with CachingModel and CachingToolset wrappers."""
+        from pydantic_ai import Agent
+
+        from airflow.providers.common.ai.durable.caching_model import CachingModel
+        from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
+        from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+        from airflow.providers.common.ai.durable.storage import DurableStorage
+
+        ti = context["task_instance"]
+        storage = DurableStorage(
+            dag_id=ti.dag_id,
+            task_id=ti.task_id,
+            run_id=ti.run_id,
+            map_index=ti.map_index if ti.map_index is not None else -1,
+        )
+        counter = DurableStepCounter()
+
+        wrapped_model = CachingModel(self.llm_hook.get_conn(), storage=storage, counter=counter)
+        cached_toolsets = [CachingToolset(ts, storage=storage, counter=counter) for ts in toolsets]
+
+        agent = Agent(
+            wrapped_model,
+            output_type=output_type,
+            instructions=instructions,
+            toolsets=cached_toolsets,
+            **self.agent_params,
+        )
+        return agent, counter, storage
+
+    def _build_plan_approval_body(self, plan: DQPlan) -> str:
+        """
+        Build a Markdown body for the HITL approval form.
+
+        Shows the LLM-generated SQL queries and validator selections so the
+        reviewer can inspect the plan before any SQL is executed or any
+        validator is applied.
+        """
+        lines = [
+            "## Data Quality Check Plan \u2014 Awaiting Approval",
+            "",
+            f"**Total checks:** {len(plan.checks)}",
+            "",
+            "Review the SQL queries and validator selections below.  ",
+            "**Approve** to execute the SQL, apply validators, and produce the quality report.  ",
+            "**Reject** to cancel without executing anything.",
+            "",
+            "| # | Check | Metric Key | Row Level | Validator | Args |",
+            "|---|-------|------------|-----------|-----------|------|",
+        ]
+        for i, cp in enumerate(plan.checks, 1):
+            args_display = json.dumps(cp.validator_args) if cp.validator_args else "{}"
+            lines.append(
+                f"| {i} | `{cp.check_name}` | `{cp.metric_key}` "
+                f"| {'Yes' if cp.row_level else 'No'} "
+                f"| `{cp.validator_name}` | `{args_display}` |"
+            )
+
+        lines += ["", "---", "", "### SQL Queries", ""]
+        for cp in plan.checks:
+            lines += [f"**{cp.check_name}**", "", "```sql", cp.sql_query.strip(), "```", ""]
+
+        return "\n".join(lines)
+
+    def _execute_plan_validators(
+        self,
+        plan: DQPlan,
+        dq_toolset: BaseDQToolset,
+        toolsets: list[AbstractToolset],
+    ) -> DQReport:
+        """
+        Phase 2 of the two-phase approval flow.
+
+        For each check in the approved :class:`~...DQPlan`:
+
+        1. Executes the SQL query via the data-source toolset.
+        2. Extracts the metric value from the result.
+        3. Applies the chosen validator in pure Python (no LLM calls).
+        4. Builds the final :class:`~...DQReport`.
+
+        :param plan: The approved :class:`~...DQPlan` from Phase 1.
+        :param dq_toolset: The DQ toolset (must implement ``_apply_validator``).
+        :param toolsets: All toolsets (used to find the SQL toolset for execution).
+        :raises ValueError: If no SQL-capable toolset or no validator executor found.
+        """
+        from airflow.providers.common.ai.utils.dataquality.models import DQCheckResult
+
+        sql_toolset = next((ts for ts in toolsets if hasattr(ts, "_query")), None)
+        if sql_toolset is None:
+            raise ValueError(
+                "require_approval Phase 2 requires a toolset with a _query() method "
+                "(e.g. SQLToolset). Add a data-source toolset to the toolsets list."
+            )
+
+        apply_fn = getattr(dq_toolset, "_apply_validator", None)
+        if apply_fn is None:
+            raise ValueError(
+                "require_approval two-phase execution requires a DQ toolset that "
+                "implements _apply_validator (e.g. SQLDQToolset)."
+            )
+
+        results: list[DQCheckResult] = []
+        for cp in plan.checks:
+            try:
+                raw_result = sql_toolset._query(cp.sql_query)  # type: ignore[union-attr]
+                result_data = json.loads(raw_result)
+                rows: list[dict[str, Any]] = result_data.get("rows", [])
+            except Exception as exc:
+                results.append(
+                    DQCheckResult(
+                        check_name=cp.check_name,
+                        passed=False,
+                        failure_reason=f"SQL execution failed: {exc}",
+                        sql_query=cp.sql_query,
+                        validator_info={"name": cp.validator_name, "args": cp.validator_args},
+                    )
+                )
+                continue
+
+            if cp.row_level:
+                metric_value: Any = [row.get(cp.metric_key) for row in rows]
+            else:
+                metric_value = rows[0].get(cp.metric_key) if rows else None
+
+            raw = apply_fn(cp.check_name, metric_value, cp.validator_name, cp.validator_args)
+            parsed = json.loads(raw)
+            passed = bool(parsed.get("passed", False))
+            result_value = parsed.get("value", metric_value)
+            results.append(
+                DQCheckResult(
+                    check_name=cp.check_name,
+                    passed=passed,
+                    value=result_value,
+                    failure_reason=parsed.get("reason") if not passed else None,
+                    metric_key=cp.metric_key,
+                    sql_query=cp.sql_query,
+                    validator_info={"name": cp.validator_name, "args": cp.validator_args},
+                )
+            )
+
+        return DQReport.build(results)
+
+    # ------------------------------------------------------------------
+    # Toolset / prompt helpers
+    # ------------------------------------------------------------------
+
+    def _resolve_toolsets(self) -> list[AbstractToolset]:
+        """
+        Return explicit toolsets or auto-create from ``db_conn_id``.
+
+        When ``toolsets`` is provided but contains no
+        :class:`BaseDQToolset`, a default
+        :class:`~airflow.providers.common.ai.toolsets.dataquality.sql.SQLDQToolset`
+        is appended automatically.
+        """
+        if self.toolsets:
+            has_dq = any(isinstance(ts, BaseDQToolset) for ts in self.toolsets)
+            if not has_dq:
+                from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+
+                return [*self.toolsets, SQLDQToolset()]
+            return list(self.toolsets)
+
+        from airflow.providers.common.ai.toolsets.dataquality.sql import SQLDQToolset
+        from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+        return [
+            SQLToolset(db_conn_id=self.db_conn_id, allowed_tables=self.table_names),  # type: ignore[arg-type]
+            SQLDQToolset(),
+        ]
+
+    @staticmethod
+    def _find_dq_toolset(toolsets: list[AbstractToolset]) -> BaseDQToolset:
+        """Return the first :class:`BaseDQToolset` found in *toolsets*."""
+        for toolset in toolsets:
+            if isinstance(toolset, BaseDQToolset):
+                return toolset
+        raise ValueError(
+            "No BaseDQToolset found in toolsets. "
+            "Add SQLDQToolset (or another BaseDQToolset subclass) to the toolsets list."
+        )
+
+    def _build_system_prompt(self, dq_toolset: BaseDQToolset, toolsets: list[AbstractToolset]) -> str:
+        """Return the full DQ system prompt for the normal (non-planning) execution path."""
+        prompt = self._make_prompt(_DQ_SYSTEM_PROMPT, toolsets)
+
+        fixed_checks = [c.name for c in self.checks if c.validator is not None]
+        if fixed_checks:
+            names = ", ".join(f'"{n}"' for n in fixed_checks)
+            prompt += (
+                "\nFIXED VALIDATORS:\n"
+                f"  Checks {names} have pre-assigned validators.\n"
+                '  For these, call apply_validator with validator_name="fixed" — '
+                "the system uses the pre-configured validator automatically.\n"
+            )
+
+        if self.system_prompt:
+            prompt += f"\nAdditional instructions:\n{self.system_prompt}\n"
+
+        return prompt
+
+    def _build_plan_system_prompt(self, toolsets: list[AbstractToolset]) -> str:
+        """Return the full DQ system prompt for Phase 1 (planning mode, no apply_validator)."""
+        prompt = self._make_prompt(_DQ_PLAN_SYSTEM_PROMPT, toolsets)
+
+        # In planning mode the FIXED VALIDATORS section uses different wording:
+        # the LLM records "fixed" in the DQPlan instead of calling apply_validator.
+        fixed_checks = [c.name for c in self.checks if c.validator is not None]
+        if fixed_checks:
+            names = ", ".join(f'"{n}"' for n in fixed_checks)
+            prompt += (
+                "\nFIXED VALIDATORS:\n"
+                f"  Checks {names} have pre-assigned validators.\n"
+                '  For these, set validator_name: "fixed" and validator_args: {} '
+                "in your DQCheckPlan output.\n"
+            )
+
+        if self.system_prompt:
+            prompt += f"\nAdditional instructions:\n{self.system_prompt}\n"
+
+        return prompt
+
+    def _make_prompt(self, base: str, toolsets: list[AbstractToolset]) -> str:
+        """Inject dialect, schema context, and fixed-validator sections into *base*."""
+        prompt = base
+
+        dialect = self._detect_sql_dialect(toolsets)
+        if dialect:
+            notes = _DIALECT_SQL_NOTES.get(dialect, "")
+            dialect_section = f"\nSQL DIALECT: {dialect.upper()}\n"
+            if notes:
+                dialect_section += "  Adapt your SQL to the following dialect-specific rules:\n" + notes
+            prompt += dialect_section
+
+        if self.schema_context:
+            table_names = _extract_schema_table_names(self.schema_context)
+            if table_names:
+                prompt += (
+                    "\nTABLE NAME CONSTRAINT:\n"
+                    f"  The ONLY tables you may reference in FROM clauses are: {', '.join(table_names)}.\n"
+                    "  Use these exact names — do not rename, abbreviate, or invent new table names.\n"
+                )
+            prompt += f"\nSchema context:\n{self.schema_context}\n"
+
+        return prompt
+
+    @staticmethod
+    def _detect_sql_dialect(toolsets: list[AbstractToolset]) -> str | None:
+        """Return the sqlglot dialect of the first dialect-aware toolset found."""
+        for toolset in toolsets:
+            dialect = getattr(toolset, "sqlglot_dialect", None)
+            if dialect:
+                return dialect
+        return None

Review Comment:
   `_detect_sql_dialect()` directly evaluates `toolset.sqlglot_dialect`. If that property triggers connection resolution (e.g. `BaseHook.get_connection`) and raises, prompt construction will fail even though dialect hints are optional. Treat dialect detection as best-effort by catching exceptions and continuing.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3105702926


##########
providers/common/ai/src/airflow/providers/common/ai/utils/db_schema.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Shared database hook and schema introspection utilities.
+
+These helpers are used by both :class:`~airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator`
+and :class:`~airflow.providers.common.ai.operators.llm_data_quality.LLMDataQualityOperator` to
+avoid code duplication while keeping both operators decoupled from each other.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from airflow.providers.common.compat.sdk import BaseHook
+
+if TYPE_CHECKING:
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+# SQLAlchemy dialect_name → sqlglot dialect mapping for names that differ.
+SQLALCHEMY_TO_SQLGLOT_DIALECT: dict[str, str] = {
+    "postgresql": "postgres",
+    "mssql": "tsql",
+}
+
+
+def get_db_hook(db_conn_id: str) -> DbApiHook:
+    """
+    Return a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` for *db_conn_id*.
+
+    :param db_conn_id: Airflow connection ID that resolves to a ``DbApiHook``.
+    :raises ValueError: If the connection does not resolve to a ``DbApiHook``.
+    """
+    # Lazy load to avoid hard dependency on common.sql
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+    connection = BaseHook.get_connection(db_conn_id)
+    hook = connection.get_hook()
+    if not isinstance(hook, DbApiHook):
+        raise ValueError(
+            f"Connection {db_conn_id!r} does not provide a DbApiHook. Got {type(hook).__name__}."
+        )
+    return hook
+
+
+def resolve_dialect(db_hook: DbApiHook | None, explicit_dialect: str | None) -> str | None:
+    """
+    Resolve the SQL dialect from an explicit parameter or a database hook.
+
+    Normalises SQLAlchemy dialect names to sqlglot equivalents
+    (e.g. ``postgresql`` → ``postgres``).
+
+    :param db_hook: Database hook to read ``dialect_name`` from when *explicit_dialect* is absent.
+    :param explicit_dialect: Caller-supplied dialect string; takes priority over the hook.
+    :return: Resolved dialect string, or ``None`` when neither source provides one.
+    """
+    raw = explicit_dialect
+    if not raw and db_hook and hasattr(db_hook, "dialect_name"):
+        candidate = db_hook.dialect_name
+        raw = candidate if isinstance(candidate, str) else None
+    if raw:
+        return SQLALCHEMY_TO_SQLGLOT_DIALECT.get(raw, raw)
+    return None
+
+
+def build_schema_context(
+    *,
+    db_hook: DbApiHook | None,
+    table_names: list[str] | None,
+    schema_context: str | None,
+    datasource_config: DataSourceConfig | None,
+) -> str:
+    """
+    Return a schema description string suitable for inclusion in an LLM prompt.
+
+    Resolution order:
+    1. *schema_context* — returned as-is when provided (manual override).
+    2. DB introspection via *db_hook* + *table_names*.
+    3. Object-storage introspection via *datasource_config*.
+    4. Empty string when none of the above are available.
+
+    :param db_hook: Hook used for relational-database schema introspection.
+    :param table_names: Table names to introspect via *db_hook*.
+    :param schema_context: Manual schema description; bypasses introspection when set.
+    :param datasource_config: DataFusion datasource config for object-storage schema.
+    :raises ValueError: If *table_names* are provided but none yield schema information.
+    """
+    if schema_context:
+        return schema_context
+
+    if (db_hook and table_names) or datasource_config:
+        return _introspect_schemas(
+            db_hook=db_hook,
+            table_names=table_names,
+            datasource_config=datasource_config,
+        )
+
+    return ""
+
+
+def _introspect_schemas(
+    *,
+    db_hook: DbApiHook | None,
+    table_names: list[str] | None,
+    datasource_config: DataSourceConfig | None,
+) -> str:
+    """Build schema context by introspecting tables and/or object-storage sources."""
+    parts: list[str] = []
+    table_to_columns: dict[str, list[dict[str, str]]] = {}
+
+    if table_names and db_hook is None:
+        raise ValueError("table_names requires db_conn_id so table schema can be introspected.")

Review Comment:
   lets validate this early in the operator init, instead of validating after task run from the execute.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.

gopidesupavan commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4274505612

   Thanks @cetingokhan i see your intent but here is my perception on the current approach after my thourlgy gone through this.
   
   
   Right now, the same check intent is represented in a few different places: `prompts`, `validators`, and then some validator-specific metadata such as `llm_context`, `check_category`, and row-level attributes. One possible simplification would be to let planning use a small set of registered DQ check tools, and have the final `DQPlan` store the selected `validation.kind` plus concrete `params` for each check.
   
   That would allow users to declare each check once, reduce how much validator behavior has to be described in the prompt, and still keep execution fully deterministic after planning. Put differently: the agent could help with check-type selection and threshold normalization, while SQL execution and final validation would remain in normal Python. I think that could make the API easier to use and reduce some of the coupling between `prompts` and `validators`.
   
   
   so user facing front end would like this:
   
   ```
   LLMDataQualityOperator(
           checks=[
              {"name": "null_order_id", "description": "order_id should never be null"},
              {"name": "duplicate_orders", "description": "duplicate order_id percentage should be 0%"},
              {"name": "min_row_count", "description": "orders table should have at least 10,000 rows"},
           ],
       )
   ```
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.

cetingokhan commented on code in PR #62963:
URL: https://github.com/apache/airflow/pull/62963#discussion_r3077684100


##########
providers/common/ai/src/airflow/providers/common/ai/utils/dq_planner.py:
##########
@@ -0,0 +1,908 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+SQL-based data-quality plan generation and execution.
+
+:class:`SQLDQPlanner` is the single entry-point for all SQL DQ logic.
+It is deliberately kept separate from the operator so it can be unit-tested
+without an Airflow context and later swapped for GEX/SODA planners without
+touching the operator.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+try:
+    from airflow.providers.common.ai.utils.sql_validation import (
+        DEFAULT_ALLOWED_TYPES,
+        SQLSafetyError,
+        validate_sql as _validate_sql,
+    )
+except ImportError as e:
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+    raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.utils.db_schema import build_schema_context, resolve_dialect
+from airflow.providers.common.ai.utils.dq_models import DQCheckGroup, DQPlan, RowLevelResult, UnexpectedResult
+from airflow.providers.common.ai.utils.logging import log_run_summary
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.messages import ModelMessage
+
+    from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+    from airflow.providers.common.sql.config import DataSourceConfig
+    from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+log = logging.getLogger(__name__)
+
+_MAX_CHECKS_PER_GROUP = 5
+# Maximum rows fetched from DB per chunk during row-level processing — avoids loading the
+# entire result set into memory at once.
+_ROW_LEVEL_CHUNK_SIZE = 10_000
+# Hard cap on violation samples stored per check — independent of SQL LIMIT and chunk size.
+_MAX_VIOLATION_SAMPLES = 100
+
+_PLANNING_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 while keeping each group \
+focused and manageable.
+
+GROUPING STRATEGY (multi-dimensional):
+  Group checks by **(target_table, check_category)**.  Checks on the same table
+  that belong to different categories MUST be in separate groups.
+
+  Allowed check_category values (assign one per check based on its description):
+    - null_check      — null / missing value counts or percentages
+    - uniqueness      — duplicate detection, cardinality checks
+    - validity        — regex / format / pattern matching on string columns
+    - numeric_range   — range, bounds, or statistical checks on numeric columns
+    - row_count       — total row counts or existence checks
+    - string_format   — length, encoding, whitespace, or character-set checks
+    - row_level       — per-row or anomaly checks that evaluate individual records
+
+  Row-level checks still follow the same grouping rule: group by (target_table, check_category="row_level").
+  MAX {max_checks_per_group} CHECKS PER GROUP:
+    If a (table, category) pair has more than {max_checks_per_group} checks,
+    split them into sub-groups of at most {max_checks_per_group}.
+
+  GROUP-ID NAMING:
+    Use the pattern "{{table}}_{{category}}_{{part}}".
+    Examples: customers_null_check_1, orders_validity_1, orders_validity_2
+
+  RATIONALE:
+    Keeping string-column checks (validity, string_format) apart from
+    numeric-column checks (numeric_range, null_check on numbers) produces
+    simpler SQL and makes failures easier to diagnose.
+
+  CORRECT (two groups for same table, different categories):
+    Group customers_null_check_1:
+      SELECT
+        (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+        (COUNT(CASE WHEN name IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_name_pct
+      FROM customers
+
+    Group customers_validity_1:
+      SELECT
+        COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+      FROM customers
+
+  WRONG (mixing null-check and regex-validity in one group):
+    SELECT
+      (COUNT(CASE WHEN email IS NULL THEN 1 END) * 100.0 / COUNT(*)) AS null_email_pct,
+      COUNT(CASE WHEN phone NOT LIKE '+___-___-____' THEN 1 END) AS invalid_phone_fmt
+    FROM customers
+
+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 {dialect} syntax.
+  6. Each check must appear in exactly ONE group.
+  7. Each check must have a check_category from the allowed list above.
+  8. Return a valid DQPlan object. No extra commentary.
+"""
+
+_DATAFUSION_SYNTAX_SECTION = """\
+
+DATAFUSION SQL SYNTAX RULES:
+  The target engine is Apache DataFusion.  Observe these syntax differences
+  from standard PostgreSQL / ANSI SQL:
+
+  1. NO "FILTER (WHERE ...)" clause.  Use CASE expressions instead:
+       WRONG:  COUNT(*) FILTER (WHERE email IS NULL)
+       RIGHT:  COUNT(CASE WHEN email IS NULL THEN 1 END)
+
+  2. Regex matching uses the tilde operator:
+       column ~ 'pattern'    (match)
+       column !~ 'pattern'   (no match)
+     Do NOT use SIMILAR TO or POSIX-style ~* (case-insensitive).
+
+  3. CAST syntax — prefer CAST(expr AS type) over :: shorthand.
+
+  4. String functions: Use CHAR_LENGTH (not LEN), SUBSTR (not SUBSTRING with FROM/FOR).
+
+  5. Integer division: DataFusion performs integer division for INT/INT.
+     Use CAST(expr AS DOUBLE) to force floating-point division.
+
+  6. Boolean literals: Use TRUE / FALSE (not 1 / 0).
+
+  7. LIMIT is supported.  OFFSET is supported.  FETCH FIRST is NOT supported.
+
+  8. NULL handling: COALESCE, NULLIF, IFNULL are all supported.
+     NVL and ISNULL are NOT supported.
+"""
+
+_UNEXPECTED_QUERY_PROMPT_SECTION = """\
+
+UNEXPECTED VALUE COLLECTION:
+  For checks whose check_category is "validity" or "string_format", also
+  generate an unexpected_query field on the DQCheck.  This query must:
+    - SELECT the primary key column(s) and the column(s) being validated
+    - WHERE the row violates the check condition (the negation of the check)
+    - LIMIT {sample_size}
+    - Use {dialect} syntax
+    - Be a standalone SELECT (not a subquery of the group query)
+
+  For all other categories (null_check, uniqueness, numeric_range, row_count),
+  set unexpected_query to null — these are aggregate checks where individual
+  violating rows are not meaningful.
+
+  Example for a phone-format validity check:
+    unexpected_query: "SELECT id, phone FROM customers WHERE phone !~ '^\\d{{4}}-\\d{{4}}-\\d{{4}}$' LIMIT 100"
+"""
+
+_ROW_LEVEL_PROMPT_SECTION = """
+
+ROW-LEVEL CHECKS:
+  Some checks are marked as row_level.  For these:
+    - Generate a SELECT that returns the primary key column(s) and the column
+      being validated.  Do NOT aggregate.
+    - Set row_level = true on the DQCheck entry.
+    - metric_key must be the name of the column containing the value to validate
+      (the Python validator will read row[metric_key] for each row).
+    - {row_level_limit_clause}
+    - Place ALL row-level checks for the same table in a single group.
+
+  Row-level check names that require this treatment: {row_level_check_names}
+"""
+
+
+class SQLDQPlanner:
+    """
+    Generates and executes a SQL-based :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+    :param llm_hook: Hook used to call the LLM for plan generation.
+    :param db_hook: Hook used to execute generated SQL against the database.
+    :param dialect: SQL dialect forwarded to the LLM prompt and ``validate_sql``.
+        Auto-detected from *db_hook* when ``None``.
+    :param max_sql_retries: Maximum number of times a failing SQL group query is sent
+        back to the LLM for correction before the error is re-raised.  Default ``2``.
+    :param validator_contexts: Pre-built LLM context string from
+        :meth:`~airflow.providers.common.ai.utils.dq_validation.ValidatorRegistry.build_llm_context`.
+        Appended to the system prompt so the LLM knows what metric format each
+        custom validator expects.
+    :param row_validators: Mapping of ``{check_name: row_level_callable}`` for
+        checks that require row-by-row Python validation.  When a check's name
+        appears here, ``execute_plan`` fetches all (or sampled) rows and applies
+        the callable to each value instead of reading a single aggregate scalar.
+    :param row_level_sample_size: Maximum number of rows to fetch for row-level
+        checks.  ``None`` (default) performs a full scan.  A positive integer
+        instructs the LLM to add ``LIMIT N`` to the generated SELECT.
+    """
+
+    def __init__(
+        self,
+        *,
+        llm_hook: PydanticAIHook,
+        db_hook: DbApiHook | None,
+        dialect: str | None = None,
+        max_sql_retries: int = 2,
+        datasource_config: DataSourceConfig | None = None,
+        system_prompt: str = "",
+        agent_params: dict[str, Any] | None = None,
+        collect_unexpected: bool = False,
+        unexpected_sample_size: int = 100,
+        validator_contexts: str = "",
+        row_validators: dict[str, Any] | None = None,
+        row_level_sample_size: int | None = None,
+    ) -> None:
+        self._llm_hook = llm_hook
+        self._db_hook = db_hook
+        self._datasource_config = datasource_config
+        self._dialect = resolve_dialect(db_hook, dialect)
+        # Track whether the execution target is DataFusion so the prompt can
+        # include DataFusion-specific syntax rules.  The dialect stays None
+        # (generic SQL) for sqlglot validation — sqlglot has no DataFusion dialect.
+        self._is_datafusion = db_hook is None and datasource_config is not None
+        # When targeting DataFusion, use PostgreSQL dialect for sqlglot validation
+        # because DataFusion shares regex operators (~, !~) that the generic SQL
+        # parser does not recognise.
+        self._validation_dialect: str | None = "postgres" if self._is_datafusion else self._dialect
+        self._max_sql_retries = max_sql_retries
+        self._extra_system_prompt = system_prompt
+        self._agent_params: dict[str, Any] = agent_params or {}
+        self._collect_unexpected = collect_unexpected
+        self._unexpected_sample_size = unexpected_sample_size
+        self._validator_contexts = validator_contexts
+        self._row_validators: dict[str, Any] = row_validators or {}
+        self._row_level_sample_size = row_level_sample_size
+        self._cached_datafusion_engine: DataFusionEngine | None = None
+        self._plan_agent: Agent[None, DQPlan] | None = None
+        self._plan_all_messages: list[ModelMessage] | None = None
+
+    def build_schema_context(
+        self,
+        table_names: list[str] | None,
+        schema_context: str | None,
+    ) -> str:
+        """
+        Return a schema description string for inclusion in the LLM prompt.
+
+        Delegates to :func:`~airflow.providers.common.ai.utils.db_schema.build_schema_context`.
+        """
+        return build_schema_context(
+            db_hook=self._db_hook,
+            table_names=table_names,
+            schema_context=schema_context,
+            datasource_config=self._datasource_config,
+        )
+
+    def generate_plan(self, prompts: dict[str, str], schema_context: str) -> DQPlan:
+        """
+        Ask the LLM to produce a :class:`~airflow.providers.common.ai.utils.dq_models.DQPlan`.
+
+        The LLM receives the user prompts, schema context, and planning instructions
+        as a structured-output call (``output_type=DQPlan``).  After generation the
+        method verifies that the returned ``check_names`` exactly match
+        ``prompts.keys()``.
+
+        :param prompts: ``{check_name: natural_language_description}`` dict.
+        :param schema_context: Schema description previously built via
+            :meth:`build_schema_context`.
+        :raises ValueError: If the LLM's plan does not cover every prompt key
+            exactly once.
+        """
+        dialect_label = self._dialect or ("DataFusion-compatible SQL" if self._is_datafusion else "SQL")
+        system_prompt = _PLANNING_SYSTEM_PROMPT.format(
+            dialect=dialect_label, max_checks_per_group=_MAX_CHECKS_PER_GROUP
+        )
+
+        if self._is_datafusion:
+            system_prompt += _DATAFUSION_SYNTAX_SECTION
+
+        if self._collect_unexpected:
+            system_prompt += _UNEXPECTED_QUERY_PROMPT_SECTION.format(
+                dialect=dialect_label, sample_size=self._unexpected_sample_size
+            )
+
+        if schema_context:
+            system_prompt += f"\nAvailable schema:\n{schema_context}\n"
+
+        if self._validator_contexts:
+            system_prompt += self._validator_contexts
+
+        if self._row_validators:
+            row_level_check_names = ", ".join(sorted(self._row_validators))
+            if self._row_level_sample_size is not None:
+                limit_clause = f"Add LIMIT {self._row_level_sample_size} to the query."
+            else:
+                limit_clause = "Do NOT add a LIMIT — return all rows."
+            system_prompt += _ROW_LEVEL_PROMPT_SECTION.format(
+                row_level_check_names=row_level_check_names,
+                row_level_limit_clause=limit_clause,
+            )
+
+        if self._extra_system_prompt:
+            system_prompt += f"\nAdditional instructions:\n{self._extra_system_prompt}\n"
+
+        user_message = self._build_user_message(prompts)
+
+        log.info("Using system prompt:\n%s", system_prompt)
+        log.info("Using user message:\n%s", user_message)
+

Review Comment:
   I've set logging to Debug mode, but I don't think it's a problem if it prints as info on the first run, since the process continues without generating a context after the cache is loaded. I've set it to debug mode for now, but we can switch it back to info depending on the situation.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] Aip 99 llmdataqualityoperator [airflow]

Posted by "potiuk (via GitHub)" <gi...@apache.org>.

potiuk commented on PR #62963:
URL: https://github.com/apache/airflow/pull/62963#issuecomment-4658957949

   @cetingokhan — A reviewer (@gopidesupavan) has requested changes on this PR, so I've removed the `ready for maintainer review` label — the next step is on your side. Could you address the review comments (push a fix, or reply in-thread explaining why the feedback doesn't apply)? Once addressed, re-request review from @gopidesupavan or re-mark the PR ready and it returns to the maintainer queue. Thank you.
   
   _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](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org