andygrove commented on code in PR #2588:
URL:
https://github.com/apache/datafusion-ballista/pull/2588#discussion_r4277112628
##########
python/python/ballista/extension.py:
##########
@@ -531,80 +200,38 @@ def __repr__(self) -> str:
return
f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
-class BallistaSessionContext(SessionContext,
metaclass=RedefiningSessionContextMeta):
- """
- A session context for connecting to and querying a Ballista cluster.
-
- This class extends DataFusion's SessionContext to work with distributed
- Ballista clusters, automatically routing query execution to the cluster
- while maintaining API compatibility with local DataFusion usage.
-
- Example:
- >>> from ballista import BallistaSessionContext
- >>> ctx = BallistaSessionContext("df://localhost:50050")
- >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
- >>> df.show()
-
- To override DataFusion / Ballista session settings on the cluster
- (e.g. the number of target partitions used by the scheduler):
-
- >>> ctx = BallistaSessionContext(
- ... "df://localhost:50050",
- ... cluster_config={"datafusion.execution.target_partitions":
"256"},
- ... )
-
- For Jupyter notebook users:
- >>> %load_ext ballista.jupyter
- >>> %ballista connect df://localhost:50050
- >>> %sql SELECT * FROM my_table
- """
-
- def __init__(
- self,
- address: str,
- config=None,
- runtime=None,
- cluster_config: Optional[Dict[str, str]] = None,
- ):
- self.cluster_config = (
- {str(k): str(v) for k, v in cluster_config.items()}
- if cluster_config is not None
- else None
- )
- # Apply overrides to the local SessionContext too: some settings
- # (e.g. datafusion.execution.listing_table_factory_infer_partitions)
- # are consulted during local table registration / planning, before
- # the plan is shipped to the scheduler. Ballista-namespaced keys
- # are not understood by the local SessionConfig and are forwarded
- # to the scheduler only.
- if self.cluster_config:
- if config is None:
- config = SessionConfig()
- for key, value in self.cluster_config.items():
- if key.startswith("ballista."):
- continue
- config = config.set(key, value)
- super().__init__(config, runtime)
- self.address = address
- self.session_id_internal = super().session_id()
-
- @property
- def session_id(self):
- return self.session_id_internal
-
- def get_tables(self) -> Optional[dict[str, List[str]]]:
- """Get tables and their respective schemas (in terms of database
schema)."""
- try:
- catalog = self.catalog()
- schema_names = list(catalog.schema_names())
- if schema_names:
- tables_info = {}
- for schema_name in schema_names:
- tables_info[schema_name] = list(
- catalog.schema(name=schema_name).table_names()
- )
- return tables_info
- except (AttributeError, NotImplementedError) as e:
- warnings.warn(f"Could not retrieve tables from catalog: {e}")
- pass
- return {}
+def BallistaSessionContext(
+ address: str,
+ config=None,
+ runtime=None,
+ cluster_config: Optional[Dict[str, str]] = None,
+):
+ cluster_config = (
+ {str(k): str(v) for k, v in cluster_config.items()}
+ if cluster_config is not None
+ else None
+ )
+ # Apply overrides to the local SessionContext too: some settings
+ # (e.g. datafusion.execution.listing_table_factory_infer_partitions)
+ # are consulted during local table registration / planning, before
+ # the plan is shipped to the scheduler. Ballista-namespaced keys
+ # are registered locally as a config extension below.
+ if config is None:
Review Comment:
I think passing a `config` here now turns off Ballista's distributed
defaults on the cluster. The defaults are only applied when `config is None`.
The FFI path then sends every entry of the session config to the scheduler, and
the scheduler applies them on top of its own defaults. So a plain
`SessionConfig()` brings back dynamic filter pushdown, physical scalar
subqueries and `prefer_hash_join = true`, which
`ballista_restricted_configuration()` turns off on purpose.
I could reproduce a failure with two 6-file Parquet tables. `SELECT count(*)
FROM a WHERE v > (SELECT avg(v) FROM b)` works with
`BallistaSessionContext(addr)`. With `config=SessionConfig()` it fails on the
executor with `ScalarSubqueryExpr can only be deserialized as part of a
surrounding ScalarSubqueryExec`, which is #1909. The join plan also flips from
`SortMergeJoinExec` to `HashJoinExec` with a `DynamicFilter`, and aggregates
repartition to 28 partitions (my core count) instead of 16. On `main` the same
queries with the same `config` keep the Ballista plans and the subquery works.
Could we apply the defaults to a user-supplied config too? That would match
what `upgrade_for_ballista()` does for the Rust client. Doing it in
`with_ballista_query_planner` would also cover anyone calling it directly, and
`cluster_config` still gets applied afterwards for explicit overrides. A test
with a custom config and a scalar subquery would catch this, since
`test_custom_session_config` only checks `count(*)`. The README line saying
`config` is used locally only needs an update too.
##########
python/src/config.rs:
##########
@@ -0,0 +1,81 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use ballista_core::config::BallistaConfig;
+use datafusion::config::{ConfigExtension, ExtensionOptions};
+use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+use std::collections::HashMap;
+
+/// Ballista configuration exposed as a DataFusion config extension, to be
+/// registered with `SessionConfig.with_extension(...)`.
+///
+/// Once registered, `ballista.*` keys are listed in
+/// `information_schema.df_settings` and can be changed with `SET`.
+#[pyclass(
+ name = "BallistaConfigExtension",
+ module = "ballista._internal_ballista",
+ skip_from_py_object
+)]
+pub(crate) struct PyBallistaConfigExtension {
+ config: BallistaConfig,
+}
+
+#[pymethods]
+impl PyBallistaConfigExtension {
+ #[new]
+ #[pyo3(signature = (settings=None))]
+ fn new(settings: Option<HashMap<String, String>>) -> PyResult<Self> {
+ let mut config = BallistaConfig::default();
+ for (key, value) in settings.unwrap_or_default() {
+ let Some(inner_key) = key
+ .strip_prefix(BallistaConfig::PREFIX)
+ .and_then(|k| k.strip_prefix('.'))
+ else {
+ return Err(PyValueError::new_err(format!(
+ "unsupported Ballista configuration key: {key}"
+ )));
+ };
+ config
+ .set(inner_key, &value)
+ .map_err(|error| PyValueError::new_err(error.to_string()))?;
+ }
+ Ok(Self { config })
+ }
+
+ fn __datafusion_extension_options__<'py>(
Review Comment:
I think `ballista.client.use_tls` and `ballista.client.pull` stop working in
this path. `DistributedQueryExec::execute` reads both from a native
`BallistaConfig` on the task context. `datafusion-ffi` rebuilds that context on
the Ballista side with only a generic `FFI_ExtensionOptions`, so both fall back
to their defaults. The docs tell clients to set `ballista.client.use_tls` for
TLS result fetches, and on `main` it works from `cluster_config`. I haven't
tried this against a TLS cluster, so please double check. Could `execute()`
fall back to `self.config` when the task context has no `BallistaConfig`? Or
could the Python side rebuild one from the `datafusion_ffi` options?
##########
python/python/ballista/jupyter.py:
##########
@@ -108,7 +110,7 @@ def decorator(func):
return decorator
-from .extension import BallistaSessionContext, DistributedDataFrame
+from .extension import BallistaSessionContext
Review Comment:
`get_tables()` went away with the old class, so `%ballista tables` (line
340) now only warns `Error listing tables: 'SessionContext' object has no
attribute 'get_tables'`. `%ballista status` on line 317 has a similar problem.
`session_id` is a method on `SessionContext`, so it prints the bound method.
What do you think about keeping `get_tables`, `explain_visual` and
`collect_with_progress` as module-level helpers that take a context or
DataFrame? None of them need a subclass. Users would keep the features, and the
README and the `getting_started`, `dataframe_api` and `distributed_queries`
notebooks would only need small changes.
##########
ballista/core/src/serde/mod.rs:
##########
@@ -1044,6 +1044,10 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
input.clone(),
)?))
}
+ PhysicalPlanType::DistributedQuery(_) =>
Err(DataFusionError::Internal(
Review Comment:
Moving the serde into `PyBallistaPhysicalExtensionCodec` makes sense given
your earlier note. Would it make sense to move the proto message there too?
Right now core's wire format has a node that its own codec rejects, and a
library user who hits this arm gets an internal error with no hint about where
the node comes from. If it needs to stay in core's oneof so the Python codec
can tell it apart from Ballista nodes, a short comment on the message and on
this arm saying only the Python client produces it would help. Keeping it out
of core would also avoid adding a `PhysicalPlanType` variant, though tags 13 to
15 went in the same way.
##########
python/src/query_planner.rs:
##########
@@ -0,0 +1,529 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::codec;
+use crate::ffi_planner_workaround::decode_file_formats_locally;
+use ballista_core::extension::SessionConfigExt;
+use ballista_core::planner::BallistaQueryPlanner;
+use datafusion::catalog::Session;
+use datafusion::common::{DataFusionError, Result};
+use datafusion::execution::config::SessionConfig;
+use datafusion::execution::context::QueryPlanner;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion_ffi::execution::FFI_TaskContextProvider;
+use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use
datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use datafusion_ffi::query_planner::FFI_QueryPlanner;
+use datafusion_proto::logical_plan::LogicalExtensionCodec;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_proto::protobuf::LogicalPlanNode;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::{PyCapsule, PyDict};
+use std::collections::HashMap;
+use std::ptr::NonNull;
+use std::sync::{Arc, OnceLock};
+use tokio::runtime::{Handle, Runtime};
+use url::Url;
+
+const DEFAULT_SCHEDULER_PORT: u16 = 50050;
+
+#[pyclass(
+ name = "BallistaExtension",
+ module = "ballista._internal_ballista",
+ skip_from_py_object
+)]
+pub(crate) struct PyBallistaExtension {
+ address: String,
+ config_overrides: Option<HashMap<String, String>>,
+}
+
+#[pymethods]
+impl PyBallistaExtension {
+ #[new]
+ #[pyo3(signature = (address, config_overrides=None))]
+ fn new(address: String, config_overrides: Option<HashMap<String, String>>)
-> Self {
+ Self {
+ address,
+ config_overrides,
+ }
+ }
+
+ /// Phase one of `SessionContext.with_extensions`: contribute the Ballista
+ /// codecs, bound against the context they will run on.
+ fn __datafusion_session_components__<'py>(
+ &self,
+ py: Python<'py>,
+ session_ctx: Bound<'py, PyAny>,
+ ) -> PyResult<Bound<'py, PyAny>> {
+ ballista_extension_components(py, session_ctx)
+ }
+
+ /// Phase two of `SessionContext.with_extensions`: contribute the Ballista
+ /// planner. `session_ctx` already carries the final codec chains.
+ ///
+ /// Ballista plans every query on the scheduler, so it replaces `fallback`
+ /// rather than delegating to it.
+ fn __datafusion_session_planner__<'py>(
+ &self,
+ py: Python<'py>,
+ session_ctx: Bound<'py, PyAny>,
+ _fallback: Bound<'py, PyAny>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ ballista_extension_planner(
+ py,
+ session_ctx,
+ self.address.clone(),
+ self.config_overrides.clone(),
+ )
+ }
+}
+
+#[pyclass(
+ name = "BallistaLogicalExtensionCodec",
+ module = "ballista._internal_ballista",
+ skip_from_py_object
+)]
+pub(crate) struct PyBallistaLogicalExtensionCodec {
+ logical_codec: Arc<dyn LogicalExtensionCodec>,
+ ffi_logical_codec: FFI_LogicalExtensionCodec,
+ // Retain the source context for the legacy low-level API. Atomic extension
+ // installation leaves this empty so the returned context owns the
provider.
+ _session_ctx: Option<Py<PyAny>>,
+}
+
+#[pymethods]
+impl PyBallistaLogicalExtensionCodec {
+ #[new]
+ fn new(session_ctx: Bound<'_, PyAny>) -> PyResult<Self> {
+ let _user_logical_codec: Arc<dyn LogicalExtensionCodec> =
+ (&ffi_logical_codec_from_python(session_ctx.clone())?).into();
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
+ Arc::new(codec::PyBallistaLogicalExtensionCodec::default());
+ let task_ctx_provider =
ffi_task_ctx_provider_from_python(session_ctx.clone())?;
+ let ffi_logical_codec = FFI_LogicalExtensionCodec::new(
+ Arc::clone(&logical_codec),
+ Some(tokio_runtime_handle()),
+ task_ctx_provider,
+ );
+
+ Ok(Self {
+ logical_codec,
+ ffi_logical_codec,
+ _session_ctx: Some(session_ctx.unbind()),
+ })
+ }
+
+ #[pyo3(signature = (_session=None))]
+ fn __datafusion_logical_extension_codec__<'py>(
+ &self,
+ py: Python<'py>,
+ _session: Option<Bound<'py, PyAny>>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ PyCapsule::new_with_value(
+ py,
+ self.ffi_logical_codec.clone(),
+ cr"datafusion_logical_extension_codec",
+ )
+ }
+}
+
+#[pyclass(
+ name = "BallistaPhysicalExtensionCodec",
+ module = "ballista._internal_ballista",
+ skip_from_py_object
+)]
+pub(crate) struct PyBallistaPhysicalExtensionCodec {
+ ffi_physical_codec: FFI_PhysicalExtensionCodec,
+ // Retain the source context for the legacy low-level API. Atomic extension
+ // installation leaves this empty so the returned context owns the
provider.
+ _session_ctx: Option<Py<PyAny>>,
+ // Retain the Python logical codec as well as its native Arc.
+ _logical_codec: Py<PyAny>,
+}
+
+impl PyBallistaPhysicalExtensionCodec {
+ fn from_logical_codec(
+ session_ctx: Bound<'_, PyAny>,
+ logical_codec: Arc<dyn LogicalExtensionCodec>,
+ logical_codec_object: Py<PyAny>,
+ ) -> PyResult<Self> {
+ let ffi_physical_codec =
+ ballista_physical_ffi_codec(session_ctx.clone(), logical_codec)?;
+
+ Ok(Self {
+ ffi_physical_codec,
+ _session_ctx: Some(session_ctx.unbind()),
+ _logical_codec: logical_codec_object,
+ })
+ }
+}
+
+#[pymethods]
+impl PyBallistaPhysicalExtensionCodec {
+ #[new]
+ fn new(
+ session_ctx: Bound<'_, PyAny>,
+ logical_codec: Bound<'_, PyAny>,
+ ) -> PyResult<Self> {
+ let native_logical_codec = {
+ let logical_codec =
+ logical_codec.extract::<PyRef<'_,
PyBallistaLogicalExtensionCodec>>()?;
+ Arc::clone(&logical_codec.logical_codec)
+ };
+ Self::from_logical_codec(
+ session_ctx,
+ native_logical_codec,
+ logical_codec.unbind(),
+ )
+ }
+
+ #[pyo3(signature = (_session=None))]
+ fn __datafusion_physical_extension_codec__<'py>(
+ &self,
+ py: Python<'py>,
+ _session: Option<Bound<'py, PyAny>>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ PyCapsule::new_with_value(
+ py,
+ self.ffi_physical_codec.clone(),
+ cr"datafusion_physical_extension_codec",
+ )
+ }
+}
+
+#[pyclass(
+ name = "BallistaQueryPlanner",
+ module = "ballista._internal_ballista",
+ skip_from_py_object
+)]
+pub(crate) struct PyBallistaQueryPlanner {
+ address: String,
+ config: ballista_core::config::BallistaConfig,
+ logical_codec: Arc<dyn LogicalExtensionCodec>,
+ fallback_logical_codec: FFI_LogicalExtensionCodec,
+ fallback_physical_codec: FFI_PhysicalExtensionCodec,
+ // Retain the source context for the legacy low-level API. Atomic extension
+ // installation leaves this empty so the returned context owns the
provider.
+ _session_ctx: Option<Py<PyAny>>,
+}
+
+#[pymethods]
+impl PyBallistaQueryPlanner {
+ #[new]
+ #[pyo3(signature = (address, session_ctx, config_overrides=None))]
+ fn new(
+ address: String,
+ session_ctx: Bound<'_, PyAny>,
+ config_overrides: Option<HashMap<String, String>>,
+ ) -> PyResult<Self> {
+ let logical_codec =
PyBallistaLogicalExtensionCodec::new(session_ctx.clone())?;
+ let fallback_physical_codec = ballista_physical_ffi_codec(
+ session_ctx.clone(),
+ Arc::clone(&logical_codec.logical_codec),
+ )?;
+ Self::from_codecs(
+ address,
+ session_ctx,
+ config_overrides,
+ logical_codec.logical_codec,
+ logical_codec.ffi_logical_codec,
+ fallback_physical_codec,
+ )
+ }
+
+ #[pyo3(signature = (_session=None))]
+ fn __datafusion_query_planner__<'py>(
+ &self,
+ py: Python<'py>,
+ _session: Option<Bound<'py, PyAny>>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ let planner: Arc<dyn QueryPlanner + Send + Sync> =
+ Arc::new(LazyBallistaQueryPlanner {
+ address: self.address.clone(),
+ config: self.config.clone(),
+ logical_codec: Arc::clone(&self.logical_codec),
+ _session_ctx: self
+ ._session_ctx
+ .as_ref()
+ .map(|session_ctx| session_ctx.clone_ref(py)),
+ });
+ let ffi = FFI_QueryPlanner::new_with_ffi_codecs(
+ planner,
+ self.fallback_logical_codec.clone(),
+ self.fallback_physical_codec.clone(),
+ );
+ let ffi = decode_file_formats_locally(ffi);
+
+ // This exact name is the ABI consumed by datafusion-python.
+ PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner")
+ }
+}
+
+impl PyBallistaQueryPlanner {
+ fn from_codecs(
+ address: String,
+ session_ctx: Bound<'_, PyAny>,
+ config_overrides: Option<HashMap<String, String>>,
+ logical_codec: Arc<dyn LogicalExtensionCodec>,
+ fallback_logical_codec: FFI_LogicalExtensionCodec,
+ fallback_physical_codec: FFI_PhysicalExtensionCodec,
+ ) -> PyResult<Self> {
+ let mut session_config = SessionConfig::new_with_ballista();
+ if let Some(overrides) = config_overrides {
+ for (key, value) in overrides {
+ if !key.starts_with("ballista.") {
+ return Err(PyValueError::new_err(format!(
+ "unsupported Ballista configuration key: {key}"
+ )));
+ }
+ session_config
+ .options_mut()
+ .set(&key, &value)
+ .map_err(|error|
PyValueError::new_err(error.to_string()))?;
+ }
+ }
+
+ Ok(Self {
+ address,
+ config: session_config.ballista_config(),
+ logical_codec,
+ fallback_logical_codec,
+ fallback_physical_codec,
+ _session_ctx: Some(session_ctx.unbind()),
+ })
+ }
+}
+
+struct LazyBallistaQueryPlanner {
+ address: String,
+ config: ballista_core::config::BallistaConfig,
+ logical_codec: Arc<dyn LogicalExtensionCodec>,
+ _session_ctx: Option<Py<PyAny>>,
+}
+
+impl std::fmt::Debug for LazyBallistaQueryPlanner {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("LazyBallistaQueryPlanner")
+ .field("address", &self.address)
+ .field("config", &self.config)
+ .finish_non_exhaustive()
+ }
+}
+
+#[async_trait::async_trait]
+impl QueryPlanner for LazyBallistaQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if matches!(logical_plan, LogicalPlan::Analyze(_)) {
Review Comment:
This takes `EXPLAIN ANALYZE` away from Python users. On `main`,
`ctx.sql("EXPLAIN ANALYZE ...")` goes to `BallistaQueryPlanner`, which wraps it
in `DistributedExplainAnalyzeExec` and pulls per-stage metrics from the
scheduler, including `output_rows_skew`. With this branch both
`ctx.sql("EXPLAIN ANALYZE ...")` and `df.explain(analyze=True)` fail with this
error, and `examples/distributed_queries.ipynb` uses the second one.
Would it work to add serde for `DistributedExplainAnalyzeExec` next to
`DistributedQueryExec` in `codec.rs`? It only carries the inner query exec, the
scheduler URL, the schema and `verbose`. If you'd rather do it as a follow-up,
could we open an issue and note the gap in the upgrade guide?
##########
python/python/ballista/__init__.py:
##########
@@ -21,19 +21,20 @@
import importlib_metadata
from ._internal_ballista import (
+ BallistaConfigExtension,
BallistaScheduler,
BallistaExecutor,
setup_test_cluster,
)
from .extension import (
BallistaSessionContext,
- DistributedDataFrame,
Review Comment:
With this import gone, `DistributedDataFrame` is still listed in `__all__` a
few lines down. So `from ballista import *` now fails with `AttributeError:
module 'ballista' has no attribute 'DistributedDataFrame'`.
--
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]
