fix: evaluate window aggregate arguments after FILTER

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

@@ -21,6 +21,7 @@ use std::ops::Range;
 use std::sync::Arc;
 
 use crate::PhysicalExpr;
+use crate::expressions::{Column, Literal};
 
 use arrow::array::BooleanArray;
 use arrow::array::{Array, ArrayRef, new_empty_array};
@@ -182,6 +183,26 @@ pub struct WindowPhysicalExpressions {
     pub order_by_exprs: Vec<Arc<dyn PhysicalExpr>>,
 }
 
+#[inline]
+fn can_evaluate_window_arg_unfiltered(expr: &dyn PhysicalExpr) -> bool {
+    expr.is::<Column>() || expr.is::<Literal>()
+}
+
+/// Evaluates safe arguments on the original batch and all other arguments
+/// using `selection`.
+fn evaluate_window_arg(
+    expr: &dyn PhysicalExpr,
+    record_batch: &RecordBatch,
+    selection: &BooleanArray,
+) -> Result<ArrayRef> {
+    let value = if can_evaluate_window_arg_unfiltered(expr) {
+        expr.evaluate(record_batch)?
+    } else {
+        expr.evaluate_selection(record_batch, selection)?
+    };
+    value.into_array_of_size(record_batch.num_rows())
+}
+
 /// Extension trait that adds common functionality to [`AggregateWindowExpr`]s
 pub trait AggregateWindowExpr: WindowExpr {
     /// Get the accumulator for the window expression. Note that distinct
@@ -314,8 +335,6 @@ pub trait AggregateWindowExpr: WindowExpr {
         mut idx: usize,
         not_end: bool,
     ) -> Result<ArrayRef> {
-        let values = self.evaluate_args(record_batch)?;
-
         // Evaluate filter mask once per record batch if present
         let filter_mask_arr: Option<ArrayRef> = match self.filter_expr() {
             Some(expr) => {
@@ -331,6 +350,18 @@ pub trait AggregateWindowExpr: WindowExpr {
             None => None,
         };
 
+        // Columns and literals keep their original alignment. Other arguments are
+        // evaluated only on selected rows and scattered back to the row indices used
+        // by window frames.
+        let values = match filter_mask {
+            Some(mask) => self
+                .expressions()
+                .iter()
+                .map(|expr| evaluate_window_arg(expr.as_ref(), record_batch, mask))
+                .collect::<Result<Vec<_>>>()?,
+            None => self.evaluate_args(record_batch)?,
+        };
+
         if self.is_constant_in_partition() {
             if not_end {
                 let field = self.field()?;

datafusion/sqllogictest/test_files/window.slt

@@ -6150,6 +6150,30 @@ LIMIT 5
 0 3 NULL NULL 0 NULL NULL
 0 4 NULL NULL 0 NULL NULL
 
+# FILTER excludes rows before evaluating fallible window aggregate arguments
+query II
+SELECT id,
+       SUM(10 / x) FILTER (WHERE x <> 0) OVER (ORDER BY id) AS running_sum
+FROM (VALUES (1, 2), (2, 0), (3, 5)) AS t(id, x)
+ORDER BY id
+----
+1 5
+2 5
+3 7
+
+# FILTER preserves row alignment while evaluating arguments for a sliding frame
+query III
+SELECT id, x,
+       SUM(10 / x) FILTER (WHERE x <> 0) OVER (
+           ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW
+       ) AS s
+FROM (VALUES (1, 2), (2, 0), (3, 5)) AS t(id, x)
+ORDER BY id
+----
+1 2 5
+2 0 5
+3 5 2
+
 # regression test for https://github.com/apache/datafusion/issues/17401
 query I
 WITH source AS (