[milenkovicm]

Which issue does this PR close?

Closes #.

Rationale for this change

Today, all communication between the Python client and a Ballista cluster happens in Rust, inside DistributedQueryExec. That covers submitting the job, tracking its status, fetching results over Arrow Flight, and cancelling the job. Python users can't change how the cluster is reached (authentication headers, TLS, custom channels, result handling) without changing Rust code.

This PR adds a second client written in Python. It sits beside the existing native client; nothing is replaced. Plans are still built and encoded in Rust, because that needs Ballista's logical extension codec. A thin Rust exec node, created during planning, passes the encoded query to a Python BallistaClient, which does all gRPC and Flight work. Users pick a client by choosing a session context, which also makes it easy to compare the two.

What changes are included in this PR?

Python client (python/python/ballista/client.py)

• BallistaClient submits jobs to the scheduler, tracks them by polling (pull) or by the ExecuteQueryPush status stream (push), and fetches results with pyarrow.flight do_get. The ticket is the same Action::FetchPartition protobuf the Rust client sends.
• It also supports:
  • pipelined jobs: final-stage partitions are read concurrently, and a job failure stops the stream
  • Flight proxies, local and external
  • the scheduler version check
  • IO retries
  • cancelling the job when its results are closed before the job finishes
• Subclasses can override create_channel, call_metadata and create_flight_client.
• Options are resolved in this order: options passed to the client override the query's ballista.* settings, which override Ballista defaults.
• Batches handed to Rust must have aligned buffers. On pyarrow 21 and later, Flight reads use ensure_alignment; older versions copy each batch.

Rust node and planner (python/src/python_query_exec.rs)

• PythonDistributedQueryExec encodes the query, calls client.execute(params, schema), and streams results through the Arrow C stream interface on the blocking pool. Dropping the stream closes the job.
• PythonDistributedExplainAnalyzeExec fetches job metrics through client.job_metrics(job_id).
• PythonClientQueryPlanner uses the same planning rules as BallistaQueryPlanner.
• The new BallistaPythonClientExtension installs the planner and codecs through SessionContext.with_extensions.

Codec (python/src/codec.rs)

• Physical plans cross the datafusion-python FFI boundary as protobuf, and a Python object can't be serialized. A physical codec created with a client therefore decodes DistributedQueryExecNode / DistributedExplainAnalyzeExecNode into the Python-backed execs. A codec without a client, such as the one used by schedulers and executors started from Python, behaves as before. No .proto changes.

ballista-core

• execute_query_params is extracted from DistributedQueryExec::execute, so both clients build the same ExecuteQueryParams. Behaviour is unchanged.
• format_metrics_as_record_batch is now public.

Python package

• New BallistaPythonSessionContext. BallistaSessionContext is unchanged.
• New exports: BallistaClient, BallistaJob, BallistaJobError.
• gRPC stubs generated from ballista/core/proto are committed under ballista/proto/. They are regenerated with dev/generate_proto.sh.

Build and CI

• New runtime dependencies: grpcio>=1.85.1, protobuf>=7.35.1,<8. pyarrow moves from a test dependency to a runtime one. The minimum versions match what the generated code checks at import time. grpcio-tools is pinned in the dev group so regenerated stubs are identical, dependencies are hidden behind python-client dependency option
• A CI step fails if the committed stubs are out of date. The stubs are excluded from ruff and from RAT.

Tests

• The ctx fixture in test_context.py now runs every context test against both clients.
• New test_client.py covers:
  • custom client hooks
  • pull and push tracking
  • configuring the default client from the session config
  • reporting job failures
  • cancelling on close
  • realigning batches on older pyarrow
  • version compatibility
• Pipelined execution was checked by hand against a push-staged scheduler, in both pull and push mode, with the same results as the native client. The in-process test cluster runs pull-staged, so pipelined mode is not covered by automated tests.

Not yet done in the Python client:

• ballista.client.grpc_connect_timeout_seconds and the HTTP/2 window size options are not mapped to gRPC options.
• The job timing metrics the native client records (job_execution_time_ms, job_scheduling_in_ms, total_query_time_ms) are not reported.

Are there any user-facing changes?

Yes, all additive:

• BallistaPythonSessionContext and BallistaClient in the Python package. Existing BallistaSessionContext users see no change.
• New runtime dependencies for the Python package: grpcio, protobuf and pyarrow.
• Two new public functions in ballista-core: execution_plans::execute_query_params and execution_plans::format_metrics_as_record_batch.

There are no breaking API changes and no wire or serialization format changes.

🤖 Generated with Claude Code

[milenkovicm]

using python grpc client to communicate with ballista instead of rust client, this might give us opportunity to configure client any way we want. any opinions @andygrove

(btw, i will review code later)

[andygrove]

Thanks for this, it's a nice addition and the extraction of execute_query_params looks clean. I left two inline notes on things that would be easier to settle before this merges.

[milenkovicm]

should be merged after #2748

[andygrove]

Thanks Marko, this is a nice addition. One question on settings, plus a couple of small things inline.

Also, the description still says grpcio/protobuf/pyarrow are runtime deps, but they're in the python-client extra now.

[andygrove]

Thanks @milenkovicm