Make DistributedExec handle udfs

gabotechs · gabotechs · commit 7a881d349d50 · 2025-10-22T13:08:44.000+02:00
diff --git a/src/execution_plans/distributed.rs b/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(context.session_config());
+        let codec = DistributedCodec::new_combined_with_user(Arc::clone(&context));
 
         let prepared = self.prepare_plan(&channel_resolver.get_urls()?, &codec)?;
         {
diff --git a/src/flight_service/do_get.rs b/src/flight_service/do_get.rs
@@ -81,8 +81,7 @@ impl ArrowFlightEndpoint {
             .await
             .map_err(|err| datafusion_error_to_tonic_status(&err))?;
 
-        let codec = DistributedCodec::new_combined_with_user(session_state.config());
-        let ctx = SessionContext::new_with_state(session_state.clone());
+        let codec = DistributedCodec::new_combined_with_user(session_state.task_ctx());
 
         // 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.
@@ -94,6 +93,7 @@ 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
diff --git a/src/protobuf/distributed_codec.rs b/src/protobuf/distributed_codec.rs
@@ -6,32 +6,40 @@ use crate::{NetworkShuffleExec, PartitionIsolatorExec};
 use bytes::Bytes;
 use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::common::{internal_datafusion_err, not_impl_err};
+use datafusion::common::internal_datafusion_err;
 use datafusion::error::DataFusionError;
-use datafusion::execution::FunctionRegistry;
+use datafusion::execution::{FunctionRegistry, TaskContext};
 use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 use datafusion::physical_expr::EquivalenceProperties;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::{ExecutionPlan, Partitioning, PlanProperties};
-use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion::prelude::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(Debug)]
-pub struct DistributedCodec;
+#[derive(Clone, Default)]
+pub struct DistributedCodec(Arc<TaskContext>);
+
+impl Debug for DistributedCodec {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "DistributedCodec")
+    }
+}
 
 impl DistributedCodec {
-    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));
+    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()));
         ComposedPhysicalExtensionCodec::new(codecs)
     }
 }
