PyAthena Documentation

PyAthena is a Python client for Amazon Athena that provides several key capabilities:

Core Capabilities

• DB API 2.0 Compliance: Full PEP 249 compatibility.
• SQLAlchemy Integration: Native dialect support including table reflection and ORM capabilities.
• Multiple Cursor Types: Support for Standard, Pandas, Arrow, Polars, S3FS, and Spark cursors.
• Async Support: Asynchronous query execution.

Data Type Support

PyAthena handles complex Amazon Athena data types by converting them to native Python structures:

• STRUCT/ROW: Support for complex nested data structures.
• ARRAY: Ordered collections converted to Python lists.
• MAP: Key-value structures converted to Python dictionaries.
• JSON: Seamless parsing and conversion.

Additional Capabilities

• Connection Management: Efficient pooling and configuration.
• Result Caching: Ability to reuse Athena query results.
• S3 Integration: Direct S3 data access and staging support.

The ArrowCursor is a specialized cursor designed for high-performance data retrieval. It works by downloading the CSV file produced by an Athena query execution to S3 and loading it directly into a pyarrow.Table object. This method is significantly faster than using the standard Cursor for large datasets.

You can instantiate ArrowCursor in three ways:

1. By passing cursor_class=ArrowCursor to the connect() function or Connection object.
2. By passing ArrowCursor as an argument to the .cursor() method.
3. By passing ArrowCursor and a custom converter to the .cursor() method.

from pyathena import connect
from pyathena.arrow.cursor import ArrowCursor

# Option 1: Via connect()
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
                 region_name="us-west-2",
                 cursor_class=ArrowCursor).cursor()

# Option 2: Via .cursor()
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
                 region_name="us-west-2").cursor(ArrowCursor)

To avoid memory exhaustion when loading large datasets, use the chunksize option. When chunksize is specified, .as_pandas() returns a PandasDataFrameIterator instead of a single DataFrame. This iterator behaves similarly to pandas' TextFileReader.

Chunking Methods:

• iter_chunks(): A convenient method to iterate over chunks of the dataset. Memory can be freed manually after processing each chunk.
• get_chunk(n): Retrieves the next n rows as a DataFrame. Raises StopIteration when no more rows are available.
• as_pandas() on the iterator: Collects all chunks into a single large DataFrame (equivalent to pandas.concat).

from pyathena import connect
from pyathena.pandas.cursor import PandasCursor

cursor = connect(s3_staging_dir="s3://...", region_name="us-west-2", cursor_class=PandasCursor).cursor()

# Method 1: Using iter_chunks()
cursor.execute("SELECT * FROM large_table", chunksize=50_000)
for chunk in cursor.iter_chunks():
    # Process chunk
    del chunk  # Free memory

# Method 2: Using get_chunk()
df_iter = cursor.execute("SELECT * FROM table LIMIT 15", chunksize=1_000_000).as_pandas()
chunk = df_iter.get_chunk(10)

You can control how Athena types are mapped to Polars types by subclassing pyathena.converter.Converter and providing a mappings dictionary. This is useful for ensuring specific precision or type handling.

Note: If the unload option is enabled, the conversion uses the schema from the Parquet file itself, and the types setting in your custom converter will be ignored.

import polars as pl
from pyathena.converter import Converter

class CustomPolarsTypeConverter(Converter):
    def __init__(self):
        super().__init__(
            mappings=None,
            types={
                "integer": pl.Int32,
                "bigint": pl.Int64,
                "float": pl.Float32,
                "double": pl.Float64,
                # ... other mappings
            }
        )

    def convert(self, type_, value):
        pass # Not used in PolarsCursor

# Usage
cursor = connect(s3_staging_dir="s3://...",
                 region_name="us-west-2",
                 converter=CustomPolarsTypeConverter()).cursor(PolarsCursor)