fix: serialize user-defined window functions to proto

datafusion/physical-plan/src/windows/mod.rs

@@ -194,7 +194,7 @@ pub fn create_udwf_window_expr(
 
 /// Implements [`BuiltInWindowFunctionExpr`] for [`WindowUDF`]
 #[derive(Clone, Debug)]
-struct WindowUDFExpr {
+pub struct WindowUDFExpr {
  fun: Arc<WindowUDF>,
  args: Vec<Arc<dyn PhysicalExpr>>,
  /// Display name
@@ -209,6 +209,12 @@ struct WindowUDFExpr {
  ignore_nulls: bool,
 }
 
+impl WindowUDFExpr {
+ pub fn fun(&self) -> &Arc<WindowUDF> {
+ &self.fun
+ }
+}
+
 impl BuiltInWindowFunctionExpr for WindowUDFExpr {
  fn as_any(&self) -> &dyn std::any::Any {
  self

datafusion/proto/proto/datafusion.proto

@@ -853,6 +853,7 @@ message PhysicalWindowExprNode {
  oneof window_function {
  // BuiltInWindowFunction built_in_function = 2;
  string user_defined_aggr_function = 3;
+ string user_defined_window_function = 10;
  }
  repeated PhysicalExprNode args = 4;
  repeated PhysicalExprNode partition_by = 5;

datafusion/proto/src/physical_plan/from_proto.rs

@@ -152,6 +152,12 @@ pub fn parse_physical_window_expr(
  None => registry.udaf(udaf_name)?
  })
  }
+ protobuf::physical_window_expr_node::WindowFunction::UserDefinedWindowFunction(udwf_name) => {
+ WindowFunctionDefinition::WindowUDF(match &proto.fun_definition {
+ Some(buf) => codec.try_decode_udwf(udwf_name, buf)?,
+ None => registry.udwf(udwf_name)?
+ })
+ }
  }
  } else {
  return Err(proto_error("Missing required field in protobuf"));

datafusion/proto/src/physical_plan/mod.rs

@@ -64,7 +64,7 @@ use datafusion::physical_plan::{
  ExecutionPlan, InputOrderMode, PhysicalExpr, WindowExpr,
 };
 use datafusion_common::{internal_err, not_impl_err, DataFusionError, Result};
-use datafusion_expr::{AggregateUDF, ScalarUDF};
+use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 
 use crate::common::{byte_to_string, str_to_byte};
 use crate::physical_plan::from_proto::{
@@ -2119,6 +2119,14 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync {
  fn try_encode_udaf(&self, _node: &AggregateUDF, _buf: &mut Vec<u8>) -> Result<()> {
  Ok(())
  }
+
+ fn try_decode_udwf(&self, name: &str, _buf: &[u8]) -> Result<Arc<WindowUDF>> {
+ not_impl_err!("PhysicalExtensionCodec is not provided for window function {name}")
+ }
+
+ fn try_encode_udwf(&self, _node: &WindowUDF, _buf: &mut Vec<u8>) -> Result<()> {
+ Ok(())
+ }
 }
 
 #[derive(Debug)]

datafusion/proto/src/physical_plan/to_proto.rs

@@ -19,14 +19,14 @@ use std::sync::Arc;
 
 #[cfg(feature = "parquet")]
 use datafusion::datasource::file_format::parquet::ParquetSink;
-use datafusion::physical_expr::window::SlidingAggregateWindowExpr;
+use datafusion::physical_expr::window::{BuiltInWindowExpr, SlidingAggregateWindowExpr};
 use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr, ScalarFunctionExpr};
 use datafusion::physical_plan::expressions::{
  BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr,
  Literal, NegativeExpr, NotExpr, TryCastExpr,
 };
 use datafusion::physical_plan::udaf::AggregateFunctionExpr;
-use datafusion::physical_plan::windows::PlainAggregateWindowExpr;
+use datafusion::physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr};
 use datafusion::physical_plan::{Partitioning, PhysicalExpr, WindowExpr};
 use datafusion::{
  datasource::{
@@ -68,7 +68,7 @@ pub fn serialize_physical_aggr_expr(
  ordering_req,
  distinct: aggr_expr.is_distinct(),
  ignore_nulls: aggr_expr.ignore_nulls(),
- fun_definition: (!buf.is_empty()).then_some(buf)
+ fun_definition: (!buf.is_empty()).then_some(buf),
  },
  )),
  })
@@ -120,6 +120,25 @@ pub fn serialize_physical_window_expr(
  window_frame,
  codec,
  )?
+ } else if let Some(built_in_window_expr) = expr.downcast_ref::<BuiltInWindowExpr>() {
+ if let Some(expr) = built_in_window_expr
+ .get_built_in_func_expr()
+ .as_any()
+ .downcast_ref::<WindowUDFExpr>()
+ {
+ let mut buf = Vec::new();
+ codec.try_encode_udwf(expr.fun(), &mut buf)?;
+ (
+ physical_window_expr_node::WindowFunction::UserDefinedWindowFunction(
+ expr.fun().name().to_string(),
+ ),
+ (!buf.is_empty()).then_some(buf),
+ )
+ } else {
+ return not_impl_err!(
+ "User-defined window function not supported: {window_expr:?}"
+ );
+ }
  } else {
  return not_impl_err!("WindowExpr not supported: {window_expr:?}");
  };

datafusion/proto/tests/cases/mod.rs

@@ -15,15 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::datatypes::{DataType, Field};
 use std::any::Any;
-
-use arrow::datatypes::DataType;
+use std::fmt::Debug;
 
 use datafusion_common::plan_err;
 use datafusion_expr::function::AccumulatorArgs;
 use datafusion_expr::{
- Accumulator, AggregateUDFImpl, ColumnarValue, ScalarUDFImpl, Signature, Volatility,
+ Accumulator, AggregateUDFImpl, ColumnarValue, PartitionEvaluator, ScalarUDFImpl,
+ Signature, Volatility, WindowUDFImpl,
 };
+use datafusion_functions_window_common::field::WindowUDFFieldArgs;
+use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
 
 mod roundtrip_logical_plan;
 mod roundtrip_physical_plan;
@@ -125,3 +128,54 @@ pub struct MyAggregateUdfNode {
  #[prost(string, tag = "1")]
  pub result: String,
 }
+
+#[derive(Debug)]
+pub(in crate::cases) struct CustomUDWF {
+ signature: Signature,
+ payload: String,
+}
+
+impl CustomUDWF {
+ pub fn new(payload: String) -> Self {
+ Self {
+ signature: Signature::exact(vec![DataType::Int64], Volatility::Immutable),
+ payload,
+ }
+ }
+}
+
+impl WindowUDFImpl for CustomUDWF {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "custom_udwf"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn partition_evaluator(
+ &self,
+ _partition_evaluator_args: PartitionEvaluatorArgs,
+ ) -> datafusion_common::Result<Box<dyn PartitionEvaluator>> {
+ Ok(Box::new(CustomUDWFEvaluator {}))
+ }
+
+ fn field(&self, field_args: WindowUDFFieldArgs) -> datafusion_common::Result<Field> {
+ Ok(Field::new(field_args.name(), DataType::UInt64, false))
+ }
+}
+
+#[derive(Debug)]
+struct CustomUDWFEvaluator;
+
+impl PartitionEvaluator for CustomUDWFEvaluator {}
+
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub(in crate::cases) struct CustomUDWFNode {
+ #[prost(string, tag = "1")]
+ pub payload: String,
+}