update PyWindowUDF

Ref: apache/datafusion#12803
Michael-J-Ward · Oct 15, 2024 · 7255ec2 · 7255ec2
1 parent df9197b
commit 7255ec2
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/src/udwf.rs b/src/udwf.rs
@@ -20,6 +20,7 @@ use std::ops::Range;
 use std::sync::Arc;
 
 use arrow::array::{make_array, Array, ArrayData, ArrayRef};
+use datafusion::logical_expr::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs};
 use datafusion::logical_expr::window_state::WindowAggState;
 use datafusion::scalar::ScalarValue;
 use pyo3::exceptions::PyValueError;
@@ -299,11 +300,21 @@ impl WindowUDFImpl for MultiColumnWindowUDF {
         &self.signature
     }
 
-    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-        Ok(self.return_type.clone())
+    fn field(&self, _field_args: WindowUDFFieldArgs) -> Result<arrow::datatypes::Field> {
+        // TODO: Should nulalble always be `true`?
+        Ok(arrow::datatypes::Field::new(
+            self.name(),
+            self.return_type.clone(),
+            true,
+        ))
     }
 
-    fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
+    // TODO: Enable passing partition_evaluator_args to python?
+    fn partition_evaluator(
+        &self,
+        _partition_evaluator_args: PartitionEvaluatorArgs,
+    ) -> Result<Box<dyn PartitionEvaluator>> {
+        let _ = _partition_evaluator_args;
         (self.partition_evaluator_factory)()
     }
 }