milenkovicm opened a new pull request, #2754:
URL: https://github.com/apache/datafusion-ballista/pull/2754
# 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.
- 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](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
