[HSTACK] - move codec() to physical_plan, now callable from python

Author: aditanase

src/dataframe.rs

@@ -56,7 +56,7 @@ use crate::catalog::PyTable;
 use crate::common::df_schema::PyDFSchema;
 use crate::errors::{py_datafusion_err, PyDataFusionError};
 use crate::expr::sort_expr::to_sort_expressions;
-use crate::physical_plan::PyExecutionPlan;
+use crate::physical_plan::{ codec, PyExecutionPlan } ;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::{get_tokio_runtime, validate_pycapsule, wait_for_future};
@@ -865,11 +865,6 @@ pub fn partition_stream(serialized_plan: &[u8], partition: usize, py: Python) ->
         .map_err(py_datafusion_err)
 }
 
-fn codec() -> &'static dyn PhysicalExtensionCodec {
-    static CODEC: DeltaPhysicalCodec = DeltaPhysicalCodec {};
-    &CODEC
-}
-
 /// Print DataFrame
 fn print_dataframe(py: Python, df: DataFrame) -> PyDataFusionResult<()> {
     // Get string representation of record batches

src/physical_plan.rs

@@ -16,11 +16,11 @@
 // under the License.
 
 use datafusion::physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties};
-use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec};
+use datafusion_proto::physical_plan::{AsExecutionPlan, PhysicalExtensionCodec};
+use deltalake::delta_datafusion::DeltaPhysicalCodec;
 use prost::Message;
-use std::sync::Arc;
-
 use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyBytes};
+use std::sync::Arc;
 
 use crate::{context::PySessionContext, errors::PyDataFusionResult};
 
@@ -59,10 +59,9 @@ impl PyExecutionPlan {
     }
 
     pub fn to_proto<'py>(&'py self, py: Python<'py>) -> PyDataFusionResult<Bound<'py, PyBytes>> {
-        let codec = DefaultPhysicalExtensionCodec {};
         let proto = datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan(
             self.plan.clone(),
-            &codec,
+            codec(),
         )?;
 
         let bytes = proto.encode_to_vec();
@@ -83,8 +82,7 @@ impl PyExecutionPlan {
                 ))
             })?;
 
-        let codec = DefaultPhysicalExtensionCodec {};
-        let plan = proto_plan.try_into_physical_plan(&ctx.ctx, &ctx.ctx.runtime_env(), &codec)?;
+        let plan = proto_plan.try_into_physical_plan(&ctx.ctx, &ctx.ctx.runtime_env(), codec())?;
         Ok(Self::new(plan))
     }
 
@@ -109,3 +107,8 @@ impl From<Arc<dyn ExecutionPlan>> for PyExecutionPlan {
         PyExecutionPlan { plan: plan.clone() }
     }
 }
+
+pub(crate) fn codec() -> &'static dyn PhysicalExtensionCodec {
+    static CODEC: DeltaPhysicalCodec = DeltaPhysicalCodec {};
+    &CODEC
+}