Fix issue in distributed_codec

Author: gabotechs

src/execution_plans/distributed.rs

@@ -138,7 +138,7 @@ impl ExecutionPlan for DistributedExec {
         }
 
         let channel_resolver = get_distributed_channel_resolver(context.session_config())?;
-        let codec = DistributedCodec::new_combined_with_user(Arc::clone(&context));
+        let codec = DistributedCodec::new_combined_with_user(context.session_config());
 
         let prepared = self.prepare_plan(&channel_resolver.get_urls()?, &codec)?;
         {

src/flight_service/do_get.rs

@@ -81,7 +81,8 @@ impl ArrowFlightEndpoint {
             .await
             .map_err(|err| datafusion_error_to_tonic_status(&err))?;
 
-        let codec = DistributedCodec::new_combined_with_user(session_state.task_ctx());
+        let codec = DistributedCodec::new_combined_with_user(session_state.config());
+        let ctx = SessionContext::new_with_state(session_state.clone());
 
         // There's only 1 `StageExec` responsible for all requests that share the same `stage_key`,
         // so here we either retrieve the existing one or create a new one if it does not exist.
@@ -93,7 +94,6 @@ impl ArrowFlightEndpoint {
         let stage_data = once
             .get_or_try_init(|| async {
                 let proto_node = PhysicalPlanNode::try_decode(doget.plan_proto.as_ref())?;
-                let ctx = SessionContext::from(session_state.clone());
                 let plan = proto_node.try_into_physical_plan(&ctx, &self.runtime, &codec)?;
 
                 // Initialize partition count to the number of partitions in the stage

src/protobuf/distributed_codec.rs

@@ -8,38 +8,29 @@ use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::internal_datafusion_err;
 use datafusion::error::DataFusionError;
-use datafusion::execution::{FunctionRegistry, TaskContext};
-use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+use datafusion::execution::{FunctionRegistry, SessionStateBuilder};
 use datafusion::physical_expr::EquivalenceProperties;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::{ExecutionPlan, Partitioning, PlanProperties};
-use datafusion::prelude::SessionContext;
+use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning;
 use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
 use datafusion_proto::physical_plan::{ComposedPhysicalExtensionCodec, PhysicalExtensionCodec};
 use datafusion_proto::protobuf;
 use datafusion_proto::protobuf::proto_error;
 use prost::Message;
-use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 use url::Url;
 
 /// DataFusion [PhysicalExtensionCodec] implementation that allows serializing and
 /// deserializing the custom ExecutionPlans in this project
-#[derive(Clone, Default)]
-pub struct DistributedCodec(Arc<TaskContext>);
-
-impl Debug for DistributedCodec {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "DistributedCodec")
-    }
-}
+#[derive(Debug)]
+pub struct DistributedCodec;
 
 impl DistributedCodec {
-    pub fn new_combined_with_user(ctx: Arc<TaskContext>) -> impl PhysicalExtensionCodec {
-        let mut codecs: Vec<Arc<dyn PhysicalExtensionCodec>> =
-            vec![Arc::new(DistributedCodec(Arc::clone(&ctx)))];
-        codecs.extend(get_distributed_user_codecs(ctx.session_config()));
+    pub fn new_combined_with_user(cfg: &SessionConfig) -> impl PhysicalExtensionCodec + use<> {
+        let mut codecs: Vec<Arc<dyn PhysicalExtensionCodec>> = vec![Arc::new(DistributedCodec {})];
+        codecs.extend(get_distributed_user_codecs(cfg));
         ComposedPhysicalExtensionCodec::new(codecs)
     }
 }
@@ -49,7 +40,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
         &self,
         buf: &[u8],
         inputs: &[Arc<dyn ExecutionPlan>],
-        _registry: &dyn FunctionRegistry,
+        registry: &dyn FunctionRegistry,
     ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         let DistributedExecProto {
             node: Some(distributed_exec_node),
@@ -63,7 +54,16 @@ impl PhysicalExtensionCodec for DistributedCodec {
         // TODO: The PhysicalExtensionCodec trait doesn't provide access to session state,
         // so we create a new SessionContext which loses any custom UDFs, UDAFs, and other
         // user configurations. This is a limitation of the current trait design.
-        let ctx = SessionContext::new();
+        let state = SessionStateBuilder::new()
+            .with_scalar_functions(
+                registry
+                    .udfs()
+                    .iter()
+                    .map(|f| registry.udf(f))
+                    .collect::<Result<Vec<_>, _>>()?,
+            )
+            .build();
+        let ctx = SessionContext::from(state);
 
         fn parse_stage_proto(
             proto: Option<StageProto>,
@@ -112,9 +112,13 @@ impl PhysicalExtensionCodec for DistributedCodec {
                     .map(|s| s.try_into())
                     .ok_or(proto_error("NetworkShuffleExec is missing schema"))??;
 
-                let partitioning =
-                    parse_protobuf_partitioning(partitioning.as_ref(), &ctx, &schema, self)?
-                        .ok_or(proto_error("NetworkShuffleExec is missing partitioning"))?;
+                let partitioning = parse_protobuf_partitioning(
+                    partitioning.as_ref(),
+                    &ctx,
+                    &schema,
+                    &DistributedCodec {},
+                )?
+                .ok_or(proto_error("NetworkShuffleExec is missing partitioning"))?;
 
                 Ok(Arc::new(new_network_hash_shuffle_exec(
                     partitioning,
@@ -132,9 +136,13 @@ impl PhysicalExtensionCodec for DistributedCodec {
                     .map(|s| s.try_into())
                     .ok_or(proto_error("NetworkCoalesceExec is missing schema"))??;
 
-                let partitioning =
-                    parse_protobuf_partitioning(partitioning.as_ref(), &ctx, &schema, self)?
-                        .ok_or(proto_error("NetworkCoalesceExec is missing partitioning"))?;
+                let partitioning = parse_protobuf_partitioning(
+                    partitioning.as_ref(),
+                    &ctx,
+                    &schema,
+                    &DistributedCodec {},
+                )?
+                .ok_or(proto_error("NetworkCoalesceExec is missing partitioning"))?;
 
                 Ok(Arc::new(new_network_coalesce_tasks_exec(
                     partitioning,
@@ -185,7 +193,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
                 schema: Some(node.schema().try_into()?),
                 partitioning: Some(serialize_partitioning(
                     node.properties().output_partitioning(),
-                    self,
+                    &DistributedCodec {},
                 )?),
                 input_stage: Some(encode_stage_proto(node.input_stage())?),
             };
@@ -200,7 +208,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
                 schema: Some(node.schema().try_into()?),
                 partitioning: Some(serialize_partitioning(
                     node.properties().output_partitioning(),
-                    self,
+                    &DistributedCodec {},
                 )?),
                 input_stage: Some(encode_stage_proto(node.input_stage())?),
             };
@@ -229,30 +237,6 @@ impl PhysicalExtensionCodec for DistributedCodec {
             Err(proto_error(format!("Unexpected plan {}", node.name())))
         }
     }
-
-    fn try_decode_udf(
-        &self,
-        name: &str,
-        _buf: &[u8],
-    ) -> datafusion::common::Result<Arc<ScalarUDF>> {
-        self.0.udf(name)
-    }
-
-    fn try_decode_udaf(
-        &self,
-        name: &str,
-        _buf: &[u8],
-    ) -> datafusion::common::Result<Arc<AggregateUDF>> {
-        self.0.udaf(name)
-    }
-
-    fn try_decode_udwf(
-        &self,
-        name: &str,
-        _buf: &[u8],
-    ) -> datafusion::common::Result<Arc<WindowUDF>> {
-        self.0.udwf(name)
-    }
 }
 
 /// A key that uniquely identifies a stage in a query
@@ -436,7 +420,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_single_flight() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec::default();
+        let codec = DistributedCodec;
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("a");
@@ -455,7 +439,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_isolator_flight() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec::default();
+        let codec = DistributedCodec;
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("b");
@@ -479,7 +463,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_isolator_union() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec::default();
+        let codec = DistributedCodec;
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("c");
@@ -509,7 +493,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_isolator_sort_flight() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec::default();
+        let codec = DistributedCodec;
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("d");