diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 7fe768cfbe329..241d91d2daeac 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/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>, } +#[inline] +fn can_evaluate_window_arg_unfiltered(expr: &dyn PhysicalExpr) -> bool { + expr.is::() || expr.is::() +} + +/// 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 { + 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 { - let values = self.evaluate_args(record_batch)?; - // Evaluate filter mask once per record batch if present let filter_mask_arr: Option = 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::>>()?, + None => self.evaluate_args(record_batch)?, + }; + if self.is_constant_in_partition() { if not_end { let field = self.field()?; diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 6374cbf4f4b80..a4b672472178d 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/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 (