Better UX for providing user defined codecs

Author: gabotechs

src/flight_service/do_get.rs

@@ -1,6 +1,7 @@
 use crate::composed_extension_codec::ComposedPhysicalExtensionCodec;
 use crate::errors::datafusion_error_to_tonic_status;
 use crate::flight_service::service::ArrowFlightEndpoint;
+use crate::get_user_codec;
 use crate::plan::DistributedCodec;
 use crate::stage::{stage_from_proto, ExecutionStageProto};
 use arrow_flight::encode::FlightDataEncoderBuilder;
@@ -48,10 +49,9 @@ impl ArrowFlightEndpoint {
             "FunctionRegistry not present in newly built SessionState",
         ))?;
 
-        let user_codec = self.session_builder.codec();
         let mut combined_codec = ComposedPhysicalExtensionCodec::default();
         combined_codec.push(DistributedCodec);
-        if let Some(ref user_codec) = user_codec {
+        if let Some(ref user_codec) = get_user_codec(state.config()) {
             combined_codec.push_arc(Arc::clone(&user_codec));
         }
 
@@ -62,9 +62,6 @@ impl ArrowFlightEndpoint {
             &combined_codec,
         )
         .map_err(|err| Status::invalid_argument(format!("Cannot decode stage proto: {err}")))?;
-        if let Some(user_codec) = user_codec {
-            stage = stage.with_user_codec(user_codec)
-        }
         let inner_plan = Arc::clone(&stage.plan);
 
         // Add the extensions that might be required for ExecutionPlan nodes in the plan

src/flight_service/session_builder.rs

@@ -1,41 +1,22 @@
 use datafusion::execution::SessionStateBuilder;
-use datafusion_proto::physical_plan::PhysicalExtensionCodec;
-use std::sync::Arc;
 
 /// Trait called by the Arrow Flight endpoint that handles distributed parts of a DataFusion
 /// plan for building a DataFusion's [datafusion::prelude::SessionContext].
 pub trait SessionBuilder {
     /// Takes a [SessionStateBuilder] and adds whatever is necessary for it to work, like
-    /// custom physical optimization rules, UDFs, UDAFs, config extensions, etc...
+    /// custom extension codecs, custom physical optimization rules, UDFs, UDAFs, config
+    /// extensions, etc...
     ///
-    /// Example:
+    /// Example: adding some custom extension plan codecs
     ///
     /// ```rust
-    /// # use datafusion::execution::{FunctionRegistry, SessionStateBuilder};
-    /// # use datafusion_distributed::{SessionBuilder};
-    ///
-    /// #[derive(Clone)]
-    /// struct CustomSessionBuilder;
-    /// impl SessionBuilder for CustomSessionBuilder {
-    ///     fn on_new_session(&self, mut builder: SessionStateBuilder) -> SessionStateBuilder {
-    ///         // add your own UDFs, optimization rules, etc...
-    ///         builder
-    ///     }
-    /// }
-    /// ```
-    fn on_new_session(&self, builder: SessionStateBuilder) -> SessionStateBuilder {
-        builder
-    }
-
-    /// Allows users to provide their own codecs.
     ///
-    /// ```rust
     /// # use std::sync::Arc;
     /// # use datafusion::execution::runtime_env::RuntimeEnv;
     /// # use datafusion::execution::{FunctionRegistry, SessionStateBuilder};
     /// # use datafusion::physical_plan::ExecutionPlan;
     /// # use datafusion_proto::physical_plan::PhysicalExtensionCodec;
-    /// # use datafusion_distributed::{SessionBuilder};
+    /// # use datafusion_distributed::{with_user_codec, SessionBuilder};
     ///
     /// #[derive(Debug)]
     /// struct CustomExecCodec;
@@ -53,14 +34,13 @@ pub trait SessionBuilder {
     /// #[derive(Clone)]
     /// struct CustomSessionBuilder;
     /// impl SessionBuilder for CustomSessionBuilder {
-    ///     fn codec(&self) -> Option<Arc<dyn PhysicalExtensionCodec + 'static>> {
-    ///         Some(Arc::new(CustomExecCodec))
+    ///     fn on_new_session(&self, mut builder: SessionStateBuilder) -> SessionStateBuilder {
+    ///         // Add your UDFs, optimization rules, etc...
+    ///         with_user_codec(builder, CustomExecCodec)
     ///     }
     /// }
     /// ```
-    fn codec(&self) -> Option<Arc<dyn PhysicalExtensionCodec + 'static>> {
-        None
-    }
+    fn on_new_session(&self, builder: SessionStateBuilder) -> SessionStateBuilder;
 }
 
 /// Noop implementation of the [SessionBuilder]. Used by default if no [SessionBuilder] is provided

src/lib.rs

@@ -10,6 +10,9 @@ mod test_utils;
 pub mod physical_optimizer;
 pub mod stage;
 pub mod task;
+mod user_provided_codec;
+
 pub use channel_manager::{BoxCloneSyncChannel, ChannelManager, ChannelResolver};
 pub use flight_service::{ArrowFlightEndpoint, SessionBuilder};
 pub use plan::ArrowFlightReadExec;
+pub use user_provided_codec::{add_user_codec, get_user_codec, with_user_codec};

src/physical_optimizer.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use super::stage::ExecutionStage;
 use crate::{plan::PartitionIsolatorExec, ArrowFlightReadExec};
 use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::error::DataFusionError;
@@ -15,15 +16,9 @@ use datafusion::{
         displayable, repartition::RepartitionExec, ExecutionPlan, ExecutionPlanProperties,
     },
 };
