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_r1856855689


##########
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:
   Thanks, @timsaucer !
   Makes sense, and I'm entirely onboard with this goal. Which is why I tried 
to preserve this guarantee by keeping the existing code within the `else` 
branch here, where no assumption is made about the `datafusion-python` 
version/compiler.
   
   The "embedded" `datafusion-python` dependency is only supposed to be an 
opt-in alternative for users who decide to call the new function for creating 
an ABI compatible context preconfigured with the enabled extensions (that's the 
only way the above downcast can succeed, if I'm not mistaken).
   So any existing or future code that doesn't switch to using the new 
`extended_session_context()` function will continue to work without any 
compatibility restrictions.
   However, if I somehow failed to preserve this guarantee or if I missed 
something that introduces any potential risks, please let me know so I'll 
revisit the approach.



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