This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7722b1b2e03 Stop `DataFusionToolset` from materializing full results
for a capped query (#73384)
7722b1b2e03 is described below
commit 7722b1b2e033c09a3994ed252e5e6beabd383876
Author: Jyun-An Chen <[email protected]>
AuthorDate: Thu Sep 24 19:22:14 2026 +0800
Stop `DataFusionToolset` from materializing full results for a capped query
(#73384)
* Stop DataFusionToolset from materializing full results for a capped query
* Document the DataFusionToolset row limit and its EXPLAIN side effect
The bounded-results documentation still described the old behaviour, where
the engine materialized the whole result before the toolset bounded it, so
it now contradicted the code and the missing total_rows. The row limit also
makes EXPLAIN fail, which is only reachable with allow_writes=True and is
easier to meet in the docstring than in an error from the agent.
---
providers/common/ai/docs/toolsets/datafusion.rst | 8 +++++---
providers/common/ai/docs/toolsets/sql.rst | 11 +++++++----
.../providers/common/ai/toolsets/datafusion.py | 16 ++++++++++------
.../unit/common/ai/toolsets/test_datafusion.py | 22 ++++++++++++----------
4 files changed, 34 insertions(+), 23 deletions(-)
diff --git a/providers/common/ai/docs/toolsets/datafusion.rst
b/providers/common/ai/docs/toolsets/datafusion.rst
index ae44d889ebc..fc83a90f8b9 100644
--- a/providers/common/ai/docs/toolsets/datafusion.rst
+++ b/providers/common/ai/docs/toolsets/datafusion.rst
@@ -103,9 +103,11 @@ Iceberg is looked up by ``db_name`` instead, and
``DataSourceConfig`` raises
this toolset "does not prevent the agent from reading any registered data
source". The registration list is therefore the whole boundary: register
exactly what the agent may read.
-- It bounds the payload, not the scan. DataFusion has already materialized the
- full result before ``max_rows`` and ``max_result_bytes`` apply, so those
limits
- protect the model's context, not the cost of the query.
+- It bounds what the engine materializes, not what it scans. The ``query`` tool
+ runs the statement with a ``LIMIT`` of ``max_rows + 1``, so DataFusion never
+ builds a larger result than that, but a plan that has to read every row
before
+ it can return one -- an aggregation, a sort, a late-matching filter -- still
+ pays for the whole scan.
- It cannot tell failure kinds apart precisely. The DataFusion Python bindings
expose no native exception types, so the retry decision is made by matching
the
error message against regular expressions, which a wording change upstream
can
diff --git a/providers/common/ai/docs/toolsets/sql.rst
b/providers/common/ai/docs/toolsets/sql.rst
index fa82aed4559..b34433ef8bd 100644
--- a/providers/common/ai/docs/toolsets/sql.rst
+++ b/providers/common/ai/docs/toolsets/sql.rst
@@ -231,9 +231,10 @@ costs. How much is saved depends on the driver: with a
server-side cursor the
remaining rows are never sent, while a client-buffering driver (psycopg2's
default
cursor, MySQLdb) has already received them and only the per-row conversion is
skipped.
Hooks whose cursor is not DBAPI 2.0 (``ExasolHook`` passes a pyexasol
statement) fall
-back to a full fetch, and ``DataFusionToolset`` materializes the full result
in the
-engine before the toolset sees it; in both the payload is bounded but the
transfer is
-not.
+back to a full fetch, where the payload is bounded but the transfer is not.
+``DataFusionToolset`` pushes the bound into the query instead: it runs the
statement
+with a DataFusion ``LIMIT`` of ``max_rows + 1``, so the engine never
materializes more
+than that and the extra row only signals truncation.
**A byte budget bounds the payload.** ``max_rows`` caps rows, which says
nothing about
size -- one row of a 3000-column table is larger than a thousand rows of a
narrow one.
@@ -251,7 +252,9 @@ it. The result says which limit it hit:
or the column names alone exceed the budget, the result carries a ``hint``
telling the
agent to narrow its projection -- the only move that helps. ``total_rows`` is
present
when the driver reports a row count for the query; several (SQLite, some
warehouse
-drivers) do not, and it is then omitted rather than guessed.
+drivers) do not, and it is then omitted rather than guessed.
``DataFusionToolset``
+never reports it, because it reads only ``max_rows + 1`` rows and so has no
total to
+report; an agent that needs one runs ``COUNT(*)``.
The default budget is deliberately generous: the columnar shape alone shrinks
a wide
result several-fold, so results that fit before still fit. Lower
``max_result_bytes``
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
index ebc99a297b1..cc771bda05a 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
@@ -99,9 +99,13 @@ class DataFusionToolset(AbstractToolset[Any]):
:param datasource_configs: One or more DataFusion data-source
configurations.
:param allow_writes: Allow data-modifying SQL (CREATE TABLE, CREATE VIEW,
INSERT INTO, etc.). Default ``False`` — only SELECT-family statements
- are permitted.
+ are permitted. ``EXPLAIN`` reaches the engine only with
``allow_writes=True``,
+ and fails there: the ``max_rows`` limit wraps the plan, and DataFusion
requires
+ ``EXPLAIN`` to be the root of the plan. The agent gets an error
result, not the
+ plan.
:param max_rows: Maximum number of rows returned from the ``query`` tool.
- Default ``50``.
+ Default ``50``. The query is limited to ``max_rows + 1`` rows, so a
large
+ result is never fully materialized; the extra row only signals
truncation.
:param max_result_bytes: Budget for the serialized ``query`` result, in
bytes.
Default 64 KiB. ``max_rows`` bounds rows, which says nothing about
size: one
row of a 3000-column table is larger than a thousand rows of a narrow
one, and
@@ -211,12 +215,13 @@ class DataFusionToolset(AbstractToolset[Any]):
_validate_sql(sql)
engine = self._get_engine()
- pydict = engine.execute_query(sql)
+ try:
+ pydict = engine.session_context.sql(sql).limit(self._max_rows
+ 1).to_pydict()
+ except Exception as e:
+ raise QueryExecutionException(f"Error while executing query:
{e}") from e
col_names = list(pydict.keys())
num_rows = len(next(iter(pydict.values()), []))
- # DataFusion has already materialised the full result, so unlike
SQLToolset
- # there is nothing left to avoid fetching -- only the payload is
bounded.
rows = [[pydict[col][i] for col in col_names] for i in
range(min(num_rows, self._max_rows))]
return build_query_result(
col_names,
@@ -224,7 +229,6 @@ class DataFusionToolset(AbstractToolset[Any]):
max_rows=self._max_rows,
max_result_bytes=self._max_result_bytes,
more_rows_available=num_rows > self._max_rows,
- total_rows=num_rows,
)
except SQLSafetyError as ex:
log.warning("query failed SQL safety validation: %s", ex)
diff --git
a/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
index f5d2572b9b1..98a8d21f853 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
@@ -65,7 +65,7 @@ def _make_mock_engine(
for tname in tables:
mock.session_context.table(tname).schema.return_value = arrow_fields
- mock.execute_query.return_value = (
+
mock.session_context.sql.return_value.limit.return_value.to_pydict.return_value
= (
query_result
if query_result is not None
else {
@@ -187,7 +187,8 @@ class TestDataFusionToolsetQuery:
def test_truncates_at_max_rows(self):
cfg = _make_mock_datasource_config()
ts = DataFusionToolset([cfg], max_rows=1)
- ts._engine = _make_mock_engine(query_result={"id": [1, 2, 3], "name":
["a", "b", "c"]})
+ engine = _make_mock_engine(query_result={"id": [1, 2], "name": ["a",
"b"]})
+ ts._engine = engine
result = asyncio.run(
ts.call_tool(
@@ -198,10 +199,11 @@ class TestDataFusionToolsetQuery:
)
)
data = json.loads(result)
+ engine.session_context.sql.assert_called_once_with("SELECT id, name
FROM sales_data")
+
engine.session_context.sql.return_value.limit.assert_called_once_with(2)
assert data["rows"] == [[1, "a"]]
assert data["truncated"] is True
assert data["truncated_by"] == "max_rows"
- assert data["total_rows"] == 3
def test_handles_empty_result(self):
cfg = _make_mock_datasource_config()
@@ -324,12 +326,10 @@ class TestDataFusionToolsetQueryErrors:
assert matches is expected
def test_query_execution_exception_returns_error_json(self):
- from airflow.providers.common.sql.datafusion.exceptions import
QueryExecutionException
-
cfg = _make_mock_datasource_config()
ts = DataFusionToolset([cfg])
engine = _make_mock_engine()
- engine.execute_query.side_effect = QueryExecutionException("execution
failed")
+ engine.session_context.sql.side_effect = RuntimeError("execution
failed")
ts._engine = engine
result = asyncio.run(
@@ -382,12 +382,14 @@ class TestDataFusionToolsetQueryErrors:
assert "get_schema" in exc_info.value.message
assert "list_tables" in exc_info.value.message
- def test_unexpected_exception_propagates(self):
+ @patch(
+ "airflow.providers.common.ai.toolsets.datafusion.build_query_result",
+ side_effect=TypeError("unexpected error"),
+ )
+ def test_unexpected_exception_propagates(self, _):
cfg = _make_mock_datasource_config()
ts = DataFusionToolset([cfg])
- engine = _make_mock_engine()
- engine.execute_query.side_effect = TypeError("unexpected error")
- ts._engine = engine
+ ts._engine = _make_mock_engine()
with pytest.raises(TypeError, match="unexpected error"):
asyncio.run(