-use datafusion_proto::physical_plan::PhysicalExtensionCodec;
-
-use super::stage::ExecutionStage;
 
 #[derive(Debug, Default)]
 pub struct DistributedPhysicalOptimizerRule {
-    /// Optional codec to assist in serializing and deserializing any custom
-    /// ExecutionPlan nodes
-    codec: Option<Arc<dyn PhysicalExtensionCodec>>,
     /// maximum number of partitions per task. This is used to determine how many
     /// tasks to create for each stage
     partitions_per_task: Option<usize>,
@@ -32,18 +27,10 @@ pub struct DistributedPhysicalOptimizerRule {
 impl DistributedPhysicalOptimizerRule {
     pub fn new() -> Self {
         DistributedPhysicalOptimizerRule {
-            codec: None,
             partitions_per_task: None,
         }
     }
 
-    /// Set a codec to use to assist in serializing and deserializing
-    /// custom ExecutionPlan nodes.
-    pub fn with_codec(mut self, codec: impl PhysicalExtensionCodec + 'static) -> Self {
-        self.codec = Some(Arc::new(codec));
-        self
-    }
-
     /// Set the maximum number of partitions per task. This is used to determine how many
     /// tasks to create for each stage.
     ///
@@ -156,9 +143,6 @@ impl DistributedPhysicalOptimizerRule {
         if let Some(partitions_per_task) = self.partitions_per_task {
             stage = stage.with_maximum_partitions_per_task(partitions_per_task);
         }
-        if let Some(codec) = self.codec.as_ref() {
-            stage = stage.with_user_codec(codec.clone());
-        }
         stage.depth = depth;
 
         Ok(stage)

src/plan/arrow_flight_read.rs

@@ -3,6 +3,7 @@ use crate::channel_manager::ChannelManager;
 use crate::composed_extension_codec::ComposedPhysicalExtensionCodec;
 use crate::errors::tonic_status_to_datafusion_error;
 use crate::flight_service::DoGet;
+use crate::get_user_codec;
 use crate::plan::DistributedCodec;
 use crate::stage::{proto_from_stage, ExecutionStage};
 use arrow_flight::decode::FlightRecordBatchStream;
@@ -174,7 +175,7 @@ impl ExecutionPlan for ArrowFlightReadExec {
 
         let mut combined_codec = ComposedPhysicalExtensionCodec::default();
         combined_codec.push(DistributedCodec {});
-        if let Some(ref user_codec) = stage.user_codec {
+        if let Some(ref user_codec) = get_user_codec(context.session_config()) {
             combined_codec.push_arc(Arc::clone(user_codec));
         }
 

src/stage/proto.rs

@@ -81,7 +81,6 @@ pub fn stage_from_proto(
         plan,
         inputs,
         tasks: msg.tasks,
-        user_codec: None,
         depth: 0,
     })
 }
@@ -147,7 +146,6 @@ mod tests {
             plan: physical_plan,
             inputs: vec![],
             tasks: vec![],
-            user_codec: Some(Arc::new(DefaultPhysicalExtensionCodec {})),
             depth: 0,
         };
 

src/stage/stage.rs

@@ -5,7 +5,6 @@ use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::TaskContext;
 use datafusion::physical_plan::{displayable, ExecutionPlan};
 use datafusion::prelude::SessionContext;
-use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 
 use itertools::Itertools;
 use rand::Rng;
@@ -32,8 +31,6 @@ pub struct ExecutionStage {
     /// Our tasks which tell us how finely grained to execute the partitions in
     /// the plan
     pub tasks: Vec<ExecutionTask>,
-    /// An optional codec to assist in serializing and deserializing this stage
-    pub user_codec: Option<Arc<dyn PhysicalExtensionCodec>>,
     /// tree depth of our location in the stage tree, used for display only
     pub depth: usize,
 }
@@ -65,7 +62,6 @@ impl ExecutionStage {
                 .map(|s| s as Arc<dyn ExecutionPlan>)
                 .collect(),
             tasks: vec![ExecutionTask::new(partition_group)],
-            user_codec: None,
             depth: 0,
         }
     }
@@ -93,13 +89,6 @@ impl ExecutionStage {
         self
     }
 
-    /// Sets the codec for this stage, which is used to serialize and deserialize the plan
-    /// and its inputs.
-    pub fn with_user_codec(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self {
-        self.user_codec = Some(codec);
-        self
-    }
-
     /// Returns the name of this stage
     pub fn name(&self) -> String {
         format!("Stage {:<3}", self.num)
@@ -174,7 +163,6 @@ impl ExecutionStage {
             plan: self.plan.clone(),
             inputs: assigned_children,
             tasks: assigned_tasks,
-            user_codec: self.user_codec.clone(),
             depth: self.depth,
         };
 
@@ -205,7 +193,6 @@ impl ExecutionPlan for ExecutionStage {
             plan: self.plan.clone(),
             inputs: children,
             tasks: self.tasks.clone(),
-            user_codec: self.user_codec.clone(),
             depth: self.depth,
         }))
     }

src/user_provided_codec.rs

@@ -0,0 +1,66 @@
+use datafusion::execution::SessionStateBuilder;
+use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use std::sync::Arc;
+
+pub struct UserProvidedCodec(Arc<dyn PhysicalExtensionCodec>);
+
+#[allow(private_bounds)]
+pub fn add_user_codec(
+    transport: &mut impl UserCodecTransport,
+    codec: impl PhysicalExtensionCodec + 'static,
+) {
+    transport.set(codec);
+}
+
+#[allow(private_bounds)]
+pub fn with_user_codec<T: UserCodecTransport>(
+    mut transport: T,
+    codec: impl PhysicalExtensionCodec + 'static,
+) -> T {
+    transport.set(codec);
+    transport
+}
+
+#[allow(private_bounds)]
+pub fn get_user_codec(
+    transport: &impl UserCodecTransport,
+) -> Option<Arc<dyn PhysicalExtensionCodec>> {
+    transport.get()
+}
+
+trait UserCodecTransport {
+    fn set(&mut self, codec: impl PhysicalExtensionCodec + 'static);
+    fn get(&self) -> Option<Arc<dyn PhysicalExtensionCodec>>;
+}
+
+impl UserCodecTransport for SessionConfig {
+    fn set(&mut self, codec: impl PhysicalExtensionCodec + 'static) {
+        self.set_extension(Arc::new(UserProvidedCodec(Arc::new(codec))));
+    }
+
+    fn get(&self) -> Option<Arc<dyn PhysicalExtensionCodec>> {
+        Some(Arc::clone(&self.get_extension::<UserProvidedCodec>()?.0))
+    }
+}
+
+impl UserCodecTransport for SessionContext {
+    fn set(&mut self, codec: impl PhysicalExtensionCodec + 'static) {
+        self.state_ref().write().config_mut().set(codec)
+    }
+
+    fn get(&self) -> Option<Arc<dyn PhysicalExtensionCodec>> {
+        self.state_ref().read().config().get()
+    }
+}
+
+impl UserCodecTransport for SessionStateBuilder {
+    fn set(&mut self, codec: impl PhysicalExtensionCodec + 'static) {
+        self.config().get_or_insert_default().set(codec);
+    }
+
+    fn get(&self) -> Option<Arc<dyn PhysicalExtensionCodec>> {
+        // Nobody will never want to retriever a user codec from a SessionStateBuilder
+        None
+    }
+}

tests/custom_extension_codec.rs

@@ -28,7 +28,9 @@ mod tests {
         displayable, execute_stream, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
     };
     use datafusion_distributed::physical_optimizer::DistributedPhysicalOptimizerRule;
-    use datafusion_distributed::{ArrowFlightReadExec, SessionBuilder};
+    use datafusion_distributed::{
+        add_user_codec, with_user_codec, ArrowFlightReadExec, SessionBuilder,
+    };
     use datafusion_proto::physical_plan::PhysicalExtensionCodec;
     use datafusion_proto::protobuf::proto_error;
     use futures::{stream, TryStreamExt};
@@ -40,17 +42,17 @@ mod tests {
     #[tokio::test]
     #[ignore]
     async fn custom_extension_codec() -> Result<(), Box<dyn std::error::Error>> {
-        // 1. The codec should be added to the extension builder.
         #[derive(Clone)]
         struct CustomSessionBuilder;
         impl SessionBuilder for CustomSessionBuilder {
-            fn codec(&self) -> Option<Arc<dyn PhysicalExtensionCodec>> {
-                Some(Arc::new(Int64ListExecCodec))
+            fn on_new_session(&self, builder: SessionStateBuilder) -> SessionStateBuilder {
+                with_user_codec(builder, Int64ListExecCodec)
             }
         }
 
-        let (ctx, _guard) =
+        let (mut ctx, _guard) =
             start_localhost_context([50050, 50051, 50052], CustomSessionBuilder).await;
+        add_user_codec(&mut ctx, Int64ListExecCodec);
 
         let single_node_plan = build_plan(false)?;
         assert_snapshot!(displayable(single_node_plan.as_ref()).indent(true).to_string(), @r"
@@ -60,10 +62,8 @@ mod tests {
         ");
 
         let distributed_plan = build_plan(true)?;
-        let distributed_plan = DistributedPhysicalOptimizerRule::default()
-            // 1. The codec should be added to the DistributedPhysicalOptimizerRule.
-            .with_codec(Int64ListExecCodec)
-            .distribute_plan(distributed_plan)?;
+        let distributed_plan =
+            DistributedPhysicalOptimizerRule::default().distribute_plan(distributed_plan)?;
 
         assert_snapshot!(displayable(&distributed_plan).indent(true).to_string(), @r"
         ┌───── Stage 3   Task: partitions: 0,unassigned]

tests/error_propagation.rs

@@ -16,7 +16,9 @@ mod tests {
         execute_stream, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
     };
     use datafusion_distributed::physical_optimizer::DistributedPhysicalOptimizerRule;
-    use datafusion_distributed::{ArrowFlightReadExec, SessionBuilder};
+    use datafusion_distributed::{
+        add_user_codec, with_user_codec, ArrowFlightReadExec, SessionBuilder,
+    };
     use datafusion_proto::physical_plan::PhysicalExtensionCodec;
     use datafusion_proto::protobuf::proto_error;
     use futures::{stream, TryStreamExt};
@@ -27,26 +29,17 @@ mod tests {
     use std::sync::Arc;
 
     #[tokio::test]
-    #[ignore]
     async fn test_error_propagation() -> Result<(), Box<dyn Error>> {
         #[derive(Clone)]
         struct CustomSessionBuilder;
         impl SessionBuilder for CustomSessionBuilder {
-            fn on_new_session(&self, mut builder: SessionStateBuilder) -> SessionStateBuilder {
-                let codec: Arc<dyn PhysicalExtensionCodec> = Arc::new(ErrorExecCodec);
-                let config = builder.config().get_or_insert_default();
-                config.set_extension(Arc::new(codec));
-                builder
+            fn on_new_session(&self, builder: SessionStateBuilder) -> SessionStateBuilder {
+                with_user_codec(builder, ErrorExecCodec)
             }
         }
-        let (ctx, _guard) =
+        let (mut ctx, _guard) =
             start_localhost_context([50050, 50051, 50053], CustomSessionBuilder).await;
-
-        let codec: Arc<dyn PhysicalExtensionCodec> = Arc::new(ErrorExecCodec);
-        ctx.state_ref()
-            .write()
-            .config_mut()
-            .set_extension(Arc::new(codec));
+        add_user_codec(&mut ctx, ErrorExecCodec);
 
         let mut plan: Arc<dyn ExecutionPlan> = Arc::new(ErrorExec::new("something failed"));