SameerMesiah97 opened a new pull request, #66220: URL: https://github.com/apache/airflow/pull/66220
**Description** This PR is a follow-up to #50454 which adds support for returning query results as a Polars DataFrame in `ElasticsearchSQLHook` by implementing `_get_polars_df`. Instead of relying on `polars.read_database`, which requires DB-API compatibility, this implementation introduces a custom reader that executes Elasticsearch SQL queries using cursor-based pagination and converts the results into a Polars DataFrame. **Rationale** `ElasticsearchSQLCursor` is not compatible with the execution model expected by `polars.read_database`, which prevents native Polars support via the existing DB-API abstraction. This change resolves the outstanding TODO identified in #50454 by implementing a dedicated reader that interacts directly with the Elasticsearch SQL API. This avoids the need to adapt the custom cursor to Polars’ internal executor interface, which would introduce unnecessary complexity and tighter coupling. **Tests** Added unit tests verifying that: * Elasticsearch SQL responses are correctly converted into Polars DataFrames with expected schema and data, including empty result sets. * Pagination using cursors is handled correctly, including aggregation of multi-page results and support for `max_rows` truncation. * Cursor cleanup is performed using the last non-null cursor and is skipped when no cursor is returned by Elasticsearch. * The hook correctly delegates to the custom reader when `df_type="polars"` is requested. **Documentation** Updated the `_get_polars_df` docstring to describe the use of a custom reader and explain why `polars.read_database` is not used. **Backwards Compatibility** This change adds support for `df_type="polars"` in `ElasticsearchSQLHook` and does not modify existing behavior for other data frame types. No breaking changes are introduced. -- 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]
