Spark

SparkCursor

SparkCursor can run Spark applications in Athena. This cursor does not follow the DB API 2.0 (PEP 249). It does not support result set iteration.

To use this cursor, you must create a Spark enabled workgroup in Athena and use that workgroup.

You can use the SparkCursor by specifying the cursor_class with the connect method or connection object.

from pyathena import connect
from pyathena.spark.cursor import SparkCursor

with connect(region_name="us-west-2",
             work_group="YOUR_SPARK_WORKGROUP",
             cursor_class=SparkCursor).cursor() as cursor:
    ...

from pyathena.connection import Connection
from pyathena.spark.cursor import SparkCursor

with Connection(region_name="us-west-2",
                work_group="YOUR_SPARK_WORKGROUP",
                cursor_class=SparkCursor).cursor() as 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.spark.cursor import SparkCursor

with connect(region_name="us-west-2",
             work_group="YOUR_SPARK_WORKGROUP").cursor(SparkCursor) as cursor:
    ...

from pyathena.connection import Connection
from pyathena.spark.cursor import SparkCursor

with Connection(region_name="us-west-2",
                work_group="YOUR_SPARK_WORKGROUP").cursor(SparkCursor) as cursor:
    ...

This cursor allows you to send PySpark code blocks and use Spark DataFrame and SQL.

Session lifecycle

When a Spark cursor is created, it starts a new session and waits until the session is idle. If session_id is specified, the cursor uses that existing session instead, after waiting until it is idle.

The session idle timeout minutes can be specified with the session_idle_timeout_minutes argument when creating the cursor and the engine DPU and Spark properties can also be specified with the engine_configuration argument. These arguments apply only to a session started by the cursor.

from pyathena import connect
from pyathena.spark.cursor import SparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP",
               cursor_class=SparkCursor)
with conn.cursor(session_idle_timeout_minutes=60,
                 engine_configuration={
                     "CoordinatorDpuSize": 1,
                     "MaxConcurrentDpus": 20,
                     "DefaultExecutorDpuSize": 1,
                     "AdditionalConfigs": {"string": "string"},
                     "SparkProperties": {"string": "string"},
                 }) as cursor:
    cursor.execute("...")
    ...

The close method of the cursor terminates the session that the cursor started. A session specified with session_id is not terminated. The terminate_session_on_close argument changes which sessions the close method terminates:

terminate_session_on_close	Session started by the cursor	Session specified with session_id
None (default)	Terminated	Not terminated
True	Terminated	Terminated
False	Not terminated	Not terminated

Once the close method has terminated the session, calling it again does not terminate the session again. If termination fails, the close method raises OperationalError, and calling it again retries the termination.

The following example keeps the session started by the first cursor, runs another calculation on it with a second cursor, and terminates it when the second cursor is closed:

from pyathena import connect
from pyathena.spark.cursor import SparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP",
               cursor_class=SparkCursor)
with conn.cursor(terminate_session_on_close=False) as cursor:
    cursor.execute("...")
    session_id = cursor.session_id

with conn.cursor(session_id=session_id, terminate_session_on_close=True) as cursor:
    cursor.execute("...")

Closing the connection does not close its cursors or terminate their sessions. This applies to Connection and AioConnection, including their context managers. Close each Spark cursor, for example with the context manager of the cursor. A session that is left running is terminated by Athena after its idle timeout.

Spark DataFrames

The Spark DataFrames code in the sample notebook that can be enabled when creating a workgroup can be executed as follows:

