Extensions framework

Author: ccciudatu

Cargo.lock

@@ -25,19 +25,21 @@ version = "0.1.0"
 edition = "2021"
 readme = "README.md"
 license = "Apache-2.0"
-rust-version = "1.62"
+rust-version = "1.70"
 build = "build.rs"
 
 [dependencies]
 datafusion = { version = "42.0.0", features = ["pyarrow", "avro"] }
 datafusion-proto = "42.0.0"
+datafusion-python = "42.0.0"
 futures = "0.3"
 glob = "0.3.1"
 log = "0.4"
 prost = "0.13"
 pyo3 = { version = "0.22", features = ["extension-module", "abi3", "abi3-py38"] }
 tokio = { version = "1.40", features = ["macros", "rt", "rt-multi-thread", "sync"] }
 uuid = "1.11.0"
+async-trait = "0.1.83"
 
 [build-dependencies]
 prost-types = "0.13"

Cargo.toml

@@ -25,6 +25,7 @@
     ExecutionGraph,
     QueryStage,
     execute_partition,
+    extended_session_context,
 )
 from .context import DatafusionRayContext
 

datafusion_ray/__init__.py

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::ext::Extensions;
 use crate::planner::{make_execution_graph, PyExecutionGraph};
-use crate::shuffle::ShuffleCodec;
 use datafusion::arrow::pyarrow::ToPyArrow;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::error::{DataFusionError, Result};
@@ -26,6 +26,7 @@ use datafusion::physical_plan::{displayable, ExecutionPlan};
 use datafusion::prelude::*;
 use datafusion_proto::physical_plan::AsExecutionPlan;
 use datafusion_proto::protobuf;
+use datafusion_python::physical_plan::PyExecutionPlan;
 use futures::StreamExt;
 use prost::Message;
 use pyo3::exceptions::PyRuntimeError;
