You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "ColtenOuO (via GitHub)" <gi...@apache.org> on 2026/09/19 17:18:26 UTC
[PR] Stop DataFusionToolset from materializing full results for a capped query [airflow]
ColtenOuO opened a new pull request, #73384:
URL: https://github.com/apache/airflow/pull/73384
### Summary
`DataFusionToolset`'s `query` tool returns at most `max_rows` rows (default 50), but it ran the query unbounded and called `to_pydict()` on the whole result before slicing, so a broad query like `SELECT * FROM big_table` could exhaust worker memory. The query is now capped at `max_rows + 1` rows through `DataFusionEngine.execute_query(max_rows=...)`; the extra row only signals truncation. What the agent sees is unchanged apart from `total_rows` (see below).
### Changes
- Pass `max_rows=self._max_rows + 1` to `execute_query` and stop reporting `total_rows`.
- Bump the `common-sql` lower bound in the `sql` / `common.sql` extras to `>=1.34.0`: the `max_rows` argument of `execute_query` was added in 1.34.0 (#64183), so with 1.33.0 every `query` call would raise `TypeError`.
- Update the existing truncation test to assert the capped call.
### Impact
Peak RSS of one `query` call on a local 3M-row parquet file (65 MiB), `max_rows=50`:
| Query | Before | After | Time before → after |
|---|---|---|---|
| `SELECT *` | 1121 MiB | 219 MiB | 4.20s → 0.05s |
| `WHERE amount > 500` | 860 MiB | 218 MiB | 1.16s → 0.07s |
| `ORDER BY amount DESC` | 1358 MiB | 279 MiB | 3.30s → 0.30s |
| `GROUP BY` (20 rows) | 219 MiB | 219 MiB | unchanged |
Before, memory grew linearly with the result size; after, it no longer depends on the result's row count.
### Trade-off
`total_rows` is no longer returned, since only `max_rows + 1` rows are read, and knowing the total would require a full scan or a second `COUNT(*)` query. An untruncated result loses nothing because `row_count` is already the total. For a truncated result, an agent that needs the count can run `COUNT(*)`, in line with the tool description and `hint`, which already tell it to aggregate in SQL rather than page through a truncated result.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "ColtenOuO (via GitHub)" <gi...@apache.org>.
ColtenOuO commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4056892264
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -101,7 +101,8 @@ class DataFusionToolset(AbstractToolset[Any]):
INSERT INTO, etc.). Default ``False`` — only SELECT-family statements
are permitted.
: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.
Review Comment:
Thanks for pointing it out, fixed!
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil merged PR #73384:
URL: https://github.com/apache/airflow/pull/73384
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4091312046
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -211,20 +215,20 @@ def _query(self, sql: str) -> str:
_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:
Review Comment:
sounds reasonable. thanks for checking!
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4056833972
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -101,7 +101,8 @@ class DataFusionToolset(AbstractToolset[Any]):
INSERT INTO, etc.). Default ``False`` — only SELECT-family statements
are permitted.
: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.
Review Comment:
The docs spellcheck job is red on `materialised`; `materialized` is what the rest of the provider docs use. While touching the wording, the "Rows are fetched, not filtered" paragraph in `providers/common/ai/docs/toolsets.rst` still says `DataFusionToolset` "materializes the full result in the engine before the toolset sees it", which is the behaviour this PR removes, so that sentence wants updating in the same change.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -211,20 +212,20 @@ def _query(self, sql: str) -> str:
_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()
Review Comment:
One side effect of putting the limit on every statement: with `allow_writes=True`, `EXPLAIN` now fails. On datafusion 51.0.0 in breeze, `ctx.sql("EXPLAIN SELECT * FROM t").limit(2).to_pydict()` raises `Internal("Unsupported logical plan: Explain must be root of the plan")`, whereas `CREATE TABLE`, `CREATE VIEW`, `INSERT INTO` and `DROP` return the same pydict with or without the limit. It cannot reach here with `allow_writes=False` (the validator only passes SELECT-family statements), and the agent gets an error JSON rather than a crash, so this is minor. A note in the `allow_writes` docstring would cover it.
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4090524952
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -211,20 +215,20 @@ def _query(self, sql: str) -> str:
_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:
Review Comment:
Is it possible for us to narrow it to only the potential exception?
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4056833972
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -101,7 +101,8 @@ class DataFusionToolset(AbstractToolset[Any]):
INSERT INTO, etc.). Default ``False`` — only SELECT-family statements
are permitted.
: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.
Review Comment:
The "Rows are fetched, not filtered" paragraph in `providers/common/ai/docs/toolsets.rst` (around line 254) still says `DataFusionToolset` "materializes the full result in the engine before the toolset sees it; in both the payload is bounded but the transfer is not". That is the behaviour this PR removes, so that sentence wants updating in the same change.
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "ColtenOuO (via GitHub)" <gi...@apache.org>.
ColtenOuO commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4056893474
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -211,20 +212,20 @@ def _query(self, sql: str) -> str:
_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()
Review Comment:
addressed in 5200f44
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "ColtenOuO (via GitHub)" <gi...@apache.org>.
ColtenOuO commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4091297055
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -211,20 +215,20 @@ def _query(self, sql: str) -> str:
_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:
Review Comment:
The DataFusion Python bindings expose no native exception types, as the toolset docs already note, that's why the retry decision is made by matching the error message against regular expressions rather than catching a specific type.
I also traced it upstream: `DataFusionEngine.execute_query` handles this the same way, with
a bare `except Exception` at L121: https://github.com/apache/airflow/blob/main/providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py#L121
So I don't think we can narrow it here, WDYT
--
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] Stop DataFusionToolset from materializing full results for a capped query [airflow]
Posted by "ColtenOuO (via GitHub)" <gi...@apache.org>.
ColtenOuO commented on PR #73384:
URL: https://github.com/apache/airflow/pull/73384#issuecomment-5806977172
Resolve the conflict :D
--
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