@@ -104,13 +112,9 @@ 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,
-                    &DistributedCodec {},
-                )?
-                .ok_or(proto_error("NetworkShuffleExec is missing partitioning"))?;
+                let partitioning =
+                    parse_protobuf_partitioning(partitioning.as_ref(), &ctx, &schema, self)?
+                        .ok_or(proto_error("NetworkShuffleExec is missing partitioning"))?;
 
                 Ok(Arc::new(new_network_hash_shuffle_exec(
                     partitioning,
@@ -128,13 +132,9 @@ 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,
-                    &DistributedCodec {},
-                )?
-                .ok_or(proto_error("NetworkCoalesceExec is missing partitioning"))?;
+                let partitioning =
+                    parse_protobuf_partitioning(partitioning.as_ref(), &ctx, &schema, self)?
+                        .ok_or(proto_error("NetworkCoalesceExec is missing partitioning"))?;
 
                 Ok(Arc::new(new_network_coalesce_tasks_exec(
                     partitioning,
@@ -185,7 +185,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
                 schema: Some(node.schema().try_into()?),
                 partitioning: Some(serialize_partitioning(
                     node.properties().output_partitioning(),
-                    &DistributedCodec {},
+                    self,
                 )?),
                 input_stage: Some(encode_stage_proto(node.input_stage())?),
             };
@@ -200,7 +200,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
                 schema: Some(node.schema().try_into()?),
                 partitioning: Some(serialize_partitioning(
                     node.properties().output_partitioning(),
-                    &DistributedCodec {},
+                    self,
                 )?),
                 input_stage: Some(encode_stage_proto(node.input_stage())?),
             };
@@ -230,16 +230,28 @@ impl PhysicalExtensionCodec for DistributedCodec {
         }
     }
 
-    fn try_encode_udf(&self, _: &ScalarUDF, _: &mut Vec<u8>) -> datafusion::common::Result<()> {
-        not_impl_err!("DistributedCodec does not encode UDFs")
+    fn try_decode_udf(
+        &self,
+        name: &str,
+        _buf: &[u8],
+    ) -> datafusion::common::Result<Arc<ScalarUDF>> {
+        self.0.udf(name)
     }
 
-    fn try_encode_udaf(&self, _: &AggregateUDF, _: &mut Vec<u8>) -> datafusion::common::Result<()> {
-        not_impl_err!("DistributedCodec does not encode UDAFs")
+    fn try_decode_udaf(
+        &self,
+        name: &str,
+        _buf: &[u8],
+    ) -> datafusion::common::Result<Arc<AggregateUDF>> {
+        self.0.udaf(name)
     }
 
-    fn try_encode_udwf(&self, _: &WindowUDF, _: &mut Vec<u8>) -> datafusion::common::Result<()> {
-        not_impl_err!("DistributedCodec does not encode UDWFs")
+    fn try_decode_udwf(
+        &self,
+        name: &str,
+        _buf: &[u8],
+    ) -> datafusion::common::Result<Arc<WindowUDF>> {
+        self.0.udwf(name)
     }
 }
 
@@ -424,7 +436,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_single_flight() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec;
+        let codec = DistributedCodec::default();
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("a");
@@ -443,7 +455,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_isolator_flight() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec;
+        let codec = DistributedCodec::default();
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("b");
@@ -467,7 +479,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_isolator_union() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec;
+        let codec = DistributedCodec::default();
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("c");
@@ -497,7 +509,7 @@ mod tests {
 
     #[test]
     fn test_roundtrip_isolator_sort_flight() -> datafusion::common::Result<()> {
-        let codec = DistributedCodec;
+        let codec = DistributedCodec::default();
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("d");
diff --git a/tests/udfs.rs b/tests/udfs.rs
@@ -0,0 +1,163 @@
+#[cfg(all(feature = "integration", test))]
+mod tests {
+    use arrow::util::pretty::pretty_format_batches;
+    use datafusion::arrow::datatypes::DataType;
+    use datafusion::common::not_impl_err;
+    use datafusion::error::DataFusionError;
+    use datafusion::execution::{FunctionRegistry, SessionState, SessionStateBuilder};
+    use datafusion::logical_expr::{
+        AggregateUDF, ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
+        Volatility, WindowUDF,
+    };
+    use datafusion::logical_expr_common::dyn_eq::DynEq;
+    use datafusion::physical_optimizer::PhysicalOptimizerRule;
+    use datafusion::physical_plan::{ExecutionPlan, execute_stream};
+    use datafusion_distributed::test_utils::localhost::start_localhost_context;
+    use datafusion_distributed::test_utils::parquet::register_parquet_tables;
+    use datafusion_distributed::{
+        DistributedExt, DistributedPhysicalOptimizerRule, DistributedSessionBuilderContext,
+        assert_snapshot,
+    };
+    use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+    use futures::TryStreamExt;
+    use std::any::Any;
+    use std::error::Error;
+    use std::sync::Arc;
+
+    #[tokio::test]
+    async fn test_udfs() -> Result<(), Box<dyn Error>> {
+        async fn build_state(
+            ctx: DistributedSessionBuilderContext,
+        ) -> Result<SessionState, DataFusionError> {
+            Ok(SessionStateBuilder::new()
+                .with_runtime_env(ctx.runtime_env)
+                .with_default_features()
+                .with_distributed_user_codec(UdfCodec)
+                .build())
+        }
+
+        let (ctx, _guard) = start_localhost_context(3, build_state).await;
+        register_parquet_tables(&ctx).await?;
+
+        let df = ctx
+            .sql(r#"SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)"#)
+            .await?;
+        let physical = df.create_physical_plan().await?;
+
+        let physical_distributed = DistributedPhysicalOptimizerRule::default()
+            .with_network_shuffle_tasks(2)
+            .with_network_coalesce_tasks(2)
+            .optimize(physical.clone(), &Default::default())?;
+
+        let batches = pretty_format_batches(
+            &execute_stream(physical_distributed, ctx.task_ctx())?
+                .try_collect::<Vec<_>>()
+                .await?,
+        )?;
+
+        assert_snapshot!(batches, @r"
+        +----------+-----------+
+        | count(*) | RainToday |
+        +----------+-----------+
+        | 66       | Yes       |
+        | 300      | No        |
+        +----------+-----------+
+        ");
+        Ok(())
+    }
+
+    #[derive(Debug, PartialEq, Eq, Hash)]
+    pub struct Udf(Signature);
+
+    impl Udf {
+        fn new() -> Self {
+            Self(Signature::any(1, Volatility::Immutable))
+        }
+    }
+
+    impl ScalarUDFImpl for Udf {
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        fn name(&self) -> &str {
+            "test_udf"
+        }
+
+        fn signature(&self) -> &Signature {
+            &self.0
+        }
+
+        fn return_type(&self, arg_types: &[DataType]) -> datafusion::common::Result<DataType> {
+            Ok(arg_types[0].clone())
+        }
+
+        fn invoke_with_args(
+            &self,
+            mut args: ScalarFunctionArgs,
+        ) -> datafusion::common::Result<ColumnarValue> {
+            Ok(args.args.remove(0))
+        }
+    }
+
+    #[derive(Debug)]
+    struct UdfCodec;
+
+    impl PhysicalExtensionCodec for UdfCodec {
+        fn try_decode(
+            &self,
+            _: &[u8],
+            _: &[Arc<dyn ExecutionPlan>],
+            _registry: &dyn FunctionRegistry,
+        ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+            not_impl_err!("not implemented")
+        }
+
+        fn try_encode(
+            &self,
+            _: Arc<dyn ExecutionPlan>,
+            _: &mut Vec<u8>,
+        ) -> datafusion::common::Result<()> {
+            not_impl_err!("not implemented")
+        }
+
+        fn try_encode_udaf(
+            &self,
+            _: &AggregateUDF,
+            _: &mut Vec<u8>,
+        ) -> datafusion::common::Result<()> {
+            not_impl_err!("not implemented")
+        }
+
+        fn try_encode_udwf(
+            &self,
+            _: &WindowUDF,
+            _: &mut Vec<u8>,
+        ) -> datafusion::common::Result<()> {
+            not_impl_err!("not implemented")
+        }
+
+        fn try_encode_udf(
+            &self,
+            node: &ScalarUDF,
+            _: &mut Vec<u8>,
+        ) -> datafusion::common::Result<()> {
+            if node.dyn_eq(node) {
+                return Ok(());
+            };
+            not_impl_err!("not implemented")
+        }
+
+        fn try_decode_udf(
+            &self,
+            name: &str,
+            _: &[u8],
+        ) -> datafusion::common::Result<Arc<ScalarUDF>> {
+            let udf = Udf::new();
+            if name == udf.name() {
+                return Ok(Arc::new(ScalarUDF::new_from_impl(udf)));
+            }
+            not_impl_err!("not implemented")
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -138,7 +138,7 @@ impl ExecutionPlan for DistributedExec {`
`138`	`138`	`}`
`139`	`139`
`140`	`140`	`let channel_resolver = get_distributed_channel_resolver(context.session_config())?;`
`141`		`- let codec = DistributedCodec::new_combined_with_user(context.session_config());`
	`141`	`+ let codec = DistributedCodec::new_combined_with_user(Arc::clone(&context));`
`142`	`142`
`143`	`143`	`let prepared = self.prepare_plan(&channel_resolver.get_urls()?, &codec)?;`
`144`	`144`	`{`