Return error rather than wrong results when aggregate without retract_batch is used as a sliding accumulator

alamb · alamb · commit 999866e281e4 · 2023-06-15T13:00:10.000-04:00
diff --git a/datafusion/core/src/physical_plan/udaf.rs b/datafusion/core/src/physical_plan/udaf.rs
@@ -41,7 +41,7 @@ pub fn create_aggregate_expr(
     input_phy_exprs: &[Arc<dyn PhysicalExpr>],
     input_schema: &Schema,
     name: impl Into<String>,
-) -> Result<Arc<AggregateFunctionExpr>> {
+) -> Result<Arc<dyn AggregateExpr>> {
     let input_exprs_types = input_phy_exprs
         .iter()
         .map(|arg| arg.data_type(input_schema))
@@ -70,11 +70,6 @@ impl AggregateFunctionExpr {
     pub fn fun(&self) -> &AggregateUDF {
         &self.fun
     }
-
-    /// Returns true if this can support sliding accumulators
-    pub fn retractable(&self) -> Result<bool> {
-        Ok((self.fun.accumulator)(&self.data_type)?.supports_retract_batch())
-    }
 }
 
 impl AggregateExpr for AggregateFunctionExpr {
@@ -114,12 +109,54 @@ impl AggregateExpr for AggregateFunctionExpr {
     fn create_sliding_accumulator(&self) -> Result<Box<dyn Accumulator>> {
         let accumulator = (self.fun.accumulator)(&self.data_type)?;
 
+        // Accumulators that have window frame startings different
+        // than `UNBOUNDED PRECEDING`, such as `1 PRECEEDING`, need to
+        // implement retract_batch method in order to run correctly
+        // currently in DataFusion.
+        //
+        // If this `retract_batches` is not present, there is no way
+        // to calculate result correctly. For example, the query
+        //
+        // ```sql
+        // SELECT
+        //  SUM(a) OVER(ORDER BY a ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS sum_a
+        // FROM
+        //  t
+        // ```
+        //
+        // 1. First sum value will be the sum of rows between `[0, 1)`,
+        //
+        // 2. Second sum value will be the sum of rows between `[0, 2)`
+        //
+        // 3. Third sum value will be the sum of rows between `[1, 3)`, etc.
+        //
+        // Since the accumulator keeps the running sum:
+        //
+        // 1. First sum we add to the state sum value between `[0, 1)`
+        //
+        // 2. Second sum we add to the state sum value between `[1, 2)`
+        // (`[0, 1)` is already in the state sum, hence running sum will
+        // cover `[0, 2)` range)
+        //
+        // 3. Third sum we add to the state sum value between `[2, 3)`
+        // (`[0, 2)` is already in the state sum).  Also we need to
+        // retract values between `[0, 1)` by this way we can obtain sum
+        // between [1, 3) which is indeed the apropriate range.
+        //
+        // When we use `UNBOUNDED PRECEDING` in the query starting
+        // index will always be 0 for the desired range, and hence the
+        // `retract_batch` method will not be called. In this case
+        // having retract_batch is not a requirement.
+        //
+        // This approach is a a bit different than window function
+        // approach. In window function (when they use a window frame)
+        // they get all the desired range during evaluation.
         if !accumulator.supports_retract_batch() {
-            return Err(DataFusionError::Internal(
-                format!(
-                    "Can't make sliding accumulator because retractable_accumulator not available for {}",
-                    self.name)
-            ));
+            return Err(DataFusionError::NotImplemented(format!(
+                "Aggregate can not be used as a sliding accumulator because \
+                     `retract_batch` is not implemented: {}",
+                self.name
+            )));
         }
         Ok(accumulator)
     }
diff --git a/datafusion/core/src/physical_plan/windows/mod.rs b/datafusion/core/src/physical_plan/windows/mod.rs
@@ -33,8 +33,9 @@ use datafusion_expr::{
     window_function::{BuiltInWindowFunction, WindowFunction},
     WindowFrame,
 };
-use datafusion_physical_expr::window::{
-    BuiltInWindowFunctionExpr, SlidingAggregateWindowExpr,
+use datafusion_physical_expr::{
+    window::{BuiltInWindowFunctionExpr, SlidingAggregateWindowExpr},
+    AggregateExpr,
 };
 use std::borrow::Borrow;
 use std::convert::TryInto;
@@ -64,28 +65,16 @@ pub fn create_window_expr(
     window_frame: Arc<WindowFrame>,
     input_schema: &Schema,
 ) -> Result<Arc<dyn WindowExpr>> {
-    // Is there a potentially unlimited sized window frame?
-    let unbounded_window = window_frame.start_bound.is_unbounded();
-
     Ok(match fun {
         WindowFunction::AggregateFunction(fun) => {
             let aggregate =
                 aggregates::create_aggregate_expr(fun, false, args, input_schema, name)?;
-            if !unbounded_window {
-                Arc::new(SlidingAggregateWindowExpr::new(
-                    aggregate,
-                    partition_by,
-                    order_by,
-                    window_frame,
-                ))
-            } else {
-                Arc::new(PlainAggregateWindowExpr::new(
-                    aggregate,
-                    partition_by,
-                    order_by,
-                    window_frame,
-                ))
-            }
+            window_expr_from_aggregate_expr(
+                partition_by,
+                order_by,
+                window_frame,
+                aggregate,
+            )
         }
         WindowFunction::BuiltInWindowFunction(fun) => Arc::new(BuiltInWindowExpr::new(
             create_built_in_window_expr(fun, args, input_schema, name)?,
@@ -96,26 +85,43 @@ pub fn create_window_expr(
         WindowFunction::AggregateUDF(fun) => {
             let aggregate =
                 udaf::create_aggregate_expr(fun.as_ref(), args, input_schema, name)?;
-
-            if !unbounded_window && aggregate.retractable()? {
-                Arc::new(SlidingAggregateWindowExpr::new(
-                    aggregate,
-                    partition_by,
-                    order_by,
-                    window_frame,
-                ))
-            } else {
-                Arc::new(PlainAggregateWindowExpr::new(
-                    aggregate,
-                    partition_by,
-                    order_by,
-                    window_frame,
-                ))
-            }
+            window_expr_from_aggregate_expr(
+                partition_by,
+                order_by,
+                window_frame,
+                aggregate,
+            )
         }
     })
 }
 
+/// Creates an appropriate [`WindowExpr`] based on the window frame and
+fn window_expr_from_aggregate_expr(
+    partition_by: &[Arc<dyn PhysicalExpr>],
+    order_by: &[PhysicalSortExpr],
+    window_frame: Arc<WindowFrame>,
+    aggregate: Arc<dyn AggregateExpr>,
+) -> Arc<dyn WindowExpr> {
+    // Is there a potentially unlimited sized window frame?
+    let unbounded_window = window_frame.start_bound.is_unbounded();
+
+    if !unbounded_window {
+        Arc::new(SlidingAggregateWindowExpr::new(
+            aggregate,
+            partition_by,
+            order_by,
+            window_frame,
+        ))
+    } else {
+        Arc::new(PlainAggregateWindowExpr::new(
+            aggregate,
+            partition_by,
+            order_by,
+            window_frame,
+        ))
+    }
+}
+
 fn get_scalar_value_from_args(
     args: &[Arc<dyn PhysicalExpr>],
     index: usize,
diff --git a/datafusion/core/tests/user_defined_aggregates.rs b/datafusion/core/tests/user_defined_aggregates.rs
@@ -40,7 +40,7 @@ use datafusion::{
     prelude::SessionContext,
     scalar::ScalarValue,
 };
-use datafusion_common::{cast::as_primitive_array, DataFusionError};
+use datafusion_common::{assert_contains, cast::as_primitive_array, DataFusionError};
 
 /// Test to show the contents of the setup
 #[tokio::test]
@@ -58,7 +58,7 @@ async fn test_setup() {
         "| 5.0   | 1970-01-01T00:00:00.000005 |",
         "+-------+----------------------------+",
     ];
-    assert_batches_eq!(expected, &execute(&ctx, sql).await);
+    assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
 }
 
 /// Basic user defined aggregate
@@ -74,7 +74,7 @@ async fn test_udaf() {
         "| 1970-01-01T00:00:00.000019 |",
         "+----------------------------+",
     ];
-    assert_batches_eq!(expected, &execute(&ctx, sql).await);
+    assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
     // normal aggregates call update_batch
     assert!(test_state.update_batch());
     assert!(!test_state.retract_batch());
@@ -96,7 +96,7 @@ async fn test_udaf_as_window() {
         "| 1970-01-01T00:00:00.000019 |",
         "+----------------------------+",
     ];
-    assert_batches_eq!(expected, &execute(&ctx, sql).await);
+    assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
     // aggregate over the entire window function call update_batch
     assert!(test_state.update_batch());
     assert!(!test_state.retract_batch());
@@ -118,35 +118,23 @@ async fn test_udaf_as_window_with_frame() {
         "| 1970-01-01T00:00:00.000010 |",
         "+----------------------------+",
     ];
-    assert_batches_eq!(expected, &execute(&ctx, sql).await);
+    assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
     // user defined aggregates with window frame should be calling retract batch
     assert!(test_state.update_batch());
     assert!(test_state.retract_batch());
 }
 
 /// Ensure that User defined aggregate used as a window function with a window
-/// frame, but that does not implement retract_batch, does not error
+/// frame, but that does not implement retract_batch, returns an error
 #[tokio::test]
 async fn test_udaf_as_window_with_frame_without_retract_batch() {
     let test_state = Arc::new(TestState::new().with_error_on_retract_batch());
 
-    let TestContext { ctx, test_state } = TestContext::new_with_test_state(test_state);
+    let TestContext { ctx, test_state: _ } = TestContext::new_with_test_state(test_state);
     let sql = "SELECT time_sum(time) OVER(ORDER BY time ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) as time_sum from t";
-    // TODO: It is not clear why this is a different value than when retract batch is used
-    let expected = vec![
-        "+----------------------------+",
-        "| time_sum                   |",
-        "+----------------------------+",
-        "| 1970-01-01T00:00:00.000005 |",
-        "| 1970-01-01T00:00:00.000009 |",
-        "| 1970-01-01T00:00:00.000014 |",
-        "| 1970-01-01T00:00:00.000019 |",
-        "| 1970-01-01T00:00:00.000019 |",
-        "+----------------------------+",
-    ];
-    assert_batches_eq!(expected, &execute(&ctx, sql).await);
-    assert!(test_state.update_batch());
-    assert!(!test_state.retract_batch());
+    // Note if this query ever does start working
+    let err = execute(&ctx, sql).await.unwrap_err();
+    assert_contains!(err.to_string(), "This feature is not implemented: Aggregate can not be used as a sliding accumulator because `retract_batch` is not implemented: AggregateUDF { name: \"time_sum\"");
 }
 
 /// Basic query for with a udaf returning a structure
@@ -161,7 +149,7 @@ async fn test_udaf_returning_struct() {
         "| {value: 2.0, time: 1970-01-01T00:00:00.000002} |",
         "+------------------------------------------------+",
     ];
-    assert_batches_eq!(expected, &execute(&ctx, sql).await);
+    assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
 }
 
 /// Demonstrate extracting the fields from a structure using a subquery
@@ -176,11 +164,11 @@ async fn test_udaf_returning_struct_subquery() {
         "| 2.0             | 1970-01-01T00:00:00.000002 |",
         "+-----------------+----------------------------+",
     ];
-    assert_batches_eq!(expected, &execute(&ctx, sql).await);
+    assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
 }
 
-async fn execute(ctx: &SessionContext, sql: &str) -> Vec<RecordBatch> {
-    ctx.sql(sql).await.unwrap().collect().await.unwrap()
+async fn execute(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
+    ctx.sql(sql).await?.collect().await
 }
 
 /// Returns an context with a table "t" and the "first" and "time_sum"
diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
@@ -465,7 +465,6 @@ impl AsExecutionPlan for PhysicalPlanNode {
                                         AggregateFunction::UserDefinedAggrFunction(udaf_name) => {
                                             let agg_udf = registry.udaf(udaf_name)?;
                                             udaf::create_aggregate_expr(agg_udf.as_ref(), &input_phy_expr, &physical_schema, name)
-                                                .map(|func| func as Arc<dyn AggregateExpr>)
                                         }
                                     }
                                 }).transpose()?.ok_or_else(|| {

Original file line number	Diff line number	Diff line change
`@@ -465,7 +465,6 @@ impl AsExecutionPlan for PhysicalPlanNode {`
`465`	`465`	`AggregateFunction::UserDefinedAggrFunction(udaf_name) => {`
`466`	`466`	`let agg_udf = registry.udaf(udaf_name)?;`
`467`	`467`	`udaf::create_aggregate_expr(agg_udf.as_ref(), &input_phy_expr, &physical_schema, name)`
`468`		`- .map(\|func\| func as Arc<dyn AggregateExpr>)`
`469`	`468`	`}`
`470`	`469`	`}`
`471`	`470`	`}).transpose()?.ok_or_else(\|\| {`