Allow passing custom codecs

Author: gabotechs

src/composed_extension_codec.rs

@@ -3,7 +3,6 @@ use datafusion::error::DataFusionError;
 use datafusion::execution::FunctionRegistry;
 use datafusion::logical_expr::{AggregateUDF, ScalarUDF};
 use datafusion::physical_plan::ExecutionPlan;
-use datafusion::prelude::SessionConfig;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use std::fmt::Debug;
 use std::sync::Arc;
@@ -25,42 +24,10 @@ impl ComposedPhysicalExtensionCodec {
         self.codecs.push(Arc::new(codec));
     }
 
-    /// Adds a new [PhysicalExtensionCodec] from DataFusion's [SessionConfig] extensions.
-    ///
-    /// If users have a custom [PhysicalExtensionCodec] for their own nodes, they should
-    /// populate the config extensions with a [PhysicalExtensionCodec] so that we can use
-    /// it while encoding/decoding plans to/from protobuf.
-    ///
-    /// Example:
-    /// ```rust
-    /// # use std::sync::Arc;
-    /// # use datafusion::execution::FunctionRegistry;
-    /// # use datafusion::physical_plan::ExecutionPlan;
-    /// # use datafusion::prelude::SessionConfig;
-    /// # use datafusion_proto::physical_plan::PhysicalExtensionCodec;
-    ///
-    /// #[derive(Debug)]
-    /// struct CustomUserCodec {}
-    ///
-    /// impl PhysicalExtensionCodec for CustomUserCodec {
-    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], registry: &dyn FunctionRegistry) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
-    ///         todo!()
-    ///     }
-    ///
-    ///     fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> datafusion::common::Result<()> {
-    ///         todo!()
-    ///     }
-    /// }
-    ///
-    /// let mut config = SessionConfig::new();
-    ///
-    /// let codec: Arc<dyn PhysicalExtensionCodec> = Arc::new(CustomUserCodec {});
-    /// config.set_extension(Arc::new(codec));
-    /// ```
-    pub(crate) fn push_from_config(&mut self, config: &SessionConfig) {
-        if let Some(user_codec) = config.get_extension::<Arc<dyn PhysicalExtensionCodec>>() {
-            self.codecs.push(user_codec.as_ref().clone());
-        }
+    /// Adds a new [PhysicalExtensionCodec] to the list. These codecs will be tried
+    /// sequentially until one works.
+    pub(crate) fn push_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>) {
+        self.codecs.push(codec);
     }
 
     fn try_any<T>(

src/flight_service/do_get.rs

@@ -1,3 +1,4 @@
+use crate::composed_extension_codec::ComposedPhysicalExtensionCodec;
 use crate::errors::datafusion_error_to_tonic_status;
 use crate::flight_service::service::ArrowFlightEndpoint;
 use crate::plan::DistributedCodec;
@@ -8,7 +9,6 @@ use arrow_flight::flight_service_server::FlightService;
 use arrow_flight::Ticket;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::optimizer::OptimizerConfig;
-use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use futures::TryStreamExt;
 use prost::Message;
 use std::sync::Arc;
@@ -48,25 +48,36 @@ impl ArrowFlightEndpoint {
             "FunctionRegistry not present in newly built SessionState",
         ))?;
 
-        let codec = DistributedCodec {};
-        let codec = Arc::new(codec) as Arc<dyn PhysicalExtensionCodec>;
+        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 {
+            combined_codec.push_arc(Arc::clone(&user_codec));
+        }
 
-        let stage = stage_from_proto(stage_msg, function_registry, &self.runtime.as_ref(), codec)
-            .map(Arc::new)
-            .map_err(|err| Status::invalid_argument(format!("Cannot decode stage proto: {err}")))?;
+        let mut stage = stage_from_proto(
+            stage_msg,
+            function_registry,
+            &self.runtime.as_ref(),
+            &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
         let config = state.config_mut();
         config.set_extension(Arc::clone(&self.channel_manager));
-        config.set_extension(stage.clone());
+        config.set_extension(Arc::new(stage));
 
-        let stream = stage
-            .plan
+        let stream = inner_plan
             .execute(doget.partition as usize, state.task_ctx())
             .map_err(|err| Status::internal(format!("Error executing stage plan: {err:#?}")))?;
 
         let flight_data_stream = FlightDataEncoderBuilder::new()
-            .with_schema(stage.plan.schema().clone())
+            .with_schema(inner_plan.schema().clone())
             .build(stream.map_err(|err| {
                 FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(&err)))
             }));

src/flight_service/session_builder.rs