import textwrap
from pyathena import connect
from pyathena.spark.cursor import SparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=SparkCursor)
with conn.cursor() as cursor:
    cursor.execute(
        textwrap.dedent(
            """
            file_name = "s3://athena-examples-us-east-1/notebooks/yellow_tripdata_2016-01.parquet"

            taxi_df = (spark.read.format("parquet")
                 .option("header", "true")
                 .option("inferSchema", "true")
                 .load(file_name))
            """
        )
    )

    cursor.execute(
        textwrap.dedent(
            """
            taxi1_df=taxi_df.groupBy("VendorID", "passenger_count").count()
            taxi1_df.show()
            """
        )
    )
    print(cursor.get_std_out())

    cursor.execute(
        textwrap.dedent(
            """
            taxi1_df.coalesce(1).write.mode('overwrite').csv("s3://YOUR_S3_BUCKET/select_taxi")
            print("Write to s3 " + "complete")
            """
        )
    )
    print(cursor.get_std_out())

The standard output and standard error of a spark application can be retrieved with the get_std_out() and get_std_error() methods in the cursor class.

Spark SQL

The Spark SQL code in the sample notebook can be executed as follows:

import textwrap
from pyathena import connect
from pyathena.spark.cursor import SparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=SparkCursor)
with conn.cursor() as cursor:
    cursor.execute(
        textwrap.dedent(
            """
            file_name = "s3://athena-examples-us-east-1/notebooks/yellow_tripdata_2016-01.parquet"

            taxi_df = (spark.read.format("parquet")
                 .option("header", "true")
                 .option("inferSchema", "true")
                 .load(file_name))
            taxi_df.createOrReplaceTempView("taxis")

            sqlDF = spark.sql("SELECT DOLocationID, sum(total_amount) as sum_total_amount FROM taxis where DOLocationID < 25 GRoup by DOLocationID ORDER BY DOLocationID")
            sqlDF.show(50)
            """
        )
    )
    print(cursor.get_std_out())

    cursor.execute(
        textwrap.dedent(
            """
            spark.sql("create database if not exists spark_demo_database")
            spark.sql("show databases").show()
            """
        )
    )
    print(cursor.get_std_out())

    cursor.execute(
        textwrap.dedent(
            """
            spark.sql("use spark_demo_database")
            taxi1_df=taxi_df.groupBy("VendorID", "passenger_count").count()
            taxi1_df.write.mode("overwrite").format("parquet").option("path","s3://YOUR_S3_BUCKET/select_taxi").saveAsTable("select_taxi_table")
            print("Create new table" + " complete")
            """
        )
    )
    print(cursor.get_std_out())

    cursor.execute(
        textwrap.dedent(
            """
            spark.sql("show tables").show()
            """
        )
    )
    print(cursor.get_std_out())

    cursor.execute(
        textwrap.dedent(
            """
            spark.sql("select * from select_taxi_table").show()
            """
        )
    )
    print(cursor.get_std_out())

Cancellation

The cancel() method sends a StopCalculationExecution request for the calculation. It does not terminate the session. Athena cancels the calculation on a best-effort basis:

• A running Spark job, such as a DataFrame action, usually stops within seconds. The calculation then ends in the CANCELED state, and the session remains usable for later calculations.

• A request sent right after the calculation starts can occasionally have no effect. The calculation then runs as if it had not been canceled.

• Python code that runs on the driver without a Spark job, such as time.sleep(), runs to completion. The calculation ends in the COMPLETED state, and the session rejects new calculations until then.

• Canceling a calculation that has already finished does not raise an error or change its state.

The execute() method raises OperationalError when the calculation ends in the CANCELED state. The following example cancels the calculation from another thread after 60 seconds:

import threading
from pyathena import connect
from pyathena.spark.cursor import SparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=SparkCursor)
with conn.cursor() as cursor:
    timer = threading.Timer(60, cursor.cancel)
    timer.start()
    try:
        cursor.execute("""print(spark.read.parquet("s3://YOUR_S3_BUCKET/large_dataset/").count())""")
    finally:
        timer.cancel()

With kill_on_interrupt enabled, which is the default, a KeyboardInterrupt while execute() waits for the calculation requests cancellation, waits until the calculation reaches a terminal state, and then propagates. The state property returns that terminal state. If the cancellation request or that wait fails, the KeyboardInterrupt propagates with the error as its cause, and the state property returns None.

