andygrove commented on code in PR #2832:
URL:
https://github.com/apache/datafusion-ballista/pull/2832#discussion_r4524565476
##########
python/python/tests/test_client.py:
##########
@@ -103,17 +104,43 @@ def test_set_ballista_config_reaches_scheduler(address,
pull):
assert status.job_name == "set_reaches_scheduler"
-def test_session_config_configures_default_client(address):
- ctx = BallistaPythonSessionContext(address,
config={"ballista.client.pull": False})
+def test_session_context_without_client_uses_native_client(address):
+ client = RecordingClient(address)
+ ctx = BallistaSessionContext(address)
+ ctx.register_csv("t", TEST_CSV)
+
+ assert ctx.sql("SELECT count(*) AS c FROM t").to_pydict() == {"c": [5]}
+ assert client.jobs == []
Review Comment:
Small thing here. The `RecordingClient` is created but never handed to the
context, so `client.jobs == []` holds whatever `BallistaSessionContext` does. I
think `test_rust_client_from_config` with the `python_jobs` fixture already
covers this properly. Maybe drop this one, or switch it to `python_jobs`?
##########
python/python/ballista/extension.py:
##########
@@ -289,10 +313,60 @@ def BallistaSessionContext(
... "df://localhost:50050",
... config=BallistaSessionConfig().with_target_partitions(32),
... )
+
+ Talk to the cluster through a Python client, here one that sends an
+ authentication header with every scheduler call:
+
+ >>> from ballista import BallistaClient
+ >>> class AuthClient(BallistaClient):
+ ... def call_metadata(self):
+ ... return super().call_metadata() + [("authorization",
"token")]
+ >>> ctx =
BallistaSessionContext(client=AuthClient("df://localhost:50050"))
+
+ Use the Python client, configured from ``config``:
+
+ >>> ctx = BallistaSessionContext(
+ ... "df://localhost:50050",
+ ... config={
+ ... "ballista.python.client": "python",
+ ... "ballista.client.use_tls": True,
+ ... },
+ ... )
"""
if cluster_config is not None:
config = _merge_cluster_config(config, cluster_config)
+ if not isinstance(config, BallistaSessionConfig):
+ config = BallistaSessionConfig(config)
+
+ if client is not None and config._python_client == "rust":
+ raise ValueError(
+ f"`client` can not be used with `{PYTHON_CLIENT_KEY}` set to
`rust`"
+ )
+
+ if client is None and config._python_client != "python":
+ return BallistaRustSessionContext(address, config, runtime)
+ # imported on use, the Python client needs the `python-client` extra
+ from .client import BallistaClient, BallistaPythonSessionContext
+
+ if client is None:
+ client = BallistaClient(address, config=config._ballista_overrides())
+ return BallistaPythonSessionContext(client, config, runtime)
Review Comment:
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.
##########
python/python/ballista/extension.py:
##########
@@ -321,6 +395,12 @@ def BallistaSessionContext(
}
+# Python only option picking the cluster client, see `BallistaSessionContext`.
+# It is kept out of `BallistaConfigExtension`, so it never reaches Rust.
+PYTHON_CLIENT_KEY = "ballista.python.client"
Review Comment:
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.
##########
python/python/ballista/extension.py:
##########
@@ -552,6 +642,17 @@ def with_grpc_client_http2_keepalive_interval_seconds(
# client
+ def with_python_client(self, kind: str) -> "BallistaSessionConfig":
Review Comment:
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.
##########
docs/source/upgrading/56.0.0.md:
##########
@@ -233,25 +233,45 @@ statuses still set `executor_meta` on every location, so
clients that read its
methods. They are functions in `ballista.jupyter`.
- Ballista settings are registered as a DataFusion config extension, exposed as
`BallistaConfigExtension`.
-- A Ballista client implemented in Python, `BallistaClient` and
- `BallistaPythonSessionContext`, is added. Its dependencies are not installed
- by default, they come with the `python-client` extra.
+- A Ballista client implemented in Python, `BallistaClient`, is added. Pass
+ it to `BallistaSessionContext(client=...)`, or set `ballista.python.client`
+ to `python` in `config`, to use it instead of the native Rust client. Its
dependencies are not installed by default, they come with
Review Comment:
Nit: this bullet has one long line that wasn't rewrapped. Prettier is fine
with it, so purely cosmetic.
##########
python/python/ballista/extension.py:
##########
@@ -552,6 +642,17 @@ def with_grpc_client_http2_keepalive_interval_seconds(
# client
+ def with_python_client(self, kind: str) -> "BallistaSessionConfig":
+ """
+ Talk to the cluster through the ``"python"`` ``BallistaClient`` or the
+ native ``"rust"`` client (default), see ``BallistaSessionContext``.
+ """
+ return self.set(PYTHON_CLIENT_KEY, kind)
+
+ def _ballista_overrides(self) -> Dict[str, str]:
Review Comment:
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.
--
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]