@@ -1,16 +1,35 @@
 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 extension codecs, custom physical optimization rules, UDFs, UDAFs, config
-    /// extensions, etc...
+    /// custom physical optimization rules, UDFs, UDAFs, config extensions, etc...
     ///
-    /// Example: adding some custom extension plan codecs
+    /// Example:
     ///
     /// ```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};
@@ -34,15 +53,14 @@ pub trait SessionBuilder {
     /// #[derive(Clone)]
     /// struct CustomSessionBuilder;
     /// impl SessionBuilder for CustomSessionBuilder {
-    ///     fn on_new_session(&self, mut builder: SessionStateBuilder) -> SessionStateBuilder {
-    ///         let config = builder.config().get_or_insert_default();
-    ///         let codec: Arc<dyn PhysicalExtensionCodec> = Arc::new(CustomExecCodec);
-    ///         config.set_extension(Arc::new(codec));
-    ///         builder
+    ///     fn codec(&self) -> Option<Arc<dyn PhysicalExtensionCodec + 'static>> {
+    ///         Some(Arc::new(CustomExecCodec))
     ///     }
     /// }
     /// ```
-    fn on_new_session(&self, builder: SessionStateBuilder) -> SessionStateBuilder;
+    fn codec(&self) -> Option<Arc<dyn PhysicalExtensionCodec + 'static>> {
+        None
+    }
 }
 
 /// Noop implementation of the [SessionBuilder]. Used by default if no [SessionBuilder] is provided

src/physical_optimizer.rs

