PyStarburst DataFrame API

PyStarburst DataFrame API allows you to query and transform data in Starburst products in a data pipeline without having to download the data locally.

Documentation

See the PyStarburst API documentation and the examples repository.

Getting started

Install pystarburst

pip install pystarburst

Connect to a Starburst server

The parameters are the same connect parameters as in Trino Python Client.

from pystarburst import Session

connection_parameters = {
    "host": "localhost",
    "port": 8080,
    "user": "admin",
    "catalog": "tpch",
    "schema": "tiny"
}

session = Session.builder.configs(connection_parameters).create()

Using SQL

from pystarburst import Session

session = Session.builder.configs({ ... }).create()

session.sql("SELECT 1 as a").show()

Querying a table

from pystarburst import Session

session = Session.builder.configs({ ... }).create()

df = session.table("nation")
print(df.schema)
df.show()

Filtering a data frame

from pystarburst import Session

session = Session.builder.configs({ ... }).create()

df = session.table("nation")
df.filter(df.col("regionkey") == 0).show()

Joining data frames

from pystarburst import Session

session = Session.builder.configs({ ... }).create()

df = session.table("nation")
df.filter(df.col("regionkey") == 0).show()

Aggregation

from pystarburst import Session
from pystarburst.functions import col

session = Session.builder.configs({ ... }).create()
df = session.table("nation")
df.agg((col("regionkey"), "max"), (col("regionkey"), "avg")).show()

Arrow spooling

When configured with Arrow encoding, DataFrame methods to_arrow_batches(), to_arrow_table() and to_pandas() use Arrow IPC spooling with parallel segment decoding for significantly faster transfer of large result sets.

pip install pystarburst[pyarrow]

from pystarburst import Session

session = Session.builder.configs({
    ...
    "encoding": "arrow-preview+zstd",
}).create()

pandas_df = session.sql("SELECT * FROM nation").to_pandas()
# or
arrow_reader = session.sql("SELECT * FROM nation").to_arrow_batches()
# or
arrow_table = session.sql("SELECT * FROM nation").to_arrow_table()

Of the three methods: to_arrow_batches(), to_arrow_table() and to_pandas(), to_arrow_batches() is the most memory efficient, as it returns pyarrow.RecordBatchReader that can iterate over record batches without materializing the entire result set in memory.

Arrow encoding is used only for those three methods. All other operations (collect(), show(), etc.) use the default encoding.