This is an automated email from the ASF dual-hosted git repository. timsaucer pushed a commit to branch feat/ffi-query-planner-55 in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
commit 0a07776f6b2c3194d743d6c6b3e9e8ed725b1138 Author: Tim Saucer <[email protected]> AuthorDate: Fri Aug 7 16:07:37 2026 -0400 feat(python): use atomic Ballista session extensions AI Disclosure: This code was written in part by an AI agent.: --- python/Cargo.lock | 4 +- python/Cargo.toml | 2 +- python/examples/getting_started.ipynb | 16 +--- python/pyproject.toml | 4 +- python/python/ballista/__init__.py | 4 + python/python/tests/test_context.py | 32 ++++++++ python/requirements.txt | 2 +- python/src/lib.rs | 1 + python/src/query_planner.rs | 145 +++++++++++++++++++++++----------- python/uv.lock | 4 +- 10 files changed, 150 insertions(+), 64 deletions(-) diff --git a/python/Cargo.lock b/python/Cargo.lock index 1a2121812..49a856bd5 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -1681,7 +1681,7 @@ dependencies = [ [[package]] name = "datafusion-python" version = "54.0.0" -source = "git+https://github.com/apache/datafusion-python.git?rev=55f43b5e231da5591e7bbd632dac85f304a22cca#55f43b5e231da5591e7bbd632dac85f304a22cca"; +source = "git+https://github.com/apache/datafusion-python.git?rev=f1526210338d75ef105246e5c948e0e4a4a0f5ce#f1526210338d75ef105246e5c948e0e4a4a0f5ce"; dependencies = [ "arrow", "arrow-select", @@ -1713,7 +1713,7 @@ dependencies = [ [[package]] name = "datafusion-python-util" version = "54.0.0" -source = "git+https://github.com/apache/datafusion-python.git?rev=55f43b5e231da5591e7bbd632dac85f304a22cca#55f43b5e231da5591e7bbd632dac85f304a22cca"; +source = "git+https://github.com/apache/datafusion-python.git?rev=f1526210338d75ef105246e5c948e0e4a4a0f5ce#f1526210338d75ef105246e5c948e0e4a4a0f5ce"; dependencies = [ "arrow", "datafusion", diff --git a/python/Cargo.toml b/python/Cargo.toml index 10527a71b..3f45c5364 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -38,7 +38,7 @@ ballista-scheduler = { path = "../ballista/scheduler", version = "54.0.0", defau datafusion = { git = "https://github.com/apache/datafusion.git";, rev = "e08aed1e5de41dcf81d529140dae07723b942a5e", features = ["avro"] } datafusion-ffi = { git = "https://github.com/apache/datafusion.git";, rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } datafusion-proto = { git = "https://github.com/apache/datafusion.git";, rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } -datafusion-python = { git = "https://github.com/apache/datafusion-python.git";, rev = "55f43b5e231da5591e7bbd632dac85f304a22cca", default-features = false } +datafusion-python = { git = "https://github.com/apache/datafusion-python.git";, rev = "f1526210338d75ef105246e5c948e0e4a4a0f5ce", default-features = false } log = { version = "0.4" } pyo3 = { version = "0.29", features = ["extension-module", "abi3", "abi3-py310"] } pyo3-log = "0.13" diff --git a/python/examples/getting_started.ipynb b/python/examples/getting_started.ipynb index b2b938ba5..60eb5c3dd 100644 --- a/python/examples/getting_started.ipynb +++ b/python/examples/getting_started.ipynb @@ -58,11 +58,9 @@ "outputs": [], "source": [ "import ballista\n", - "from ballista._internal_ballista import ballista_datafusion_config_defaults\n", "from ballista import (\n", - " BallistaQueryPlanner,\n", - " BallistaLogicalExtensionCodec,\n", - " BallistaPhysicalExtensionCodec,\n", + " BallistaExtension,\n", + " ballista_datafusion_config_defaults,\n", " setup_test_cluster,\n", ")\n", "from datafusion import SessionConfig, SessionContext\n", @@ -87,14 +85,8 @@ "address = f\"df://{host}:{port}\"\n", "\n", "ballista_config = ballista_datafusion_config_defaults()\n", - "ctx = SessionContext(SessionConfig(ballista_config))\n", - "logical_codec = BallistaLogicalExtensionCodec(ctx)\n", - "physical_codec = BallistaPhysicalExtensionCodec(ctx, logical_codec)\n", - "\n", - "ctx = (\n", - " ctx.with_logical_extension_codec(logical_codec)\n", - " .with_physical_extension_codec(physical_codec)\n", - " .with_query_planner(BallistaQueryPlanner(address, ctx))\n", + "ctx = SessionContext(SessionConfig(ballista_config)).with_extensions(\n", + " BallistaExtension(address)\n", ")\n", "\n", "print(f\"Connected to Ballista at {host}:{port}\")\n", diff --git a/python/pyproject.toml b/python/pyproject.toml index 4e85fbb07..dbc82006c 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -41,8 +41,8 @@ classifiers = [ ] dependencies = [ "pyarrow>=23.0.1", - # Temporary PR #1672 pin. Use its released version before publishing. - "datafusion @ git+https://github.com/apache/datafusion-python.git@55f43b5e231da5591e7bbd632dac85f304a22cca";, + # Temporary PR #1679 pin. Use its released version before publishing. + "datafusion @ git+https://github.com/apache/datafusion-python.git@f1526210338d75ef105246e5c948e0e4a4a0f5ce";, "typing-extensions;python_version<'3.13'", ] dynamic = ["version"] diff --git a/python/python/ballista/__init__.py b/python/python/ballista/__init__.py index e892ff8ea..97ba8ef65 100644 --- a/python/python/ballista/__init__.py +++ b/python/python/ballista/__init__.py @@ -22,10 +22,12 @@ except ImportError: from ._internal_ballista import ( BallistaExecutor, + BallistaExtension, BallistaLogicalExtensionCodec, BallistaPhysicalExtensionCodec, BallistaQueryPlanner, BallistaScheduler, + ballista_datafusion_config_defaults, setup_test_cluster, ) from .extension import ( @@ -38,8 +40,10 @@ __version__ = importlib_metadata.version(__name__) __all__ = [ "setup_test_cluster", + "ballista_datafusion_config_defaults", "BallistaScheduler", "BallistaExecutor", + "BallistaExtension", "BallistaLogicalExtensionCodec", "BallistaPhysicalExtensionCodec", "BallistaQueryPlanner", diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py index 7fc016806..ab8a2a759 100644 --- a/python/python/tests/test_context.py +++ b/python/python/tests/test_context.py @@ -15,7 +15,10 @@ # specific language governing permissions and limitations # under the License. +import gc + from ballista import ( + BallistaExtension, BallistaLogicalExtensionCodec, BallistaPhysicalExtensionCodec, BallistaQueryPlanner, @@ -92,6 +95,35 @@ def test_manual_query_planner_and_codec_composition(): assert df.collect()[0].column(0).to_pylist() == [5] +def test_ballista_extension_uses_atomic_context_installation(): + address, port = setup_test_cluster() + source_ctx = SessionContext() + source_ctx.register_csv("registered", "testdata/test.csv", has_header=True) + + configured_ctx = source_ctx.with_extensions( + BallistaExtension(f"df://{address}:{port}") + ) + + df = configured_ctx.sql("SELECT COUNT(*) FROM registered") + assert type(df) is DataFrame + assert_uses_ballista(df) + assert df.collect()[0].column(0).to_pylist() == [5] + + +def test_ballista_extension_context_must_outlive_dataframe(): + address, port = setup_test_cluster() + configured_ctx = SessionContext().with_extensions( + BallistaExtension(f"df://{address}:{port}") + ) + df = configured_ctx.sql("SELECT 1") + + del configured_ctx + gc.collect() + + with pytest.raises(Exception, match="went out of scope"): + df.collect() + + def test_low_level_helper_returns_standard_distributed_dataframe(): address, port = setup_test_cluster() source_ctx = SessionContext() diff --git a/python/requirements.txt b/python/requirements.txt index bc9cfb34d..ae3a415c6 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -1,4 +1,4 @@ -datafusion==52.0.0 +datafusion @ git+https://github.com/apache/datafusion-python.git@f1526210338d75ef105246e5c948e0e4a4a0f5ce pyarrow pytest maturin>=1.8.0,<2.0.0 diff --git a/python/src/lib.rs b/python/src/lib.rs index 7e44e2f20..810c2dae9 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -39,6 +39,7 @@ fn _internal_ballista(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { m.add_class::<PyScheduler>()?; m.add_class::<PyExecutor>()?; + m.add_class::<query_planner::PyBallistaExtension>()?; m.add_class::<query_planner::PyBallistaLogicalExtensionCodec>()?; m.add_class::<query_planner::PyBallistaPhysicalExtensionCodec>()?; m.add_class::<query_planner::PyBallistaQueryPlanner>()?; diff --git a/python/src/query_planner.rs b/python/src/query_planner.rs index c934bb9ce..0d10423cc 100644 --- a/python/src/query_planner.rs +++ b/python/src/query_planner.rs @@ -36,7 +36,7 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_proto::protobuf::LogicalPlanNode; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::PyCapsule; +use pyo3::types::{PyCapsule, PyDict}; use std::collections::HashMap; use std::ptr::NonNull; use std::sync::{Arc, OnceLock}; @@ -45,6 +45,41 @@ 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, + } + } + + fn __datafusion_session_extension__<'py>( + &self, + py: Python<'py>, + session_ctx: Bound<'py, PyAny>, + ) -> PyResult<Bound<'py, PyAny>> { + ballista_extension_components( + py, + session_ctx, + self.address.clone(), + self.config_overrides.clone(), + ) + } +} + #[pyclass( name = "BallistaLogicalExtensionCodec", module = "ballista._internal_ballista", @@ -53,8 +88,9 @@ const DEFAULT_SCHEDULER_PORT: u16 = 50050; pub(crate) struct PyBallistaLogicalExtensionCodec { logical_codec: Arc<dyn LogicalExtensionCodec>, ffi_logical_codec: FFI_LogicalExtensionCodec, - // The FFI codec holds a weak task-context provider. - _session_ctx: Py<PyAny>, + // 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] @@ -75,7 +111,7 @@ impl PyBallistaLogicalExtensionCodec { Ok(Self { logical_codec, ffi_logical_codec, - _session_ctx: session_ctx.unbind(), + _session_ctx: Some(session_ctx.unbind()), }) } @@ -98,8 +134,9 @@ impl PyBallistaLogicalExtensionCodec { )] pub(crate) struct PyBallistaPhysicalExtensionCodec { ffi_physical_codec: FFI_PhysicalExtensionCodec, - // The FFI codec holds a weak task-context provider. - _session_ctx: Py<PyAny>, + // 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>, } @@ -115,7 +152,7 @@ impl PyBallistaPhysicalExtensionCodec { Ok(Self { ffi_physical_codec, - _session_ctx: session_ctx.unbind(), + _session_ctx: Some(session_ctx.unbind()), _logical_codec: logical_codec_object, }) } @@ -163,9 +200,9 @@ pub(crate) struct PyBallistaQueryPlanner { logical_codec: Arc<dyn LogicalExtensionCodec>, fallback_logical_codec: FFI_LogicalExtensionCodec, fallback_physical_codec: FFI_PhysicalExtensionCodec, - // The fallback FFI codecs hold a weak task-context provider. Retain its - // source context for as long as Python retains this planner. - _session_ctx: Py<PyAny>, + // 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] @@ -201,7 +238,10 @@ impl PyBallistaQueryPlanner { address: self.address.clone(), config: self.config.clone(), logical_codec: Arc::clone(&self.logical_codec), - _session_ctx: self._session_ctx.clone_ref(py), + _session_ctx: self + ._session_ctx + .as_ref() + .map(|session_ctx| session_ctx.clone_ref(py)), }); let ffi = FFI_QueryPlanner::new_with_ffi_codecs( planner, @@ -209,7 +249,7 @@ impl PyBallistaQueryPlanner { self.fallback_physical_codec.clone(), ); - // This exact name is the ABI consumed by datafusion-python 55f43b5e. + // This exact name is the ABI consumed by datafusion-python. PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") } } @@ -244,7 +284,7 @@ impl PyBallistaQueryPlanner { logical_codec, fallback_logical_codec, fallback_physical_codec, - _session_ctx: session_ctx.unbind(), + _session_ctx: Some(session_ctx.unbind()), }) } } @@ -253,7 +293,7 @@ struct LazyBallistaQueryPlanner { address: String, config: ballista_core::config::BallistaConfig, logical_codec: Arc<dyn LogicalExtensionCodec>, - _session_ctx: Py<PyAny>, + _session_ctx: Option<Py<PyAny>>, } impl std::fmt::Debug for LazyBallistaQueryPlanner { @@ -317,6 +357,48 @@ pub(crate) fn ballista_datafusion_config_defaults() -> HashMap<String, String> { .collect() } +fn ballista_extension_components<'py>( + py: Python<'py>, + session_ctx: Bound<'py, PyAny>, + address: String, + config_overrides: Option<HashMap<String, String>>, +) -> PyResult<Bound<'py, PyAny>> { + let mut logical_codec = PyBallistaLogicalExtensionCodec::new(session_ctx.clone())?; + let logical_native = Arc::clone(&logical_codec.logical_codec); + let fallback_logical_codec = logical_codec.ffi_logical_codec.clone(); + logical_codec._session_ctx = None; + let logical_codec = Py::new(py, logical_codec)?; + + let mut physical_codec = PyBallistaPhysicalExtensionCodec::from_logical_codec( + session_ctx.clone(), + Arc::clone(&logical_native), + logical_codec.clone_ref(py).into_any(), + )?; + let fallback_physical_codec = physical_codec.ffi_physical_codec.clone(); + physical_codec._session_ctx = None; + let physical_codec = Py::new(py, physical_codec)?; + + let mut planner = PyBallistaQueryPlanner::from_codecs( + address, + session_ctx, + config_overrides, + logical_native, + fallback_logical_codec, + fallback_physical_codec, + )?; + planner._session_ctx = None; + let planner = Py::new(py, planner)?; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("logical_extension_codecs", (logical_codec,))?; + kwargs.set_item("physical_extension_codecs", (physical_codec,))?; + kwargs.set_item("query_planner", planner)?; + components.call((), Some(&kwargs)) +} + #[pyfunction] #[pyo3(signature = (session_ctx, address, config_overrides=None))] pub(crate) fn with_ballista_query_planner( @@ -340,37 +422,12 @@ pub(crate) fn with_ballista_query_planner( } } - let logical_codec = PyBallistaLogicalExtensionCodec::new(session_ctx.clone())?; - let logical_native = Arc::clone(&logical_codec.logical_codec); - let fallback_logical_codec = logical_codec.ffi_logical_codec.clone(); - let logical_codec = Py::new(py, logical_codec)?; - - let physical_codec = PyBallistaPhysicalExtensionCodec::from_logical_codec( - session_ctx.clone(), - Arc::clone(&logical_native), - logical_codec.clone_ref(py).into_any(), + let extension = Py::new( + py, + PyBallistaExtension::new(address, Some(ballista_overrides)), )?; - let fallback_physical_codec = physical_codec.ffi_physical_codec.clone(); - let physical_codec = Py::new(py, physical_codec)?; - - let planner = PyBallistaQueryPlanner::from_codecs( - address, - session_ctx.clone(), - Some(ballista_overrides), - logical_native, - fallback_logical_codec, - fallback_physical_codec, - )?; - let planner = Py::new(py, planner)?; - - // Install each Ballista codec once before the planner. datafusion-python - // then rebinds the planner to the receiving context's codec chains. - let ctx = session_ctx - .call_method1("with_logical_extension_codec", (logical_codec.bind(py),))?; - let ctx = - ctx.call_method1("with_physical_extension_codec", (physical_codec.bind(py),))?; - Ok(ctx - .call_method1("with_query_planner", (planner.bind(py),))? + Ok(session_ctx + .call_method1("with_extensions", (extension.bind(py),))? .unbind()) } diff --git a/python/uv.lock b/python/uv.lock index f467ad317..4966fc4be 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -52,7 +52,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "datafusion", git = "https://github.com/apache/datafusion-python.git?rev=55f43b5e231da5591e7bbd632dac85f304a22cca"; }, + { name = "datafusion", git = "https://github.com/apache/datafusion-python.git?rev=f1526210338d75ef105246e5c948e0e4a4a0f5ce"; }, { name = "ipython", marker = "extra == 'jupyter'", specifier = ">=8.0.0" }, { name = "pyarrow", specifier = ">=23.0.1" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, @@ -328,7 +328,7 @@ wheels = [ [[package]] name = "datafusion" version = "54.0.0" -source = { git = "https://github.com/apache/datafusion-python.git?rev=55f43b5e231da5591e7bbd632dac85f304a22cca#55f43b5e231da5591e7bbd632dac85f304a22cca"; } +source = { git = "https://github.com/apache/datafusion-python.git?rev=f1526210338d75ef105246e5c948e0e4a4a0f5ce#f1526210338d75ef105246e5c948e0e4a4a0f5ce"; } dependencies = [ { name = "cloudpickle" }, { name = "pyarrow" }, --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
