diff --git a/datafusion/expr/src/partition_evaluator.rs b/datafusion/expr/src/partition_evaluator.rs index 5a4e20e5ac9ac..9ab7d1520203b 100644 --- a/datafusion/expr/src/partition_evaluator.rs +++ b/datafusion/expr/src/partition_evaluator.rs @@ -200,6 +200,33 @@ pub trait PartitionEvaluator: Debug + Send + std::any::Any { not_impl_err!("evaluate is not implemented by default") } + /// Evaluate window function on the window frames of consecutive rows + /// in an input partition, returning one result per range. + /// + /// This is called instead of [`Self::evaluate`] for window functions + /// that [use the window frame](Self::uses_window_frame), with the same + /// `values`. + /// + /// The default implementation calls [`Self::evaluate`] for each range. + /// Implementing this function is an optimization: a function such as + /// `FIRST_VALUE`, whose result is one of the input rows, can gather + /// all results with a single `take` instead of creating a + /// [`ScalarValue`] for every row. + fn evaluate_ranges( + &mut self, + values: &[ArrayRef], + ranges: &[Range], + ) -> Result { + if let [range] = ranges { + return self.evaluate(values, range)?.to_array(); + } + let mut results = Vec::with_capacity(ranges.len()); + for range in ranges { + results.push(self.evaluate(values, range)?); + } + ScalarValue::iter_to_array(results) + } + /// [`PartitionEvaluator::evaluate_all_with_rank`] is called for window /// functions that only need the rank of a row within its window /// frame. diff --git a/datafusion/functions-window/src/nth_value.rs b/datafusion/functions-window/src/nth_value.rs index 6724dbf9c6a60..7b98f77b43f77 100644 --- a/datafusion/functions-window/src/nth_value.rs +++ b/datafusion/functions-window/src/nth_value.rs @@ -19,7 +19,10 @@ use crate::utils::{get_scalar_value_from_args, get_signed_integer}; +use arrow::array::{Array, UInt64Array}; use arrow::buffer::BooleanBuffer; +use arrow::compute::kernels::zip::zip; +use arrow::compute::{is_not_null, take}; use arrow::datatypes::FieldRef; use arrow::util::bit_chunk_iterator::UnalignedBitChunk; use datafusion_common::arrow::array::ArrayRef; @@ -283,6 +286,8 @@ impl WindowUDFImpl for NthValue { state, ignore_nulls: partition_evaluator_args.ignore_nulls(), n: 0, + last_frame: None, + last_valid: None, })); } @@ -316,6 +321,8 @@ impl WindowUDFImpl for NthValue { state, ignore_nulls: partition_evaluator_args.ignore_nulls(), n, + last_frame: None, + last_valid: None, })) } @@ -378,6 +385,13 @@ pub(crate) struct NthValueEvaluator { state: NthValueState, ignore_nulls: bool, n: i64, + /// With `IGNORE NULLS`, whether FIRST_VALUE or NTH_VALUE selected a row in the + /// last frame passed to `evaluate_ranges`, and the first row of that frame that + /// later frames with the same start can still need. `memoize` consumes it. + last_frame: Option<(bool, usize)>, + /// With `IGNORE NULLS`, the last non-NULL result of LAST_VALUE that `memoize` + /// saw, which is also the result of later frames without valid rows + last_valid: Option, } impl PartitionEvaluator for NthValueEvaluator { @@ -390,31 +404,28 @@ impl PartitionEvaluator for NthValueEvaluator { let out = &state.out_col; let size = out.len(); if self.ignore_nulls { - match self.state.kind { - // Prune on first non-null output in case of FIRST_VALUE - NthValueKind::First => { - if let Some(nulls) = out.nulls() { - if self.state.finalized_result.is_none() { - if let Some(valid_index) = nulls.valid_indices().next() { - let result = - ScalarValue::try_from_array(out, valid_index)?; - self.state.finalized_result = Some(result); - } else { - // The output is empty or all nulls, ignore - } - } - if state.window_frame_range.start < state.window_frame_range.end { - state.window_frame_range.start = - state.window_frame_range.end - 1; - } - return Ok(()); - } else { - // Fall through to the main case because there are no nulls - } + if self.state.kind == NthValueKind::Last { + // LAST_VALUE returns its last valid value for a later frame without + // valid rows, so none of the frame is needed + if size > 0 && out.is_valid(size - 1) { + self.last_valid = Some(ScalarValue::try_from_array(out, size - 1)?); + } + state.window_frame_range.start = state.window_frame_range.end; + } else if let Some((selected, needed_from)) = self.last_frame.take() { + if selected && (self.state.kind == NthValueKind::First || self.n > 0) { + // FIRST_VALUE and NTH_VALUE(expr, n) select the same row in + // every later frame + let result = ScalarValue::try_from_array(out, size - 1)?; + self.state.finalized_result = Some(result); + } else { + state.window_frame_range.start = needed_from; } - // Do not memoize for other kinds when nulls are ignored - NthValueKind::Last | NthValueKind::Nth => return Ok(()), } + if self.state.finalized_result.is_some() { + state.window_frame_range.start = + state.window_frame_range.end.saturating_sub(1); + } + return Ok(()); } let mut buffer_size = 1; // Decide if we arrived at a final result yet: @@ -464,18 +475,55 @@ impl PartitionEvaluator for NthValueEvaluator { } else { // FIRST_VALUE, LAST_VALUE, NTH_VALUE window functions take a single column, values will have size 1. let arr = &values[0]; - let n_range = range.end - range.start; - if n_range == 0 { - // We produce None if the window is empty. - return ScalarValue::try_from(arr.data_type()); - } - match self.valid_index(arr, range) { - Some(index) => ScalarValue::try_from_array(arr, index), - None => ScalarValue::try_from(arr.data_type()), + match (self.valid_index(arr, range), &self.last_valid) { + (Some(index), _) => ScalarValue::try_from_array(arr, index), + (None, Some(last_valid)) => Ok(last_valid.clone()), + (None, None) => ScalarValue::try_from(arr.data_type()), } } } + fn evaluate_ranges( + &mut self, + values: &[ArrayRef], + ranges: &[Range], + ) -> Result { + if let Some(result) = &self.state.finalized_result { + return result.to_array_of_size(ranges.len()); + } + let arr = &values[0]; + if self.ignore_nulls + && self.state.kind != NthValueKind::Last + && let Some(range) = ranges.last() + { + // NTH_VALUE(expr, -n) selects the same row or a later one in later frames + // with the same start. Without a selected row, the frame has too few + // valid rows, all of which later frames can need. + let selected = self.valid_index(arr, range); + let needed_from = selected + .or_else(|| range.clone().find(|&idx| arr.is_valid(idx))) + .unwrap_or(range.end); + self.last_frame = Some((selected.is_some(), needed_from)); + } + if let [range] = ranges { + // A scalar is cheaper than a gather for a single row, as in one-row partitions + return self.evaluate(values, range)?.to_array(); + } + let indices = ranges + .iter() + .map(|range| self.valid_index(arr, range).map(|index| index as u64)) + .collect::(); + let result = take(arr, &indices, None)?; + match &self.last_valid { + Some(last_valid) if indices.null_count() > 0 => Ok(zip( + &is_not_null(&indices)?, + &result, + &last_valid.to_scalar()?, + )?), + _ => Ok(result), + } + } + fn supports_bounded_execution(&self) -> bool { true } @@ -488,6 +536,10 @@ impl PartitionEvaluator for NthValueEvaluator { impl NthValueEvaluator { fn valid_index(&self, array: &ArrayRef, range: &Range) -> Option { let n_range = range.end - range.start; + if n_range == 0 { + // We produce None if the window is empty. + return None; + } if self.ignore_nulls && let Some(nulls) = array.nulls() { @@ -732,6 +784,56 @@ mod tests { ); } + #[test] + fn memoize_ignore_nulls_moves_frame_start() -> Result<()> { + let values: ArrayRef = Arc::new(Int32Array::from(vec![ + None, + Some(1), + None, + Some(2), + None, + Some(3), + None, + ])); + let all_nulls: ArrayRef = Arc::new(Int32Array::new_null(7)); + let ranges = (1..=7).map(|end| 0..end).collect::>(); + // The frame start moves to the row NTH_VALUE(v, -n) selects, to the first valid + // row if there are fewer than `n`, to the last row once FIRST_VALUE or + // NTH_VALUE(v, n) is final, and past the frame without valid rows or for + // LAST_VALUE, which keeps its last valid value + let cases = [ + (NthValue::last(), None, &values, 7), + (NthValue::nth(), Some(-2), &values, 3), + (NthValue::nth(), Some(4), &values, 1), + (NthValue::nth(), Some(2), &values, 6), + (NthValue::first(), None, &values, 6), + (NthValue::nth(), Some(-2), &all_nulls, 7), + ]; + for (func, n, arr, expected_start) in cases { + let mut exprs = vec![Arc::new(Column::new("c", 0)) as Arc]; + if let Some(n) = n { + exprs.push(Arc::new(Literal::new(ScalarValue::Int64(Some(n))))); + } + let mut evaluator = func.partition_evaluator(PartitionEvaluatorArgs::new( + &exprs, + &[Field::new("f", DataType::Int32, true).into()], + false, + true, + ))?; + let mut state = WindowAggState::new(&DataType::Int32)?; + state.out_col = evaluator.evaluate_ranges(&[Arc::clone(arr)], &ranges)?; + state.window_frame_range = 0..7; + evaluator.memoize(&mut state)?; + assert_eq!( + state.window_frame_range.start, + expected_start, + "{} {n:?}", + func.name() + ); + } + Ok(()) + } + #[test] fn nth_set_bit_from_end_matches_a_scan_from_the_end() { // The slices cover every bit offset within a word, byte-aligned and unaligned starts, diff --git a/datafusion/physical-expr/src/window/standard.rs b/datafusion/physical-expr/src/window/standard.rs index 97e22f223ceec..92b72ac7d0b97 100644 --- a/datafusion/physical-expr/src/window/standard.rs +++ b/datafusion/physical-expr/src/window/standard.rs @@ -18,7 +18,6 @@ //! Physical exec for standard window function expressions. use std::any::Any; -use std::ops::Range; use std::sync::Arc; use super::{StandardWindowFunctionExpr, WindowExpr}; @@ -117,7 +116,7 @@ impl WindowExpr for StandardWindowExpr { let num_rows = batch.num_rows(); if evaluator.uses_window_frame() { let sort_options = self.order_by.iter().map(|o| o.options).collect(); - let mut row_wise_results = Vec::with_capacity(num_rows); + let mut ranges = Vec::with_capacity(num_rows); let mut values = self.evaluate_args(batch)?; let order_bys = get_orderby_values(self.order_by_columns(batch)?); @@ -127,20 +126,17 @@ impl WindowExpr for StandardWindowExpr { let mut window_frame_ctx = WindowFrameContext::new(Arc::clone(&self.window_frame), sort_options); - let mut last_range = Range { start: 0, end: 0 }; - // We iterate on each row to calculate window frame range and window function result + // We iterate on each row to calculate its window frame range for idx in 0..num_rows { let range = window_frame_ctx.calculate_range( order_bys_ref, - &last_range, + ranges.last().unwrap_or(&(0..0)), num_rows, idx, )?; - let value = evaluator.evaluate(&values, &range)?; - row_wise_results.push(value); - last_range = range; + ranges.push(range); } - ScalarValue::iter_to_array(row_wise_results) + evaluator.evaluate_ranges(&values, &ranges) } else if evaluator.include_rank() { let columns = self.order_by_columns(batch)?; let sort_partition_points = evaluate_partition_ranges(num_rows, &columns)?; @@ -191,7 +187,8 @@ impl WindowExpr for StandardWindowExpr { let batch_ref = &partition_batch_state.record_batch; let mut values = self.evaluate_args(batch_ref)?; - let order_bys = if evaluator.uses_window_frame() || evaluator.include_rank() { + let uses_window_frame = evaluator.uses_window_frame(); + let order_bys = if uses_window_frame || evaluator.include_rank() { get_orderby_values(self.order_by_columns(batch_ref)?) } else { vec![] @@ -203,15 +200,22 @@ impl WindowExpr for StandardWindowExpr { // We iterate on each row to perform a running calculation. let record_batch = &partition_batch_state.record_batch; let num_rows = record_batch.num_rows(); - let mut row_wise_results: Vec = - Vec::with_capacity(num_rows.saturating_sub(state.last_calculated_index)); - let is_causal = if evaluator.uses_window_frame() { + let n_new_rows = num_rows.saturating_sub(state.last_calculated_index); + // Functions that use the window frame are evaluated on all frames after + // the loop. Other functions are evaluated row by row, because their + // `get_range` can depend on state that `evaluate` updates. + let (mut frame_ranges, mut row_wise_results) = if uses_window_frame { + (Vec::with_capacity(n_new_rows), Vec::new()) + } else { + (Vec::new(), Vec::with_capacity(n_new_rows)) + }; + let is_causal = if uses_window_frame { self.window_frame.is_causal() } else { evaluator.is_causal() }; for idx in state.last_calculated_index..num_rows { - let frame_range = if evaluator.uses_window_frame() { + let frame_range = if uses_window_frame { state .window_frame_ctx .get_or_insert_with(|| { @@ -240,10 +244,16 @@ impl WindowExpr for StandardWindowExpr { } // Update last range state.window_frame_range = frame_range; - row_wise_results - .push(evaluator.evaluate(&values, &state.window_frame_range)?); + if uses_window_frame { + frame_ranges.push(state.window_frame_range.clone()); + } else { + row_wise_results + .push(evaluator.evaluate(&values, &state.window_frame_range)?); + } } - let out_col = if row_wise_results.is_empty() { + let out_col = if !frame_ranges.is_empty() { + evaluator.evaluate_ranges(&values, &frame_ranges)? + } else if row_wise_results.is_empty() { new_empty_array(out_type) } else if row_wise_results.len() == 1 { // fast path when the result only has a single row diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 864a6e56bf574..7d050ec4d3bb3 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -4724,6 +4724,17 @@ NULL NULL NULL +# FIRST_VALUE and LAST_VALUE are NULL over an empty frame +query III +SELECT i, + FIRST_VALUE(i) OVER (ORDER BY i ROWS BETWEEN 2 PRECEDING AND 1 PRECEDING), + LAST_VALUE(i) OVER (ORDER BY i ROWS BETWEEN 1 FOLLOWING AND 2 FOLLOWING) +FROM (VALUES (1), (2), (3)) t(i); +---- +1 NULL 3 +2 1 3 +3 1 NULL + # Test for ignore nulls in FIRST_VALUE statement ok CREATE TABLE t AS VALUES (null::bigint), (3), (4); @@ -5014,17 +5025,23 @@ SELECT i % 4 AS p, i, END AS v FROM (SELECT value AS i FROM generate_series(1, 30000)); -query IIII +# `s` wraps `v` in a struct, so its results must come from the same rows +query IIIIII SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END), - count(fv), sum(CASE WHEN fv IS DISTINCT FROM mn THEN 1 ELSE 0 END) + count(fv), sum(CASE WHEN fv IS DISTINCT FROM mn THEN 1 ELSE 0 END), + sum(CASE WHEN ls['v'] IS DISTINCT FROM mx THEN 1 ELSE 0 END), + sum(CASE WHEN fs['v'] IS DISTINCT FROM mn THEN 1 ELSE 0 END) FROM (SELECT LAST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS lv, max(v) OVER (PARTITION BY p ORDER BY i) AS mx, FIRST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS fv, - min(v) OVER (PARTITION BY p ORDER BY i) AS mn -FROM ignore_nulls_t); + min(v) OVER (PARTITION BY p ORDER BY i) AS mn, + LAST_VALUE(s) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS ls, + FIRST_VALUE(s) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS fs +FROM (SELECT *, CASE WHEN v IS NOT NULL THEN named_struct('v', v) END AS s + FROM ignore_nulls_t)); ---- -22499 0 22499 0 +22499 0 22499 0 0 0 query IIII SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END), @@ -5071,6 +5088,24 @@ FROM r LEFT JOIN r e ON e.p = r.p AND e.v IS NOT NULL AND e.c = r.c - 64; ---- 14807 7372 0 +# NTH_VALUE(v, 4) and NTH_VALUE(v, -4) over growing frames, whose start `memoize` moves forward +# between batches. In partition 1, batches end with fewer than 4 valid rows in the frame. The +# expected values are the valid rows whose running count is 4, and 3 less than the current row's. +query IIII +WITH r AS (SELECT p, v, + count(v) OVER (PARTITION BY p ORDER BY i) AS c, + NTH_VALUE(v, 4) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS n4, + NTH_VALUE(v, -4) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS nm4 +FROM ignore_nulls_t) +SELECT count(r.n4), count(r.nm4), + sum(CASE WHEN r.n4 IS DISTINCT FROM f.v THEN 1 ELSE 0 END), + sum(CASE WHEN r.nm4 IS DISTINCT FROM e.v THEN 1 ELSE 0 END) +FROM r +LEFT JOIN r f ON f.p = r.p AND f.v IS NOT NULL AND f.c = 4 AND r.c >= 4 +LEFT JOIN r e ON e.p = r.p AND e.v IS NOT NULL AND e.c = r.c - 3; +---- +21990 21990 0 0 + query II SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END) FROM (SELECT diff --git a/datafusion/sqllogictest/test_files/window_ignore_nulls_matrix.slt b/datafusion/sqllogictest/test_files/window_ignore_nulls_matrix.slt new file mode 100644 index 0000000000000..eedf0619bcd2b --- /dev/null +++ b/datafusion/sqllogictest/test_files/window_ignore_nulls_matrix.slt @@ -0,0 +1,35 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# FIRST_VALUE, LAST_VALUE and NTH_VALUE with IGNORE NULLS, whose `memoize` drops rows between +# batches. One row per batch evaluates a single frame, more rows gather results for several. +# Matrix rules: no EXPLAIN, no in-file SET of a swept knob, rowsort every multi-row query. + +# configMatrix: datafusion.execution.batch_size=1,2,100,8192 + +# The RANGE frame end stalls after each jump of `k`, so a frame can be empty once `memoize` has +# dropped the rows before it. `mx` comes from a separate operator, which does not keep them. +query II +SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END) +FROM (SELECT mx, LAST_VALUE(v) IGNORE NULLS OVER w AS lv + FROM (SELECT k, v, max(v) OVER w AS mx + FROM (SELECT i + 100 * (i / 10) AS k, CASE WHEN i % 7 < 3 THEN i END AS v + FROM (SELECT value AS i FROM generate_series(1, 200))) + WINDOW w AS (ORDER BY k RANGE BETWEEN UNBOUNDED PRECEDING AND 5 PRECEDING)) + WINDOW w AS (ORDER BY k RANGE BETWEEN UNBOUNDED PRECEDING AND 5 PRECEDING)); +---- +195 0