@@ -45,22 +46,30 @@ pub struct PyContext {
 
 pub(crate) fn execution_plan_from_pyany(
     py_plan: &Bound<PyAny>,
+    py: Python,
 ) -> PyResult<Arc<dyn ExecutionPlan>> {
-    let py_proto = py_plan.call_method0("to_proto")?;
-    let plan_bytes: &[u8] = py_proto.extract()?;
-    let plan_node = protobuf::PhysicalPlanNode::try_decode(plan_bytes).map_err(|e| {
-        PyRuntimeError::new_err(format!(
-            "Unable to decode physical plan protobuf message: {}",
-            e
-        ))
-    })?;
-
-    let codec = ShuffleCodec {};
-    let runtime = RuntimeEnv::default();
-    let registry = SessionContext::new();
-    plan_node
-        .try_into_physical_plan(&registry, &runtime, &codec)
-        .map_err(|e| e.into())
+    if let Ok(py_plan) = py_plan.to_object(py).downcast_bound::<PyExecutionPlan>(py) {
+        // For session contexts created with datafusion_ray.extended_session_context(), the inner
+        // execution plan can be used as such (and the enabled extensions are all available).
+        Ok(py_plan.borrow().plan.clone())
+    } else {
+        // The session context originates from outside our library, so we'll grab the protobuf plan
+        // by calling the python method with no extension codecs.
+        let py_proto = py_plan.call_method0("to_proto")?;
+        let plan_bytes: &[u8] = py_proto.extract()?;
+        let plan_node = protobuf::PhysicalPlanNode::try_decode(plan_bytes).map_err(|e| {
+            PyRuntimeError::new_err(format!(
+                "Unable to decode physical plan protobuf message: {}",
+                e
+            ))
+        })?;
+
+        let runtime = RuntimeEnv::default();
+        let registry = SessionContext::new();
+        plan_node
+            .try_into_physical_plan(&registry, &runtime, Extensions::codec())
+            .map_err(|e| e.into())
+    }
 }
 
 #[pymethods]
@@ -87,14 +96,14 @@ impl PyContext {
     }
 
     /// Plan a distributed SELECT query for executing against the Ray workers
-    pub fn plan(&self, plan: &Bound<PyAny>) -> PyResult<PyExecutionGraph> {
+    pub fn plan(&self, plan: &Bound<PyAny>, py: Python) -> PyResult<PyExecutionGraph> {
         // println!("Planning {}", sql);
         // let df = wait_for_future(py, self.ctx.sql(sql))?;
         // let py_df = self.run_sql(sql, py)?;
         // let py_plan = py_df.call_method0(py, "execution_plan")?;
         // let py_plan = py_plan.bind(py);
 
-        let plan = execution_plan_from_pyany(plan)?;
+        let plan = execution_plan_from_pyany(plan, py)?;
         let graph = make_execution_graph(plan.clone())?;
 
         // debug logging
@@ -140,9 +149,10 @@ pub fn serialize_execution_plan(
     plan: Arc<dyn ExecutionPlan>,
     py: Python,
 ) -> PyResult<Bound<'_, PyBytes>> {
-    let codec = ShuffleCodec {};
-    let proto =
-        datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan(plan.clone(), &codec)?;
+    let proto = datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan(
+        plan.clone(),
+        Extensions::codec(),
+    )?;
 
     let bytes = proto.encode_to_vec();
     Ok(PyBytes::new_bound(py, &bytes))
@@ -159,9 +169,8 @@ pub fn deserialize_execution_plan(proto_msg: &Bound<PyBytes>) -> PyResult<Arc<dy
         })?;
 
     let ctx = SessionContext::new();
-    let codec = ShuffleCodec {};
     let plan = proto_plan
-        .try_into_physical_plan(&ctx, &ctx.runtime_env(), &codec)
+        .try_into_physical_plan(&ctx, &ctx.runtime_env(), Extensions::codec())
         .map_err(DataFusionError::from)?;
 
     Ok(plan)

src/context.rs

@@ -0,0 +1,121 @@
+use async_trait::async_trait;
+use datafusion::common::DataFusionError;
+use datafusion::common::Result;
+use datafusion::execution::FunctionRegistry;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::SessionContext;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_python::context::PySessionContext;
+use datafusion_python::utils::wait_for_future;
+use pyo3::{pyfunction, PyResult, Python};
+use std::collections::HashMap;
+use std::fmt::Debug;
+use std::sync::{Arc, OnceLock};
+
+mod built_in;
+
+/// Creates a datafusion session context preconfigured with the enabled extensions
+/// that will register additional table providers, catalogs etc.
+/// If no extensions are required, the plain `datafusion.SessionContext()` will work just fine.
+/// # Arguments
+/// * `settings` - dictionary containing extension-specific key/value config options
+#[pyfunction]
+pub fn extended_session_context(
+    settings: HashMap<String, String>,
+    py: Python,
+) -> PyResult<PySessionContext> {
+    let future_context = Extensions::session_context(&settings);
+    let ctx = wait_for_future(py, future_context)?;
+    Ok(ctx.into())
+}
+
+/// Allows third party table/catalog providers, object stores, etc.
+/// to be registered with the DataFusion context.
+#[async_trait]
+trait Extension: Debug + Send + Sync + 'static {
+    /// SessionContext initialization, using the provided key/value settings if needed.
+    /// Declared async to allow implementers to perform network or other I/O operations.
+    async fn init(&self, ctx: &SessionContext, settings: &HashMap<String, String>) -> Result<()> {
+        let _ = ctx;
+        let _ = settings;
+        Ok(())
+    }
+
+    /// Codecs for the custom physical plan nodes created by this extension, if any.
+    fn codecs(&self) -> Vec<Box<dyn PhysicalExtensionCodec>> {
+        vec![]
+    }
+}
+
+/// A composite extension registry for enabled extensions.
+#[derive(Debug)]
+pub(crate) struct Extensions(Box<[Box<dyn Extension>]>);
+
+#[async_trait]
+impl Extension for Extensions {
+    async fn init(&self, ctx: &SessionContext, settings: &HashMap<String, String>) -> Result<()> {
+        for ext in &self.0 {
+            ext.init(ctx, settings).await?;
+        }
+        Ok(())
+    }
+
+    fn codecs(&self) -> Vec<Box<dyn PhysicalExtensionCodec>> {
+        self.0.iter().flat_map(|ext| ext.codecs()).collect()
+    }
+}
+
+impl Extensions {
+    fn new() -> Self {
+        Self(Box::new([
+            Box::new(built_in::DefaultExtension::default()),
+        ]))
+    }
+
+    fn singleton() -> &'static Self {
+        static EXTENSIONS: OnceLock<Extensions> = OnceLock::new();
+        EXTENSIONS.get_or_init(Self::new)
+    }
+
+    pub(crate) async fn session_context(
+        settings: &HashMap<String, String>,
+    ) -> Result<SessionContext> {
+        let ctx = SessionContext::new();
+        Self::singleton().init(&ctx, settings).await?;
+        Ok(ctx)
+    }
+
+    pub(crate) fn codec() -> &'static CompositeCodec {
+        static COMPOSITE_CODEC: OnceLock<CompositeCodec> = OnceLock::new();
+        COMPOSITE_CODEC.get_or_init(|| CompositeCodec(Extensions::singleton().codecs().into()))
+    }
+}
+
+/// For both encoding and decoding, tries all the registered extension codecs and returns the first successful result.
+#[derive(Debug)]
+pub(crate) struct CompositeCodec(Box<[Box<dyn PhysicalExtensionCodec>]>);
+
+impl PhysicalExtensionCodec for CompositeCodec {
+    fn try_decode(
+        &self,
+        buf: &[u8],
+        inputs: &[Arc<dyn ExecutionPlan>],
+        registry: &dyn FunctionRegistry,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.0
+            .iter()
+            .filter_map(|codec| codec.try_decode(buf, inputs, registry).ok())
+            .next()
+            .ok_or_else(|| DataFusionError::Execution("No compatible codec found".into()))
+    }
+
+    fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
+        self.0
+            .iter()
+            .filter_map(|codec| codec.try_encode(node.clone(), buf).ok())
+            .next()
+            .ok_or_else(|| {
+                DataFusionError::Execution(format!("No compatible codec found for {}", node.name()))
+            })
+    }
+}

