Updates

Author: alamb

datafusion-examples/examples/simple_udwf.rs

@@ -17,6 +17,10 @@
 
 use std::sync::Arc;
 
+use arrow::{
+    array::{AsArray, Float64Array},
+    datatypes::Float64Type,
+};
 use arrow_schema::DataType;
 use datafusion::datasource::file_format::options::CsvReadOptions;
 
@@ -113,26 +117,19 @@ fn make_partition_evaluator() -> Result<Box<dyn PartitionEvaluator>> {
     Ok(Box::new(MyPartitionEvaluator::new()))
 }
 
-
-
 /// This implements the lowest level evaluation for a window function
 ///
 /// It handles calculating the value of the window function for each
 /// distinct values of `PARTITION BY` (each car type in our example)
 #[derive(Clone, Debug)]
-struct MyPartitionEvaluator {
-}
+struct MyPartitionEvaluator {}
 
 impl MyPartitionEvaluator {
-    fn new() -> Self
-    {
+    fn new() -> Self {
         Self {}
     }
 }
 
-
-
-
 /// These different evaluation methods are called depending on the various settings of WindowUDF
 impl PartitionEvaluator for MyPartitionEvaluator {
     fn get_range(&self, _idx: usize, _n_rows: usize) -> Result<std::ops::Range<usize>> {
@@ -142,14 +139,48 @@ impl PartitionEvaluator for MyPartitionEvaluator {
     }
 
     /// This function is given the values of each partition
-    fn evaluate(&self, values: &[arrow::array::ArrayRef], num_rows: usize) -> Result<arrow::array::ArrayRef> {
-        println!("processing num_rows={num_rows}, values:\n{values:#?}");
-        Err(DataFusionError::NotImplemented(
-            "evaluate is not implemented by default".into(),
-        ))
+    fn evaluate(
+        &self,
+        values: &[arrow::array::ArrayRef],
+        _num_rows: usize,
+    ) -> Result<arrow::array::ArrayRef> {
+        // datafusion has handled ensuring we get the correct input argument
+        assert_eq!(values.len(), 1);
+
+        // For this example, we convert convert the input argument to an
+        // array of floating point numbers to calculate a moving average
+        let arr: &Float64Array = values[0].as_ref().as_primitive::<Float64Type>();
+
+        // implement a simple moving average by averaging the current
+        // value with the previous value
+        //
+        // value | avg
+        // ------+------
+        //  10   | 10
+        //  20   | 15
+        //  30   | 25
+        //  30   | 30
+        //
+        let mut previous_value = None;
+        let new_values: Float64Array = arr
+            .values()
+            .iter()
+            .map(|&value| {
+                let new_value = previous_value
+                    .map(|previous_value| (value + previous_value) / 2.0)
+                    .unwrap_or(value);
+                previous_value = Some(value);
+                new_value
+            })
+            .collect();
+
+        Ok(Arc::new(new_values))
     }
 
-    fn evaluate_stateful(&mut self, _values: &[arrow::array::ArrayRef]) -> Result<datafusion_common::ScalarValue> {
+    fn evaluate_stateful(
+        &mut self,
+        _values: &[arrow::array::ArrayRef],
+    ) -> Result<datafusion_common::ScalarValue> {
         Err(DataFusionError::NotImplemented(
             "evaluate_stateful is not implemented by default".into(),
         ))
@@ -176,5 +207,4 @@ impl PartitionEvaluator for MyPartitionEvaluator {
     }
 }
 
-
 // TODO show how to use other evaluate methods

datafusion/core/src/execution/context.rs

@@ -815,7 +815,6 @@ impl SessionContext {
             .insert(f.name.clone(), Arc::new(f));
     }
 
-
     /// Creates a [`DataFrame`] for reading a data source.
     ///
     /// For more control such as reading multiple files, you can use
@@ -2042,7 +2041,6 @@ impl<'a> ContextProvider for SessionContextProvider<'a> {
     fn options(&self) -> &ConfigOptions {
         self.state.config_options()
     }
-
 }
 
 impl FunctionRegistry for SessionState {
@@ -2079,7 +2077,6 @@ impl FunctionRegistry for SessionState {
             ))
         })
     }
-
 }
 
 impl OptimizerConfig for SessionState {

datafusion/core/src/physical_plan/windows/mod.rs

@@ -245,6 +245,10 @@ impl BuiltInWindowFunctionExpr for WindowUDFExpr {
     fn create_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
         (self.fun.partition_evaluator)()
     }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
 }
 
 pub(crate) fn calc_requirements<

datafusion/expr/src/udwf.rs

@@ -17,7 +17,7 @@
 
 //! Support for user-defined window (UDWF) window functions
 
-use std::fmt::{self, Debug, Formatter};
+use std::fmt::{self, Debug, Display, Formatter};
 
 use crate::{function::PartitionEvaluatorFunctionFactory, ReturnTypeFunction, Signature};
 
@@ -41,6 +41,13 @@ impl Debug for WindowUDF {
     }
 }
 
+/// Defines how the WindowUDF is shown to users
+impl Display for WindowUDF {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(f, "{}", self.name)
+    }
+}
+
 impl PartialEq for WindowUDF {
     fn eq(&self, other: &Self) -> bool {
         self.name == other.name && self.signature == other.signature

datafusion/expr/src/window_function.rs

@@ -74,7 +74,7 @@ impl fmt::Display for WindowFunction {
             WindowFunction::AggregateFunction(fun) => fun.fmt(f),
             WindowFunction::BuiltInWindowFunction(fun) => fun.fmt(f),
             WindowFunction::AggregateUDF(fun) => std::fmt::Debug::fmt(fun, f),
-            WindowFunction::WindowUDF(fun) => std::fmt::Debug::fmt(fun, f),
+            WindowFunction::WindowUDF(fun) => fun.fmt(f),
         }
     }
 }

datafusion/physical-expr/src/window/built_in_window_function_expr.rs

@@ -46,9 +46,7 @@ pub trait BuiltInWindowFunctionExpr: Send + Sync + std::fmt::Debug {
 
     /// Human readable name such as `"MIN(c2)"` or `"RANK()"`. The default
     /// implementation returns placeholder text.
-    fn name(&self) -> &str {
-        "BuiltInWindowFunctionExpr: default name"
-    }
+    fn name(&self) -> &str;
 
     /// Evaluate window function's arguments against the input window
     /// batch and return an [`ArrayRef`].

datafusion/sql/src/expr/function.rs

@@ -189,7 +189,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                     .get_aggregate_meta(name)
                     .map(WindowFunction::AggregateUDF)
             })
-        // next check user defined window functions
+            // next check user defined window functions
             .or_else(|| {
                 self.schema_provider
                     .get_window_meta(name)