[milenkovicm]

Which issue does this PR close?

Closes #.

Rationale for this change

The Python package has two entry points that do the same job: BallistaSessionContext, which uses the native Rust client, and BallistaPythonSessionContext, which uses the Python BallistaClient. Users have to pick one, even though the only real difference is which client talks to the cluster. Planning, configuration and the returned datafusion.SessionContext are the same for both.

We haven't decided yet which implementation to keep, so this PR keeps both and leaves a single entry point, BallistaSessionContext. A client argument or the Python-only option ballista.python.client picks the client.

Passing a client alone has a gap: the client is built outside the session, so the ballista.* options in config never configure it (TLS, message size, keep-alive). With ballista.python.client = "python", BallistaSessionContext builds the client itself from address and config.

What changes are included in this PR?

• BallistaSessionContext picks the client as follows:
  • With client: BallistaClient, the Python client handles job submission, status tracking and Flight result fetching. address is ignored, because the client has its own address.
  • With ballista.python.client = "python" in config, a BallistaClient is created for address, with the ballista.* options set in config.
  • Otherwise, or with ballista.python.client = "rust", it uses the native Rust client, same as before.
  • client together with ballista.python.client = "rust" raises ValueError.
• ballista.python.client is used only in Python:
  • BallistaSessionConfig.set (and so a dict config) handles it before the ballista.* key validation, and keeps it out of BallistaConfigExtension. It never reaches Rust or the cluster, so there are no Rust changes.
  • Valid values are rust and python, case-insensitive. Any other value raises ValueError.
  • New typed setter: BallistaSessionConfig.with_python_client(kind).
  • It can't be read from a plain SessionConfig. In that case the native client is used.
• The previous BallistaSessionContext is renamed to BallistaRustSessionContext in ballista/extension.py. It and BallistaPythonSessionContext in ballista/client.py are now internal and are not exported from ballista.
• BallistaPythonSessionContext now requires client. It no longer builds a default client itself; BallistaSessionContext does that when ballista.python.client = "python".
• BallistaSessionContext handles the deprecated cluster_config argument, so it works with both clients and the deprecation warning points at the caller.
• ballista.client is imported only when the Python client is used, so import ballista and the native path still work without the python-client extra.
• Tests:
  • The ctx fixture in test_context.py runs every context test against both clients through BallistaSessionContext.
  • test_client.py uses the unified context.
  • New tests check that:
    • the native path is used when no client is given;
    • address is ignored when a client is given;
    • cluster_config works together with a client.
  • New tests for ballista.python.client:
    • Selecting the Python client from a dict and from with_python_client runs queries through BallistaClient.
    • The config's ballista.* options reach the client (pull and push job tracking).
    • The default and "rust" use the native client.
    • client together with "rust" raises.
    • The key is kept out of the built config, and invalid values are rejected.
• Docs:
  • the 56.0.0 upgrade guide;
  • the extras table in the Python quickstart;
  • the README (new "Python Cluster Client" section, covering both forms);
  • examples/python_client.py;
  • the comment on the pyproject.toml extra.

Are there any user-facing changes?

Yes, in the Python API:

• BallistaSessionContext selects the client and replaces BallistaPythonSessionContext(client=...):

  from ballista import BallistaClient, BallistaSessionContext
  
  # Python client configured from the session config
  ctx = BallistaSessionContext(
      "df://localhost:50050",
      config={"ballista.python.client_impl": "python", "ballista.client.use_tls": True},
  )
  
  # or a custom client, e.g. one that adds authentication metadata
  ctx = BallistaSessionContext(client=BallistaClient("df://localhost:50050"))

• BallistaPythonSessionContext is no longer exported. It was added after 55.0.0 and only appears in the 56.0.0 upgrade guide, which hasn't been released, so no released API breaks. The 56.0.0 upgrade guide is updated to match.

• Option precedence with ballista.python.client = "python":

  • ballista.* options set in config become the client's own config, so they take precedence over later SET statements in that session.
  • Options left out of config can still be changed with SET.
  • This matches passing the same options to BallistaClient(config=...).

• When a client is passed, the ballista.* options in config don't configure it. They still reach the scheduler with each query. Client-level options, such as ballista.client.pull, go in the BallistaClient config.

• Existing BallistaSessionContext calls without client or ballista.python.client behave exactly as before.

[andygrove]

Thanks @milenkovicm, this is a nice cleanup. One entry point that mirrors DataFusion's SessionContext is much easier to explain than two. I traced the dispatch for the native default, a passed client, ballista.python.client = "python", client plus "rust", and cluster_config with both clients, and it all looks right. The lazy import of .client keeps import ballista working without the extra too. A few small suggestions inline, none of them blocking.

[andygrove]

When someone passes client= and also puts ballista.client.pull or ballista.client.use_tls in config, those keys quietly don't apply to the client. The docs say so, but it's an easy trap. Would a UserWarning for ballista.client.* keys in that case be worth it? Happy to leave it documented if you prefer.

[andygrove]

Question on the key name. ballista.python.client sits in the ballista.* namespace but isn't a BallistaConfigExtension key. So SET ballista.python.client = 'python' fails with an unknown key error from Rust, and a future Rust key in the same namespace could clash. Also ballista.python.client = rust reads a little oddly. Have you considered something like ballista.python.client_impl? Or a short note in the docs that it can't be changed with SET? Not blocking either way.

[andygrove]

Nit: with_python_client("rust") reads a bit like it enables the Python client. Something like with_cluster_client("rust" | "python") might read more naturally. Totally optional.

[andygrove]

Nit: _ballista_overrides sits in the # client group of public with_* setters. Moving it next to build() would keep the setter list easier to scan.

[andygrove]

Nit: this bullet has one long line that wasn't rewrapped. Prettier is fine with it, so purely cosmetic.

[milenkovicm]

thanks @andygrove i believe i have addressed all your comments in 0d64002