src/ext.rs

@@ -0,0 +1,14 @@
+use crate::ext::Extension;
+use crate::shuffle::ShuffleCodec;
+use async_trait::async_trait;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+
+#[derive(Debug, Default)]
+pub(super) struct DefaultExtension {}
+
+#[async_trait]
+impl Extension for DefaultExtension {
+    fn codecs(&self) -> Vec<Box<dyn PhysicalExtensionCodec>> {
+        vec![Box::new(ShuffleCodec {})]
+    }
+}

src/ext/built_in.rs

@@ -23,6 +23,9 @@ mod proto;
 use crate::context::execute_partition;
 pub use proto::generated::protobuf;
 
+mod ext;
+use crate::ext::extended_session_context;
+
 pub mod context;
 pub mod planner;
 pub mod query_stage;
@@ -36,5 +39,6 @@ fn _datafusion_ray_internal(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<planner::PyExecutionGraph>()?;
     m.add_class::<query_stage::PyQueryStage>()?;
     m.add_function(wrap_pyfunction!(execute_partition, m)?)?;
+    m.add_function(wrap_pyfunction!(extended_session_context, m)?)?;
     Ok(())
 }

src/lib.rs

@@ -16,7 +16,8 @@
 // under the License.
 
 use crate::context::serialize_execution_plan;
-use crate::shuffle::{ShuffleCodec, ShuffleReaderExec};
+use crate::ext::Extensions;
+use crate::shuffle::ShuffleReaderExec;
 use datafusion::error::Result;
 use datafusion::physical_plan::{ExecutionPlan, Partitioning};
 use datafusion::prelude::SessionContext;
@@ -41,8 +42,8 @@ impl PyQueryStage {
     #[new]
     pub fn new(id: usize, bytes: Vec<u8>) -> Result<Self> {
         let ctx = SessionContext::new();
-        let codec = ShuffleCodec {};
-        let plan = physical_plan_from_bytes_with_extension_codec(&bytes, &ctx, &codec)?;
+        let plan =
+            physical_plan_from_bytes_with_extension_codec(&bytes, &ctx, Extensions::codec())?;
         Ok(PyQueryStage {
             stage: Arc::new(QueryStage { id, plan }),
         })