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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