A KeyboardInterrupt while execute() is still starting the calculation first waits for the StartCalculationExecution request to finish, and then cancels the calculation it started in the same way. The calculation_id property returns that calculation’s ID. If execute() has not begun the request when the interrupt is handled, the request is never sent and calculation_id is None. A second KeyboardInterrupt during this wait propagates at once without cancelling the calculation. A cancellation request sent right after a calculation starts can occasionally have no effect, so the calculation can still end in the COMPLETED state.

Unless client_request_token is passed to execute(), the cursor sends a generated ClientRequestToken with each calculation. A retried start request then returns the calculation that an earlier attempt started instead of starting another one. A token passed to execute() must be unique for each calculation: Athena returns the earlier calculation for a reused token, even when the code differs.

AsyncSparkCursor

AsyncSparkCursor is an asynchronous Spark cursor that, like AsyncCursor, returns future objects.

You can use the AsyncSparkCursor by specifying the cursor_class with the connect method or connection object.

from pyathena import connect
from pyathena.spark.async_cursor import AsyncSparkCursor

with connect(region_name="us-west-2",
             work_group="YOUR_SPARK_WORKGROUP",
             cursor_class=AsyncSparkCursor).cursor() as cursor:
    ...

from pyathena.connection import Connection
from pyathena.spark.async_cursor import AsyncSparkCursor

with Connection(region_name="us-west-2",
                work_group="YOUR_SPARK_WORKGROUP",
                cursor_class=AsyncSparkCursor).cursor() as 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.spark.async_cursor import AsyncSparkCursor

with connect(region_name="us-west-2",
             work_group="YOUR_SPARK_WORKGROUP").cursor(AsyncSparkCursor) as cursor:
    ...

from pyathena.connection import Connection
from pyathena.spark.async_cursor import AsyncSparkCursor

with Connection(region_name="us-west-2",
                work_group="YOUR_SPARK_WORKGROUP").cursor(AsyncSparkCursor) as cursor:
    ...

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.spark.async_cursor import AsyncSparkCursor

with connect(region_name="us-west-2",
             work_group="YOUR_SPARK_WORKGROUP",
             cursor_class=AsyncSparkCursor).cursor(max_workers=10) as cursor:
    ...

The execute method of the AsyncSparkCursor returns the tuple of the calculation ID and the future object.

from pyathena import connect
from pyathena.spark.async_cursor import AsyncSparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=AsyncSparkCursor)
with conn.cursor() as cursor:
    calculation_id, future = cursor.execute("""spark.sql("SELECT * FROM many_rows")""")

The return value of the future object is an AthenaCalculationExecution object.

from pyathena import connect
from pyathena.spark.async_cursor import AsyncSparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=AsyncSparkCursor)
with conn.cursor() as cursor:
    calculation_id, future = cursor.execute("""spark.sql("SELECT * FROM many_rows")""")
    calculation_execution = future.result()
    print(calculation_execution.session_id)
    print(calculation_execution.calculation_id)
    print(calculation_execution.description)
    print(calculation_execution.working_directory)
    print(calculation_execution.state)
    print(calculation_execution.state_change_reason)
    print(calculation_execution.submission_date_time)
    print(calculation_execution.completion_date_time)
    print(calculation_execution.dpu_execution_in_millis)
    print(calculation_execution.progress)
    print(calculation_execution.std_out_s3_uri)
    print(calculation_execution.std_error_s3_uri)
    print(calculation_execution.result_s3_uri)
    print(calculation_execution.result_type)

Standard output and standard error can be retrieved by passing this object to the cursor class. get_std_out() and get_std_error() return None instead of a future when the calculation has no standard output or standard error location.

