sketch out implementation

Author: alamb

datafusion-examples/examples/simple_udwf.rs

@@ -70,14 +70,14 @@ async fn main() -> Result<()> {
     // print the results
     df.show().await?;
 
-    // ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING: Run the window functon so that each invocation only sees 5 rows: the 2 before and 2 after) using
-    let df = ctx.sql("SELECT car, \
-                      speed, \
-                      lag(speed, 1) OVER (PARTITION BY car ORDER BY time ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING),\
-                      time \
-                      from cars").await?;
-    // print the results
-    df.show().await?;
+    // // ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING: Run the window functon so that each invocation only sees 5 rows: the 2 before and 2 after) using
+    // let df = ctx.sql("SELECT car, \
+    //                   speed, \
+    //                   lag(speed, 1) OVER (PARTITION BY car ORDER BY time ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING),\
+    //                   time \
+    //                   from cars").await?;
+    // // print the results
+    // df.show().await?;
 
     // todo show how to run dataframe API as well
 
@@ -108,7 +108,73 @@ fn return_type(arg_types: &[DataType]) -> Result<Arc<DataType>> {
     Ok(Arc::new(arg_types[0].clone()))
 }
 
-/// create a partition evaluator for this argument
+/// Create a partition evaluator for this argument
 fn make_partition_evaluator() -> Result<Box<dyn PartitionEvaluator>> {
-    todo!();
+    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 {
+}
+
+impl MyPartitionEvaluator {
+    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>> {
+        Err(DataFusionError::NotImplemented(
+            "get_range is not implemented for this window function".to_string(),
+        ))
+    }
+
+    /// 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_stateful(&mut self, _values: &[arrow::array::ArrayRef]) -> Result<datafusion_common::ScalarValue> {
+        Err(DataFusionError::NotImplemented(
+            "evaluate_stateful is not implemented by default".into(),
+        ))
+    }
+
+    fn evaluate_with_rank(
+        &self,
+        _num_rows: usize,
+        _ranks_in_partition: &[std::ops::Range<usize>],
+    ) -> Result<arrow::array::ArrayRef> {
+        Err(DataFusionError::NotImplemented(
+            "evaluate_partition_with_rank is not implemented by default".into(),
+        ))
+    }
+
+    fn evaluate_inside_range(
+        &self,
+        _values: &[arrow::array::ArrayRef],
+        _range: &std::ops::Range<usize>,
+    ) -> Result<datafusion_common::ScalarValue> {
+        Err(DataFusionError::NotImplemented(
+            "evaluate_inside_range is not implemented by default".into(),
+        ))
+    }
+}
+
+
+// TODO show how to use other evaluate methods