Add PolarsCursor for native Polars DataFrame support#636
Merged
laughingman7743 merged 4 commits intomasterfrom Jan 3, 2026
Merged
Add PolarsCursor for native Polars DataFrame support#636laughingman7743 merged 4 commits intomasterfrom
laughingman7743 merged 4 commits intomasterfrom
Conversation
This commit introduces PolarsCursor and AsyncPolarsCursor for native Polars DataFrame support without requiring PyArrow as a mandatory dependency. Key features: - PolarsCursor: Synchronous cursor that returns results as Polars DataFrames - AsyncPolarsCursor: Asynchronous cursor for non-blocking query execution - Support for both regular queries and UNLOAD operations - Proper type conversion for all Athena data types including decimal with precision/scale - CSV reading via fsspec with S3 support - Parquet reading via Polars native object_store Implementation details: - DefaultPolarsTypeConverter: Converts Athena types to Polars dtypes - DefaultPolarsUnloadTypeConverter: Minimal converter for UNLOAD operations - AthenaPolarsResultSet: Result set with Polars DataFrame conversion - SQLAlchemy dialect support via AthenaPolarsDialect Additional improvements: - Added get_dtype method to base Converter class for extensible type handling - Moved is_unload property to AthenaResultSet base class - Updated pandas and arrow result sets to use get_dtype method - Added as_polars() method to Arrow cursors for interoperability - Comprehensive documentation and tests Closes #436 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Add comprehensive tests for the as_polars() method in ArrowCursor: - test_as_polars: Basic single row test (with and without UNLOAD) - test_many_as_polars: Many rows (10000) test (with and without UNLOAD) - test_complex_as_polars: Complex data types test (standard mode) - test_complex_unload_as_polars: Complex data types test (UNLOAD mode) - Updated test_fetch_no_data to verify as_polars raises ProgrammingError - Updated test_executemany_fetch to verify as_polars raises ProgrammingError 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Add docstrings to AsyncPolarsCursor.get_default_converter() and arraysize property to match PolarsCursor documentation - Change default type from "varchar" to "string" in polars/util.py to align with arrow/util.py for consistency 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Add assertions to verify query_id is not None in tests where the variable was previously unused, improving test coverage. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
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
1 participant
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.
Summary
This PR introduces
PolarsCursorandAsyncPolarsCursorfor native Polars DataFrame support without requiring PyArrow as a mandatory dependency.Key Features
AthenaPolarsDialectImplementation Details
DefaultPolarsTypeConverter: Converts Athena types to Polars dtypesDefaultPolarsUnloadTypeConverter: Minimal converter for UNLOAD operationsAthenaPolarsResultSet: Result set with Polars DataFrame conversionget_dtypemethod to baseConverterclass for extensible type handlingis_unloadproperty toAthenaResultSetbase classAdditional Improvements
get_dtypemethodas_polars()method to Arrow cursors for interoperabilityas_polars()method in Arrow cursorsFuture Work (separate PR)
iter_chunks())Test plan
make chk- all quality checks passCloses #436
🤖 Generated with Claude Code