from pyathena import connect
from pyathena.spark.async_cursor import AsyncSparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=AsyncSparkCursor)
with conn.cursor() as cursor:
    calculation_id, future = cursor.execute("""spark.sql("SELECT * FROM many_rows")""")
    calculation_execution = future.result()
    std_out = cursor.get_std_out(calculation_execution)
    if std_out:
        print(std_out.result())
    std_error = cursor.get_std_error(calculation_execution)
    if std_error:
        print(std_error.result())

As with AsyncCursor, you need a calculation ID to cancel a calculation. The cancel method returns a future object that returns nothing.

from pyathena import connect
from pyathena.spark.async_cursor import AsyncSparkCursor

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=AsyncSparkCursor)
with conn.cursor() as cursor:
    calculation_id, future = cursor.execute("""spark.sql("SELECT * FROM many_rows").show()""")
    cursor.cancel(calculation_id).result()
    calculation_execution = future.result()
    print(calculation_execution.state)

Cancellation follows the rules described in Cancellation. The future returned by execute() completes when the calculation reaches a terminal state, which can be COMPLETED if the calculation was not canceled.

AioSparkCursor

AioSparkCursor is a native asyncio cursor for executing PySpark code on Athena. Unlike AsyncSparkCursor which uses concurrent.futures, this cursor uses native asyncio for polling and API calls, keeping the event loop free.

Since SparkBaseCursor.__init__ performs I/O (session management), cursor creation must be wrapped in asyncio.to_thread:

import asyncio
from pyathena import aio_connect
from pyathena.aio.spark.cursor import AioSparkCursor

async with await aio_connect(work_group="YOUR_SPARK_WORKGROUP",
                          cursor_class=AioSparkCursor) as conn:
    cursor = await asyncio.to_thread(conn.cursor)
    try:
        await cursor.execute("""spark.sql("SELECT 1").show()""")
        print(await cursor.get_std_out())
    finally:
        await cursor.close()

The close method follows the rules described in Session lifecycle. The cursor supports the async context manager, which calls the close method:

import asyncio
import textwrap
from pyathena import aio_connect
from pyathena.aio.spark.cursor import AioSparkCursor

async with await aio_connect(work_group="YOUR_SPARK_WORKGROUP",
                          cursor_class=AioSparkCursor) as conn:
    cursor = await asyncio.to_thread(conn.cursor)
    async with cursor:
        await cursor.execute(
            textwrap.dedent(
                """
                file_name = "s3://athena-examples-us-east-1/notebooks/yellow_tripdata_2016-01.parquet"

                taxi_df = (spark.read.format("parquet")
                     .option("header", "true")
                     .option("inferSchema", "true")
                     .load(file_name))
                taxi1_df=taxi_df.groupBy("VendorID", "passenger_count").count()
                taxi1_df.show()
                """
            )
        )
        print(await cursor.get_std_out())
        print(await cursor.get_std_error())

To cancel a running calculation, call cancel() while execute() is waiting for it. Cancellation follows the rules described in Cancellation.

import asyncio
from pyathena import aio_connect
from pyathena.aio.spark.cursor import AioSparkCursor
from pyathena.error import OperationalError

async with await aio_connect(work_group="YOUR_SPARK_WORKGROUP",
                          cursor_class=AioSparkCursor) as conn:
    cursor = await asyncio.to_thread(conn.cursor)
    async with cursor:
        task = asyncio.create_task(cursor.execute("""spark.sql("SELECT * FROM many_rows").show()"""))
        await asyncio.sleep(60)
        await cursor.cancel()
        try:
            await task
        except OperationalError:
            print(cursor.state)

With kill_on_interrupt enabled, which is the default, cancelling the task while execute() waits for the calculation requests cancellation of the calculation, waits until it reaches a terminal state, and then raises asyncio.CancelledError. Cancelling the task while execute() is still starting the calculation first waits for the start request to finish, and then cancels the calculation it started in the same way. If the task is cancelled before execute() begins the request, the request is never sent. Cancelling the task again during this wait raises asyncio.CancelledError at once without cancelling the calculation.