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

[via GitHub]

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


##########
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:
   I don't want to hold up progress, but this *feels* like a step in the wrong 
direction. But I also don't have enough time right now to give a better 
solution. I think I'd want to do something like use the FFI_ExecutionPlan in 
df43 to share these across packages.
   



-- 
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