Re: [PR] Extensions API for third-party table/catalog providers [datafusion-ray]

[via GitHub]

ccciudatu commented on code in PR #43:
URL: https://github.com/apache/datafusion-ray/pull/43#discussion_r1857074011


##########
src/context.rs:
##########
@@ -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())
+    }

Review Comment:
   Got it, I'll look into that. Thanks!



-- 
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: github-unsubscr...@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org