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