Ballista

Ballista support for DataFusion Python: run DataFusion Python queries on a distributed Ballista cluster.

pip install ballista

User documentation is in the Python client user guide.

This project is tracked under its own Cargo.toml and is intentionally not part of the default Cargo workspace so that it doesn't cause overhead for maintainers of the main Ballista codebase. It is released separately from the Rust crates, built against the ballista crates published on crates.io, once datafusion-python has published a matching release. See the Python client release process.

Creating a SessionContext

Important

Ballista supports the DataFusion Python API natively: BallistaSessionContext returns a standard datafusion.SessionContext and queries created from it are executed on the Ballista cluster.

Some features are still being worked on, such as support for custom extension codecs. See the known limitations below.

Progress has been tracked at #2618

Creates a new context which connects to a Ballista scheduler process.

from datafusion import col, lit
from datafusion import DataFrame
# we do not need datafusion context
# it will be replaced by BallistaSessionContext
# from datafusion import SessionContext
from ballista import BallistaSessionContext

# Change from:
#
# ctx = SessionContext()
#
# to: 

ctx = BallistaSessionContext("df://localhost:50050")

# all other classes and functions are from
# the datafusion module
ctx.sql("create external table t stored as parquet location './testdata/test.parquet'")
df : DataFrame = ctx.sql("select * from t limit 5")

df.show()

Session and Ballista settings are set with the config argument, either as a dictionary or with BallistaSessionConfig. They are applied to the local session and propagated to the cluster:

from ballista import BallistaSessionConfig

ctx = BallistaSessionContext(
    "df://localhost:50050",
    config={"datafusion.execution.target_partitions": 256},
)

ctx = BallistaSessionContext(
    "df://localhost:50050",
    config=BallistaSessionConfig()
    .with_job_name("daily report")
    .with_target_partitions(256),
)

Python UDFs (scalar, aggregate and window) are serialised with the plan using cloudpickle, so they do not need to be registered on the cluster. They require a scheduler and executors started from Python (see Scheduler and Executor), running the same Python minor version as the client, with everything the UDF imports installed.

Known limitations and inefficiencies of the current approach:

• Anything requiring custom datafusion_proto::logical_plan::LogicalExtensionCodec.
• A Ballista connection will be created for each request.
• The scheduler address is only used when a query runs, so connection errors are raised by the first query, not by BallistaSessionContext.

Example DataFrame Usage

ctx = BallistaSessionContext("df://localhost:50050")
df = ctx.read_parquet('./testdata/test.parquet').filter(col("id") > lit(4)).limit(5)

pyarrow_batches = df.collect()

See the DataFusion Python documentation for more examples.

Python Cluster Client

By default the native Rust client talks to the cluster. A client written in Python, BallistaClient, can be used instead. It needs the python-client extra (pip install "ballista[python-client]"). Set ballista.python.client_impl to python to use it, configured with the ballista.* options of config (ballista.python.client_impl is not sent to the cluster and can not be changed with SET):

from ballista import BallistaSessionContext

ctx = BallistaSessionContext(
    "df://localhost:50050",
    config={"ballista.python.client_impl": "python", "ballista.client.use_tls": True},
)

Pass a BallistaClient to customize how the cluster is reached, for example to add authentication metadata:

from ballista import BallistaClient, BallistaSessionContext

class AuthClient(BallistaClient):
    def call_metadata(self):
        return super().call_metadata() + [("authorization", "token")]

ctx = BallistaSessionContext(client=AuthClient("df://localhost:50050"))

Jupyter Notebook Support

PyBallista provides first-class Jupyter notebook support with SQL magic commands and rich HTML rendering.

Install the Jupyter extras:

pip install "ballista[jupyter]"

DataFrames render as HTML tables, and the ballista.jupyter extension adds SQL magic commands (%ballista connect, %register, %sql, %%sql), plan visualization (explain_visual) and progress reporting (collect_with_progress):

%load_ext ballista.jupyter
%ballista connect df://localhost:50050
%register parquet test_data ./testdata/test.parquet
%sql SELECT COUNT(*) FROM test_data

See Jupyter Notebooks in the user guide for details.

Example Notebooks

See the examples/ directory for Jupyter notebooks demonstrating various features:

• getting_started.ipynb - Basic connection and queries
• dataframe_api.ipynb - DataFrame transformations
• distributed_queries.ipynb - Multi-stage distributed query examples

Scheduler and Executor

Scheduler and executors can be configured and started from python code. They use Python aware codecs, so unlike the Rust binaries and the Docker images they can run Python UDFs.

To start scheduler:

from ballista import BallistaScheduler

scheduler = BallistaScheduler(bind_host="0.0.0.0", bind_port=50050)

scheduler.start()
scheduler.wait_for_termination()

For executor:

from ballista import BallistaExecutor

executor = BallistaExecutor(
    bind_host="0.0.0.0",
    bind_port=50051,
    scheduler_host="localhost",
    scheduler_port=50050,
    vcores=8,
    work_dir="/tmp/ballista-executor",
)

executor.start()
executor.wait_for_termination()

All constructor arguments are optional. close() stops a started process. work_dir holds the executor's shuffle and spill files; when omitted, the executor uses a temporary directory.

For tests and notebooks, setup_test_cluster() starts a scheduler and an executor in the current process and returns the scheduler (host, port).

Trust Model for Python UDFs

A scheduler or executor started from Python decodes Python UDFs in query plans with cloudpickle.loads, which can run arbitrary Python code. Anyone who can reach the scheduler or executor ports can therefore run arbitrary code on the cluster, as the user the processes run as. This is the same trust model as PySpark.

Only expose such a cluster to trusted clients, for example by binding to localhost or a private network, or restricting access with a firewall.

Python UDFs only work on clusters started from Python. The ballista-scheduler and ballista-executor Rust binaries and the Docker images can not decode them.

Development Process

Detailed development process explanation can be found in datafusion python documentation. Improving build speed section can be relevant.

Creating Virtual Environment

pip

python3 -m venv .venv
source .venv/bin/activate
pip3 install -r requirements.txt

uv

uv sync --dev --no-install-package ballista

Developing & Building

pip

maturin develop

Note that you can also run maturin develop --release to get a release build locally.

uv

uv run --no-project maturin develop --uv

Or uv run --no-project maturin build --release --strip to get a release build.

Testing

pip

python3 -m pytest

uv

uv run --no-project pytest