Pandas¶
As DataFrame¶
You can use the pandas.read_sql_query to handle the query results as a pandas.DataFrame object.
from pyathena import connect
import pandas as pd
conn = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2")
df = pd.read_sql_query("SELECT * FROM many_rows", conn)
print(df.head())
NOTE: Poor performance when using pandas.read_sql #222
The pyathena.pandas.util package also has helper methods.
from pyathena import connect
from pyathena.pandas.util import as_pandas
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2").cursor()
cursor.execute("SELECT * FROM many_rows")
df = as_pandas(cursor)
print(df.describe())
If you want to use the query results output to S3 directly, you can use PandasCursor. This cursor fetches query results faster than the default cursor. (See benchmark results.)
To SQL¶
You can use pandas.DataFrame.to_sql to write records stored in DataFrame to Amazon Athena. pandas.DataFrame.to_sql uses SQLAlchemy, so you need to install it.
import pandas as pd
from sqlalchemy import create_engine
conn_str = "awsathena+rest://:@athena.{region_name}.amazonaws.com:443/"\
"{schema_name}?s3_staging_dir={s3_staging_dir}&location={location}&compression=snappy"
engine = create_engine(conn_str.format(
region_name="us-west-2",
schema_name="YOUR_SCHEMA",
s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
location="s3://YOUR_S3_BUCKET/path/to/"))
df = pd.DataFrame({"a": [1, 2, 3, 4, 5]})
df.to_sql("YOUR_TABLE", engine, schema="YOUR_SCHEMA", index=False, if_exists="replace", method="multi")
The location of the Amazon S3 table is specified by the location parameter in the connection string.
If location is not specified, s3_staging_dir parameter will be used. The following rules apply.
s3://{location or s3_staging_dir}/{schema}/{table}/
The file format, row format, and compression settings are specified in the connection string.
The pyathena.pandas.util package also has helper methods.
import pandas as pd
from pyathena import connect
from pyathena.pandas.util import to_sql
conn = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2")
df = pd.DataFrame({"a": [1, 2, 3, 4, 5]})
to_sql(df, "YOUR_TABLE", conn, "s3://YOUR_S3_BUCKET/path/to/",
schema="YOUR_SCHEMA", index=False, if_exists="replace")
This helper method supports partitioning.
import pandas as pd
from datetime import date
from pyathena import connect
from pyathena.pandas.util import to_sql
conn = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2")
df = pd.DataFrame({
"a": [1, 2, 3, 4, 5],
"dt": [
date(2020, 1, 1), date(2020, 1, 1), date(2020, 1, 1),
date(2020, 1, 2),
date(2020, 1, 3)
],
})
to_sql(df, "YOUR_TABLE", conn, "s3://YOUR_S3_BUCKET/path/to/",
schema="YOUR_SCHEMA", partitions=["dt"])
cursor = conn.cursor()
cursor.execute("SHOW PARTITIONS YOUR_TABLE")
print(cursor.fetchall())
Conversion to Parquet and upload to S3 use ThreadPoolExecutor by default. It is also possible to use ProcessPoolExecutor.
import pandas as pd
from concurrent.futures.process import ProcessPoolExecutor
from pyathena import connect
from pyathena.pandas.util import to_sql
conn = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2")
df = pd.DataFrame({"a": [1, 2, 3, 4, 5]})
to_sql(df, "YOUR_TABLE", conn, "s3://YOUR_S3_BUCKET/path/to/",
schema="YOUR_SCHEMA", index=False, if_exists="replace",
chunksize=1, executor_class=ProcessPoolExecutor, max_workers=5)
PandasCursor¶
PandasCursor directly handles the CSV file of the query execution result output to S3. This cursor is to download the CSV file after executing the query, and then loaded into pandas.DataFrame object. Performance is better than fetching data with Cursor.
You can use the PandasCursor by specifying the cursor_class
with the connect method or connection object.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
from pyathena.connection import Connection
from pyathena.pandas.cursor import PandasCursor
cursor = Connection(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
It can also be used by specifying the cursor class when calling the connection object’s cursor method.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2").cursor(PandasCursor)
from pyathena.connection import Connection
from pyathena.pandas.cursor import PandasCursor
cursor = Connection(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2").cursor(PandasCursor)
The as_pandas method returns a pandas.DataFrame object.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
df = cursor.execute("SELECT * FROM many_rows").as_pandas()
print(df.describe())
print(df.head())
Support fetch and iterate query results.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
cursor.execute("SELECT * FROM many_rows")
print(cursor.fetchone())
print(cursor.fetchmany())
print(cursor.fetchall())
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
cursor.execute("SELECT * FROM many_rows")
for row in cursor:
print(row)
The DATE and TIMESTAMP of Athena’s data type are returned as pandas.Timestamp type.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
cursor.execute("SELECT col_timestamp FROM one_row_complex")
print(type(cursor.fetchone()[0])) # <class 'pandas._libs.tslibs.timestamps.Timestamp'>
Execution information of the query can also be retrieved.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
cursor.execute("SELECT * FROM many_rows")
print(cursor.state)
print(cursor.state_change_reason)
print(cursor.completion_date_time)
print(cursor.submission_date_time)
print(cursor.data_scanned_in_bytes)
print(cursor.engine_execution_time_in_millis)
print(cursor.query_queue_time_in_millis)
print(cursor.total_execution_time_in_millis)
print(cursor.query_planning_time_in_millis)
print(cursor.service_processing_time_in_millis)
print(cursor.output_location)
If you want to customize the pandas.Dataframe object dtypes and converters, create a converter class like this:
from pyathena.converter import Converter
class CustomPandasTypeConverter(Converter):
def __init__(self):
super().__init__(
mappings=None,
types={
"boolean": object,
"tinyint": float,
"smallint": float,
"integer": float,
"bigint": float,
"float": float,
"real": float,
"double": float,
"decimal": float,
"char": str,
"varchar": str,
"array": str,
"map": str,
"row": str,
"varbinary": str,
"json": str,
}
)
def convert(self, type_, value):
# Not used in PandasCursor.
pass
Specify the combination of converter functions in the mappings argument and the dtypes combination in the types argument.
Then you simply specify an instance of this class in the convertes argument when creating a connection or cursor.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2").cursor(PandasCursor, converter=CustomPandasTypeConverter())
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
converter=CustomPandasTypeConverter()).cursor(PandasCursor)
If the unload option is enabled, the Parquet file itself has a schema, so the conversion is done to the dtypes according to that schema,
and the mappings and types settings of the Converter class are not used.
If you want to change the NaN behavior of pandas.Dataframe,
you can do so by using the keep_default_na, na_values and quoting arguments of the cursor object’s execute method.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
df = cursor.execute("SELECT * FROM many_rows",
keep_default_na=False,
na_values=[""]).as_pandas()
NOTE: PandasCursor handles the CSV file on memory. Pay attention to the memory capacity.
Chunksize options¶
The Pandas cursor can read the CSV file for each specified number of rows by using the chunksize option. This option should reduce memory usage.
The chunksize option can be enabled by specifying an integer value in the cursor_kwargs argument of the connect method or as an argument to the cursor method.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor,
cursor_kwargs={
"chunksize": 1_000_000
}).cursor()
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor(chunksize=1_000_000)
It can also be specified in the execution method when executing the query.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
cursor.execute("SELECT * FROM many_rows", chunksize=1_000_000)
SQLAlchemy allows this option to be specified in the connection string.
awsathena+pandas://:@athena.{region_name}.amazonaws.com:443/{schema_name}?s3_staging_dir={s3_staging_dir}&chunksize=1000000...
When this option is used, the object returned by the as_pandas method is a DataFrameIterator object.
This object has exactly the same interface as the TextFileReader object and can be handled in the same way.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
df_iter = cursor.execute("SELECT * FROM many_rows", chunksize=1_000_000).as_pandas()
for df in df_iter:
print(df.describe())
print(df.head())
You can also concatenate them into a single pandas.DataFrame object using pandas.concat.
import pandas
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
df_iter = cursor.execute("SELECT * FROM many_rows", chunksize=1_000_000).as_pandas()
df = pandas.concat((df for df in df_iter), ignore_index=True)
You can use the get_chunk method to retrieve a pandas.DataFrame object for each specified number of rows.
When all rows have been read, calling the get_chunk method will raise StopIteration.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor()
df_iter = cursor.execute("SELECT * FROM many_rows LIMIT 15", chunksize=1_000_000).as_pandas()
df_iter.get_chunk(10)
df_iter.get_chunk(10)
df_iter.get_chunk(10) # raise StopIteration
Unload options¶
PandasCursor also supports the unload option, as does ArrowCursor.
See Unload options for more information.
The unload option can be enabled by specifying it in the cursor_kwargs argument of the connect method or as an argument to the cursor method.
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor,
cursor_kwargs={
"unload": True
}).cursor()
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=PandasCursor).cursor(unload=True)
SQLAlchemy allows this option to be specified in the connection string.
awsathena+pandas://:@athena.{region_name}.amazonaws.com:443/{schema_name}?s3_staging_dir={s3_staging_dir}&unload=true...
AsyncPandasCursor¶
AsyncPandasCursor is an AsyncCursor that can handle pandas.DataFrame object. This cursor directly handles the CSV of query results output to S3 in the same way as PandasCursor.
You can use the AsyncPandasCursor by specifying the cursor_class
with the connect method or connection object.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
from pyathena.connection import Connection
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = Connection(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
It can also be used by specifying the cursor class when calling the connection object’s cursor method.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2").cursor(AsyncPandasCursor)
from pyathena.connection import Connection
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = Connection(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2").cursor(AsyncPandasCursor)
The default number of workers is 5 or cpu number * 5. If you want to change the number of workers you can specify like the following.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor(max_workers=10)
The execute method of the AsyncPandasCursor returns the tuple of the query ID and the future object.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
query_id, future = cursor.execute("SELECT * FROM many_rows")
The return value of the future object is an AthenaPandasResultSet object.
This object has an interface similar to AthenaResultSetObject.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
query_id, future = cursor.execute("SELECT * FROM many_rows")
result_set = future.result()
print(result_set.state)
print(result_set.state_change_reason)
print(result_set.completion_date_time)
print(result_set.submission_date_time)
print(result_set.data_scanned_in_bytes)
print(result_set.engine_execution_time_in_millis)
print(result_set.query_queue_time_in_millis)
print(result_set.total_execution_time_in_millis)
print(result_set.query_planning_time_in_millis)
print(result_set.service_processing_time_in_millis)
print(result_set.output_location)
print(result_set.description)
for row in result_set:
print(row)
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
query_id, future = cursor.execute("SELECT * FROM many_rows")
result_set = future.result()
print(result_set.fetchall())
This object also has an as_pandas method that returns a pandas.DataFrame object similar to the PandasCursor.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
query_id, future = cursor.execute("SELECT * FROM many_rows")
result_set = future.result()
df = result_set.as_pandas()
print(df.describe())
print(df.head())
The DATE and TIMESTAMP of Athena’s data type are returned as pandas.Timestamp type.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
query_id, future = cursor.execute("SELECT col_timestamp FROM one_row_complex")
result_set = future.result()
print(type(result_set.fetchone()[0])) # <class 'pandas._libs.tslibs.timestamps.Timestamp'>
As with AsyncPandasCursor, you need a query ID to cancel a query.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor()
query_id, future = cursor.execute("SELECT * FROM many_rows")
cursor.cancel(query_id)
As with AsyncPandasCursor, the unload option is also available.
from pyathena import connect
from pyathena.pandas.async_cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor,
cursor_kwargs={
"unload": True
}).cursor()
from pyathena import connect
from pyathena.pandas.cursor import AsyncPandasCursor
cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
region_name="us-west-2",
cursor_class=AsyncPandasCursor).cursor(unload=True)