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 9e813516a9cec047a7497e73e902e5ec6ef0b957 Author: Tim Saucer <[email protected]> AuthorDate: Fri Aug 7 13:27:44 2026 -0400 feat(python): expose Ballista extension codecs AI Disclosure: This code was written in part by an AI agent. --- python/python/ballista/__init__.py | 8 +- python/python/tests/test_context.py | 45 +++++-- python/src/lib.rs | 2 + python/src/query_planner.rs | 251 +++++++++++++++++++++++++----------- 4 files changed, 223 insertions(+), 83 deletions(-) diff --git a/python/python/ballista/__init__.py b/python/python/ballista/__init__.py index 5a5ef7bb0..e892ff8ea 100644 --- a/python/python/ballista/__init__.py +++ b/python/python/ballista/__init__.py @@ -21,8 +21,11 @@ except ImportError: import importlib_metadata from ._internal_ballista import ( - BallistaScheduler, BallistaExecutor, + BallistaLogicalExtensionCodec, + BallistaPhysicalExtensionCodec, + BallistaQueryPlanner, + BallistaScheduler, setup_test_cluster, ) from .extension import ( @@ -37,6 +40,9 @@ __all__ = [ "setup_test_cluster", "BallistaScheduler", "BallistaExecutor", + "BallistaLogicalExtensionCodec", + "BallistaPhysicalExtensionCodec", + "BallistaQueryPlanner", "BallistaSessionContext", "DistributedDataFrame", "ExecutionPlanVisualization", diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py index e823a5ed8..7fc016806 100644 --- a/python/python/tests/test_context.py +++ b/python/python/tests/test_context.py @@ -15,9 +15,14 @@ # specific language governing permissions and limitations # under the License. -from ballista import BallistaSessionContext, setup_test_cluster -from ballista._internal_ballista import ( +from ballista import ( + BallistaLogicalExtensionCodec, + BallistaPhysicalExtensionCodec, BallistaQueryPlanner, + BallistaSessionContext, + setup_test_cluster, +) +from ballista._internal_ballista import ( ballista_datafusion_config_defaults, with_ballista_query_planner, ) @@ -25,9 +30,8 @@ from ballista.extension import ( DataFrame, DistributedDataFrame, SessionConfig, - SessionContext, ) -from datafusion import col, lit +from datafusion import SessionContext, col, lit import ctypes import pytest import pyarrow as pa @@ -43,25 +47,50 @@ def assert_uses_ballista(df): assert "DistributedQueryExec" in str(df.execution_plan()) -def test_query_planner_capsules_use_current_datafusion_abi(): +def test_ballista_ffi_capsules_use_current_datafusion_abi(): source_ctx = SessionContext() + logical_codec = BallistaLogicalExtensionCodec(source_ctx) + physical_codec = BallistaPhysicalExtensionCodec(source_ctx, logical_codec) planner = BallistaQueryPlanner("df://localhost:50050", source_ctx) is_valid = ctypes.pythonapi.PyCapsule_IsValid is_valid.argtypes = [ctypes.py_object, ctypes.c_char_p] is_valid.restype = ctypes.c_int capsules = { - b"datafusion_query_planner": planner.__datafusion_query_planner__(), b"datafusion_logical_extension_codec": ( - planner.__datafusion_logical_extension_codec__() + logical_codec.__datafusion_logical_extension_codec__() ), b"datafusion_physical_extension_codec": ( - planner.__datafusion_physical_extension_codec__() + physical_codec.__datafusion_physical_extension_codec__() ), + b"datafusion_query_planner": planner.__datafusion_query_planner__(), } for name, capsule in capsules.items(): assert is_valid(capsule, name) == 1 + assert not hasattr(logical_codec, "__datafusion_physical_extension_codec__") + assert not hasattr(physical_codec, "__datafusion_logical_extension_codec__") + assert not hasattr(planner, "__datafusion_logical_extension_codec__") + assert not hasattr(planner, "__datafusion_physical_extension_codec__") + + +def test_manual_query_planner_and_codec_composition(): + address, port = setup_test_cluster() + source_ctx = SessionContext() + source_ctx.register_csv("registered", "testdata/test.csv", has_header=True) + logical_codec = BallistaLogicalExtensionCodec(source_ctx) + physical_codec = BallistaPhysicalExtensionCodec(source_ctx, logical_codec) + planner = BallistaQueryPlanner(f"df://{address}:{port}", source_ctx) + + configured_ctx = source_ctx.with_logical_extension_codec(logical_codec) + configured_ctx = configured_ctx.with_physical_extension_codec(physical_codec) + configured_ctx = configured_ctx.with_query_planner(planner) + + 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_low_level_helper_returns_standard_distributed_dataframe(): address, port = setup_test_cluster() diff --git a/python/src/lib.rs b/python/src/lib.rs index ba58070d3..7e44e2f20 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -39,6 +39,8 @@ fn _internal_ballista(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { m.add_class::<PyScheduler>()?; m.add_class::<PyExecutor>()?; + m.add_class::<query_planner::PyBallistaLogicalExtensionCodec>()?; + m.add_class::<query_planner::PyBallistaPhysicalExtensionCodec>()?; m.add_class::<query_planner::PyBallistaQueryPlanner>()?; m.add_class::<datafusion_python::dataframe::PyParquetWriterOptions>()?; diff --git a/python/src/query_planner.rs b/python/src/query_planner.rs index 403ca99ae..c934bb9ce 100644 --- a/python/src/query_planner.rs +++ b/python/src/query_planner.rs @@ -18,7 +18,8 @@ use ballista_core::extension::SessionConfigExt; use ballista_core::planner::BallistaQueryPlanner; use ballista_core::serde::{ - BallistaLogicalExtensionCodec, BallistaPhysicalExtensionCodec, + BallistaLogicalExtensionCodec as NativeBallistaLogicalExtensionCodec, + BallistaPhysicalExtensionCodec as NativeBallistaPhysicalExtensionCodec, }; use datafusion::catalog::Session; use datafusion::common::{DataFusionError, Result}; @@ -30,9 +31,7 @@ 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::{ - DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; +use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_proto::protobuf::LogicalPlanNode; use pyo3::exceptions::PyValueError; @@ -47,82 +46,35 @@ use url::Url; const DEFAULT_SCHEDULER_PORT: u16 = 50050; #[pyclass( - name = "BallistaQueryPlanner", + name = "BallistaLogicalExtensionCodec", module = "ballista._internal_ballista", skip_from_py_object )] -pub(crate) struct PyBallistaQueryPlanner { - address: String, - config: ballista_core::config::BallistaConfig, +pub(crate) struct PyBallistaLogicalExtensionCodec { logical_codec: Arc<dyn LogicalExtensionCodec>, ffi_logical_codec: FFI_LogicalExtensionCodec, - ffi_physical_codec: FFI_PhysicalExtensionCodec, - // The FFI codecs hold a weak task-context provider. Retain its source - // context for as long as Python retains this planner. + // The FFI codec holds a weak task-context provider. _session_ctx: Py<PyAny>, } #[pymethods] -impl PyBallistaQueryPlanner { +impl PyBallistaLogicalExtensionCodec { #[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> { + fn new(session_ctx: Bound<'_, PyAny>) -> PyResult<Self> { let user_logical_codec: Arc<dyn LogicalExtensionCodec> = (&ffi_logical_codec_from_python(session_ctx.clone())?).into(); - // The planner-side codec must retain the source context's Python-aware - // codecs because DistributedQueryExec embeds a serialized logical plan. let logical_codec: Arc<dyn LogicalExtensionCodec> = - Arc::new(BallistaLogicalExtensionCodec::new(user_logical_codec)); - let physical_codec: Arc<dyn PhysicalExtensionCodec> = Arc::new( - BallistaPhysicalExtensionCodec::new(Arc::clone(&logical_codec)), - ); - - // The receiving SessionContext composes an installed codec with its own - // codecs. Export only the Ballista layer here; wrapping the source codec - // again would make the receiver traverse the same codec stack twice. - let installed_logical_codec: Arc<dyn LogicalExtensionCodec> = Arc::new( - BallistaLogicalExtensionCodec::new(Arc::new( - DefaultLogicalExtensionCodec {}, - )), - ); + Arc::new(NativeBallistaLogicalExtensionCodec::new(user_logical_codec)); let task_ctx_provider = ffi_task_ctx_provider_from_python(session_ctx.clone())?; - let runtime = tokio_runtime_handle(); let ffi_logical_codec = FFI_LogicalExtensionCodec::new( - installed_logical_codec, - Some(runtime.clone()), - task_ctx_provider.clone(), - ); - let ffi_physical_codec = FFI_PhysicalExtensionCodec::new( - physical_codec, - Some(runtime), - task_ctx_provider.clone(), + Arc::clone(&logical_codec), + Some(tokio_runtime_handle()), + task_ctx_provider, ); - 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, ffi_logical_codec, - ffi_physical_codec, _session_ctx: session_ctx.unbind(), }) } @@ -137,6 +89,56 @@ impl PyBallistaQueryPlanner { 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, + // The FFI codec holds a weak task-context provider. + _session_ctx: 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: 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(), + ) + } fn __datafusion_physical_extension_codec__<'py>( &self, @@ -148,6 +150,47 @@ impl PyBallistaQueryPlanner { 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, + // 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>, +} + +#[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, + ) + } fn __datafusion_query_planner__<'py>( &self, @@ -162,8 +205,8 @@ impl PyBallistaQueryPlanner { }); let ffi = FFI_QueryPlanner::new_with_ffi_codecs( planner, - self.ffi_logical_codec.clone(), - self.ffi_physical_codec.clone(), + self.fallback_logical_codec.clone(), + self.fallback_physical_codec.clone(), ); // This exact name is the ABI consumed by datafusion-python 55f43b5e. @@ -171,6 +214,41 @@ impl PyBallistaQueryPlanner { } } +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: session_ctx.unbind(), + }) + } +} + struct LazyBallistaQueryPlanner { address: String, config: ballista_core::config::BallistaConfig, @@ -262,29 +340,54 @@ pub(crate) fn with_ballista_query_planner( } } - let planner = PyBallistaQueryPlanner::new( + 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 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)?; - // datafusion-python deliberately rebinds imported planners to the codecs - // installed on the receiving context. Install Ballista's codecs first so - // that rebind keeps DistributedQueryExec serialization active. - let ctx = session_ctx.call_method1( - "with_logical_extension_codec", - (planner.bind(py),), - )?; - let ctx = ctx.call_method1( - "with_physical_extension_codec", - (planner.bind(py),), - )?; + // 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),))? .unbind()) } +fn ballista_physical_ffi_codec( + session_ctx: Bound<'_, PyAny>, + logical_codec: Arc<dyn LogicalExtensionCodec>, +) -> PyResult<FFI_PhysicalExtensionCodec> { + let physical_codec: Arc<dyn PhysicalExtensionCodec> = + Arc::new(NativeBallistaPhysicalExtensionCodec::new(logical_codec)); + let task_ctx_provider = ffi_task_ctx_provider_from_python(session_ctx)?; + Ok(FFI_PhysicalExtensionCodec::new( + physical_codec, + Some(tokio_runtime_handle()), + task_ctx_provider, + )) +} + fn apply_datafusion_override( session_ctx: &Bound<'_, PyAny>, key: &str, --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