@@ -39,8 +39,8 @@ impl DistributedPhysicalOptimizerRule {
 
     /// Set a codec to use to assist in serializing and deserializing
     /// custom ExecutionPlan nodes.
-    pub fn with_codec(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self {
-        self.codec = Some(codec);
+    pub fn with_codec(mut self, codec: impl PhysicalExtensionCodec + 'static) -> Self {
+        self.codec = Some(Arc::new(codec));
         self
     }
 
@@ -157,7 +157,7 @@ impl DistributedPhysicalOptimizerRule {
             stage = stage.with_maximum_partitions_per_task(partitions_per_task);
         }
         if let Some(codec) = self.codec.as_ref() {
-            stage = stage.with_codec(codec.clone());
+            stage = stage.with_user_codec(codec.clone());
         }
         stage.depth = depth;
 

src/plan/arrow_flight_read.rs

@@ -1,8 +1,10 @@
 use super::combined::CombinedRecordBatchStream;
 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::stage::{ExecutionStage, ExecutionStageProto};
+use crate::plan::DistributedCodec;
+use crate::stage::{proto_from_stage, ExecutionStage};
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
@@ -170,12 +172,14 @@ impl ExecutionPlan for ArrowFlightReadExec {
                 this.stage_num
             ))?;
 
-        let child_stage_tasks = child_stage.tasks.clone();
-        let child_stage_proto = ExecutionStageProto::try_from(child_stage).map_err(|e| {
-            internal_datafusion_err!(
-                "ArrowFlightReadExec: failed to convert stage to proto: {}",
-                e
-            )
+        let mut combined_codec = ComposedPhysicalExtensionCodec::default();
+        combined_codec.push(DistributedCodec {});
+        if let Some(ref user_codec) = stage.user_codec {
+            combined_codec.push_arc(Arc::clone(user_codec));
+        }
+
+        let child_stage_proto = proto_from_stage(child_stage, &combined_codec).map_err(|e| {
+            internal_datafusion_err!("ArrowFlightReadExec: failed to convert stage to proto: {e}")
         })?;
 
         let ticket_bytes = DoGet {
@@ -191,6 +195,7 @@ impl ExecutionPlan for ArrowFlightReadExec {
 
         let schema = child_stage.plan.schema();
 
+        let child_stage_tasks = child_stage.tasks.clone();
         let stream = async move {
             let futs = child_stage_tasks.iter().enumerate().map(|(i, task)| async {
                 let url = task.url()?.ok_or(internal_datafusion_err!(

src/stage/mod.rs

@@ -3,5 +3,5 @@ mod proto;
 mod stage;
 
 pub use display::display_stage_graphviz;
-pub use proto::{stage_from_proto, ExecutionStageProto};
+pub use proto::{proto_from_stage, stage_from_proto, ExecutionStageProto};
 pub use stage::ExecutionStage;

src/stage/proto.rs

@@ -11,7 +11,7 @@ use datafusion_proto::{
     protobuf::PhysicalPlanNode,
 };
 
-use crate::{plan::DistributedCodec, task::ExecutionTask};
+use crate::task::ExecutionTask;
 
 use super::ExecutionStage;
 
@@ -35,54 +35,42 @@ pub struct ExecutionStageProto {
     pub tasks: Vec<ExecutionTask>,
 }
 
-impl TryFrom<&ExecutionStage> for ExecutionStageProto {
-    type Error = DataFusionError;
-
-    fn try_from(stage: &ExecutionStage) -> Result<Self, Self::Error> {
-        let codec = stage.codec.clone().unwrap_or(Arc::new(DistributedCodec {}));
-
-        let proto_plan =
-            PhysicalPlanNode::try_from_physical_plan(stage.plan.clone(), codec.as_ref())?;
-        let inputs = stage
-            .child_stages_iter()
-            .map(|s| Box::new(ExecutionStageProto::try_from(s).unwrap()))
-            .collect();
-
-        Ok(ExecutionStageProto {
-            num: stage.num as u64,
-            name: stage.name(),
-            plan: Some(Box::new(proto_plan)),
-            inputs,
-            tasks: stage.tasks.clone(),
-        })
-    }
-}
-
-impl TryFrom<ExecutionStage> for ExecutionStageProto {
-    type Error = DataFusionError;
+pub fn proto_from_stage(
+    stage: &ExecutionStage,
+    codec: &dyn PhysicalExtensionCodec,
+) -> Result<ExecutionStageProto, DataFusionError> {
+    let proto_plan = PhysicalPlanNode::try_from_physical_plan(stage.plan.clone(), codec)?;
+    let inputs = stage
+        .child_stages_iter()
+        .map(|s| Ok(Box::new(proto_from_stage(s, codec)?)))
+        .collect::<Result<Vec<_>>>()?;
 
-    fn try_from(stage: ExecutionStage) -> Result<Self, Self::Error> {
-        ExecutionStageProto::try_from(&stage)
-    }
+    Ok(ExecutionStageProto {
+        num: stage.num as u64,
+        name: stage.name(),
+        plan: Some(Box::new(proto_plan)),
+        inputs,
+        tasks: stage.tasks.clone(),
+    })
 }
 
 pub fn stage_from_proto(
     msg: ExecutionStageProto,
     registry: &dyn FunctionRegistry,
     runtime: &RuntimeEnv,
-    codec: Arc<dyn PhysicalExtensionCodec>,
+    codec: &dyn PhysicalExtensionCodec,
 ) -> Result<ExecutionStage> {
     let plan_node = msg.plan.ok_or(internal_datafusion_err!(
         "ExecutionStageMsg is missing the plan"
     ))?;
 
-    let plan = plan_node.try_into_physical_plan(registry, runtime, codec.as_ref())?;
+    let plan = plan_node.try_into_physical_plan(registry, runtime, codec)?;
 
     let inputs = msg
         .inputs
         .into_iter()
         .map(|s| {
-            stage_from_proto(*s, registry, runtime, codec.clone())
+            stage_from_proto(*s, registry, runtime, codec)
                 .map(|s| Arc::new(s) as Arc<dyn ExecutionPlan>)
         })
         .collect::<Result<Vec<_>>>()?;
@@ -93,7 +81,7 @@ pub fn stage_from_proto(
         plan,
         inputs,
         tasks: msg.tasks,
-        codec: Some(codec),
+        user_codec: None,
         depth: 0,
     })
 }
@@ -116,6 +104,7 @@ mod tests {
     use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
     use prost::Message;
 
+    use crate::stage::proto::proto_from_stage;
     use crate::stage::{proto::stage_from_proto, ExecutionStage, ExecutionStageProto};
 
     // create a simple mem table
@@ -158,12 +147,12 @@ mod tests {
             plan: physical_plan,
             inputs: vec![],
             tasks: vec![],
-            codec: Some(Arc::new(DefaultPhysicalExtensionCodec {})),
+            user_codec: Some(Arc::new(DefaultPhysicalExtensionCodec {})),
             depth: 0,
         };
 
         // Convert to proto message
-        let stage_msg = ExecutionStageProto::try_from(&stage)?;
+        let stage_msg = proto_from_stage(&stage, &DefaultPhysicalExtensionCodec {})?;
 
         // Serialize to bytes
         let mut buf = Vec::new();
@@ -180,7 +169,7 @@ mod tests {
             decoded_msg,
             &ctx,
             ctx.runtime_env().as_ref(),
-            Arc::new(DefaultPhysicalExtensionCodec {}),
+            &DefaultPhysicalExtensionCodec {},
         )?;
 
         // Compare original and round-tripped stages