PyArrow and Pandas#

Returning a Pandas DataFrame#

PyStarburst can export any DataFrame to a pandas DataFrame via to_pandas().

df = session.table("tpch.sf1.orders")
pandas_df = df.filter(df.orderstatus == "O").to_pandas()

Returning a PyArrow batch reader or table#

A PyStarburst DataFrame can be exported to a pyarrow.RecordBatchReader via to_arrow_batches().

df = session.table("tpch.sf1.orders")
batch_reader = df.filter(df.orderstatus == "O").to_arrow_batches()

The method to_arrow_table() exports a PyStarburst DataFrame to a pyarrow.Table.

table = df.to_arrow_table()
# is equivalent to
table = df.to_arrow_batches().read_all()

Arrow-accelerated conversion#

When the server has Arrow spooling protocol enabled, to_pandas(), to_arrow_batches() and to_arrow_table() fetch data as Apache Arrow Columnar Format segments and decode them in parallel, which is up to 7× faster than using the direct protocol.

Prerequisites#

Install the pyarrow extra to enable to_arrow_batches() and to_arrow_table(), or pandas extra to enable all three:
```
pip install "pystarburst[pyarrow]"
```
Starburst Enterprise server must have Arrow spooling enabled:

Configure support for the spooling protocol on a Starburst cluster - Configuration steps

Add property to enable arrow or arrow with zstd compression:
```
protocol.spooling.encoding.arrow.enabled=true
# or
protocol.spooling.encoding.arrow+zstd.enabled=true
```
Add the following to the JVM configuration:
```
--add-opens=java.base/java.nio=ALL-UNNAMED
```

Configuration#

Pass encoding when creating the session. arrow_max_workers controls the thread pool size used for parallel IPC decoding.

import trino
from pystarburst import Session

session = Session.builder.configs({
    "host": "<host>",
    "port": "<port>",
    "http_scheme": "https",
    "auth": trino.auth.BasicAuthentication("<user>", "<password>"),
    "encoding": "arrow-preview+zstd",   # or "arrow-preview"
    "arrow_max_workers": 8,             # optional
}).create()

pandas_df = session.sql("SELECT * FROM tpch.sf1.lineitem LIMIT 2_000_000").to_pandas()
# or
arrow_table = session.sql("SELECT * FROM tpch.sf1.lineitem LIMIT 2_000_000").to_arrow_table()

You can also override arrow_max_workers for a single call:

pandas_df = df.to_pandas(arrow_max_workers=4)
# or
arrow_table = df.to_arrow_table(arrow_max_workers=4)