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]"  # or: pip install "pystarburst[pandas]"
```
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
```

Client Configuration#

No additional configuration on the client side is needed to use to_arrow_batches(), to_arrow_table() or to_pandas() with Arrow spooling once the Starburst Enterprise server is configured as described above.

Arrow spooling may be customized using the optional arrow_encoding and arrow_max_workers configuration parameters when creating the client session:

import trino
from pystarburst import Session

session = Session.builder.configs({
    "host": "<host>",
    "port": "<port>",
    "http_scheme": "https",
    "auth": trino.auth.BasicAuthentication("<user>", "<password>"),
    # Specify segment encoding preference. If not set, the default "arrow-preview+zstd,arrow-preview" is used.
    # For example, to prefer uncompressed segments (not recommended) but accept also zstd-compressed ones:
    "arrow_encoding": "arrow-preview,arrow-preview+zstd",
    # Specify the number of threads used for segment downloading and decoding, the default is based on the number of CPU cores:
    "arrow_max_workers": 4,
}).create()

pandas_df = session.table("tpch.sf1.lineitem").limit(1_000_000).to_pandas()
# or
arrow_batches = session.table("tpch.sf1.lineitem").to_arrow_batches()

Arrow segment compression is transparent for the client code, to_arrow_batches() and other Arrow-enabled methods will return the same data type independent of whether Arrow segment compression is used or not in the transport layer.

The number of worker threads can be also overridden per call:

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