ElasticsearchSQLHook: add chunked Polars DataFrame support via custom SQL reader - #68411
Merged
shahar1 merged 3 commits intoJun 29, 2026
Conversation
SameerMesiah97
force-pushed
the
66290-ElasticSearchSQLHook-Get-Polars-DF-Chunks
branch
from
June 11, 2026 22:55
87b9026 to
68aa072
Compare
SameerMesiah97
marked this pull request as ready for review
June 12, 2026 00:56
justinpakzad
left a comment
justinpakzad
left a comment
Contributor
There was a problem hiding this comment.
Nice PR. Left a couple of small non-blocking comments. Curious to hear your thoughts.
added 2 commits
June 14, 2026 19:26
Implement _get_polars_df_by_chunks using Elasticsearch SQL cursor-based pagination and add unit tests covering chunking behavior, pagination across cursor pages, and cursor cleanup.
exception and remove redundant test setup identified during review.
SameerMesiah97
force-pushed
the
66290-ElasticSearchSQLHook-Get-Polars-DF-Chunks
branch
from
June 14, 2026 18:43
68aa072 to
4f97981
Compare
Contributor
Author
|
I have addressed your feedback. |
justinpakzad
approved these changes
Jun 14, 2026
potiuk
added
the
ready for maintainer review
shahar1
approved these changes
Jun 25, 2026
shahar1
left a comment
shahar1
left a comment
Contributor
There was a problem hiding this comment.
Nice feature!
Small comment, after fixing it's good to merge from my prespective.
shahar1
approved these changes
Jun 25, 2026
shahar1
left a comment
shahar1
left a comment
Contributor
There was a problem hiding this comment.
Good to merge once CI is green :)
Contributor
Author
|
Can this be merged? |
shahar1
approved these changes
Jun 29, 2026
Contributor
Done |
karenbraganz
pushed a commit
to karenbraganz/airflow
that referenced
this pull request
Jun 30, 2026
… SQL reader (apache#68411) Implement _get_polars_df_by_chunks using Elasticsearch SQL cursor-based pagination and add unit tests covering chunking behavior, pagination across cursor pages, and cursor cleanup.
75 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR is a follow-up to #66220, which added support for returning Elasticsearch SQL query results as a Polars DataFrame via a custom cursor-based reader.
This change extends that implementation to support
get_df_by_chunks(df_type="polars")by introducing a chunked reader that uses Elasticsearch SQL cursor-based pagination to yield Polars DataFrames incrementally.Rationale
The Common SQL implementation of
get_df_by_chunks(df_type="polars")relies onpolars.read_database, which is not compatible with Elasticsearch's SQL API and custom cursor implementation.To provide feature parity with the existing Polars DataFrame support added in #66220, this change introduces an Elasticsearch-specific chunked reader that interacts directly with the Elasticsearch SQL API and yields Polars DataFrames of the requested chunk size.
Tests
Added unit tests verifying that:
get_df_by_chunks(df_type="polars")is requested.Documentation
Added docstrings for the chunked reader describing the use of Elasticsearch SQL cursor-based pagination and chunked Polars DataFrame generation.
Backwards Compatibility
This change adds support for
get_df_by_chunks(df_type="polars")in ElasticsearchSQLHook and does not modify existing behavior. No breaking changes are introduced.Closes: #66290
Was generative AI tooling used to co-author this PR?
Generated-by: [GPT 5.5] following the guidelines