Skip to content

ci: run Python client tests against a cluster built from the branch - #2374

Draft
andygrove wants to merge 2 commits into
mainfrom
ci/python-client-integration-tests
Draft

ci: run Python client tests against a cluster built from the branch#2374
andygrove wants to merge 2 commits into
mainfrom
ci/python-client-integration-tests

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #2372.

Rationale for this change

The Python client is effectively untested against the code we change. There is a pytest job, but two things stop it covering Rust changes:

  1. It almost never runs. .github/workflows/build.yml ("Python Release Build") triggers on pull_request with paths: ["python/**"]. Any PR touching ballista/core, ballista/scheduler or ballista/executor skips it.
  2. When it does run, it does not exercise the branch. Its tests call setup_test_cluster(), which starts the scheduler and executor in-process from the crates.io release that python/Cargo.toml pins. The code under review is never loaded.

Note that simply adding ballista/** to the existing job's path filter fixes neither problem. It would spend CI minutes and still test the released crates against themselves.

This is not hypothetical. #2367 is a report of a query failing with Failed to open partition file at ".../data.arrow": NotFound after the job reported Completed, and it comes down to a client and a cluster disagreeing about shuffle file naming. Building this workflow reproduced it on the first run.

The client cannot be moved forward to close the gap. pyballista re-exports datafusion-python types (datafusion_python::dataframe::PyDataFrame and friends in python/src/lib.rs), so datafusion-python, datafusion and ballista-core must all link the same datafusion. The workspace is on the 55.0.0-rc3 git tag while crates.io datafusion-python is still at 54.0.0, so the bindings cannot move until there is a matching datafusion-python release. In practice the Python client normally lags the cluster by a DataFusion release cycle, which means client/cluster skew is the steady state for Python users rather than an edge case, and it is worth testing on purpose.

So this PR keeps the client on the pinned release and swaps the cluster for one built from the branch.

What changes are included in this PR?

.github/workflows/python-integration.yml (new). Builds ballista-scheduler and ballista-executor from the working tree, starts them, and runs the existing pytest suite against them with the client built the same way the release job builds it. Triggers on ballista/** as well as python/**.

A few details worth calling out for review:

  • It is a separate workflow rather than an extra job in build.yml, because that workflow's path filter also gates the Mac, Windows and manylinux wheel jobs. Widening it there would build wheels on every Rust PR.
  • The cluster is a debug build. The test data is a handful of rows, so runtime speed is irrelevant and a release build would roughly double an already long job.
  • The scheduler and executor are started from python/ so the relative paths the tests register (testdata/test.parquet and friends) resolve executor-side as well as client-side.
  • Cluster logs are dumped on failure, and startup polls for executor registration rather than sleeping a fixed amount.

python/python/ballista/__init__.py. setup_test_cluster() now honours BALLISTA_TEST_SCHEDULER=host:port and returns that instead of starting an in-process cluster. Wrapping it here rather than in the tests means all six call sites across the three test files pick it up unchanged, and so will any test added later.

python/python/tests/test_context.py. Adds test_multi_stage_query. Every existing test is single-stage. Those still round-trip the final stage's shuffle output through Arrow Flight, but only a repartitioning query also exercises the intermediate shuffle write and read between executors. Confirmed in the executor log that this query runs stages 1, 2 and 3 while the others run stage 1 only.

Are there any user-facing changes?

No. setup_test_cluster() is a test helper and its default behaviour is unchanged. The new environment variable is opt-in.

One thing for reviewers to weigh: this job deliberately runs an older client against a current cluster. If #2370 lands and starts refusing mismatched clients at connect time, this job goes red by design and the two need to be reconciled. I have left a note on that issue.

The existing "Python Release Build" job cannot catch Rust regressions.
Its tests call setup_test_cluster(), which starts the scheduler and
executor in-process from the crates.io release that python/Cargo.toml
pins, so the code under review is never loaded. It also only triggers on
python/** changes.

Add a workflow that keeps the client on that pinned release and swaps
the cluster for one built from the branch, which is the combination
users actually run between releases. Trigger it on ballista/** too.

Add BALLISTA_TEST_SCHEDULER to setup_test_cluster() so the existing
tests can be pointed at an already-running scheduler, and add a test
whose query crosses a shuffle boundary so intermediate shuffle write
and read are covered rather than just the final stage fetch.
@andygrove
andygrove marked this pull request as ready for review August 25, 2026 22:09
@andygrove
andygrove marked this pull request as draft August 26, 2026 15:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CI: Python client tests do not run on Rust changes and do not test the working tree

1 participant