Skip to content

perf: memoize LAST_VALUE and NTH_VALUE with IGNORE NULLS #26664

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
comphead merged 2 commits into from
Oct 24, 2026
+262 −53
Merged

perf: memoize LAST_VALUE and NTH_VALUE with IGNORE NULLS #26664

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
40c848f
perf: memoize `LAST_VALUE` and `NTH_VALUE` with `IGNORE NULLS`, gathe…
comphead Oct 23, 2026
b8518a5
test: cover empty frames after `LAST_VALUE IGNORE NULLS` memoize acro…
comphead Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions datafusion/expr/src/partition_evaluator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<usize>],
) -> Result<ArrayRef> {
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.
Expand Down
164 changes: 133 additions & 31 deletions datafusion/functions-window/src/nth_value.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -283,6 +286,8 @@ impl WindowUDFImpl for NthValue {
state,
ignore_nulls: partition_evaluator_args.ignore_nulls(),
n: 0,
last_frame: None,
last_valid: None,
}));
}

Expand Down Expand Up @@ -316,6 +321,8 @@ impl WindowUDFImpl for NthValue {
state,
ignore_nulls: partition_evaluator_args.ignore_nulls(),
n,
last_frame: None,
last_valid: None,
}))
}

Expand Down Expand Up @@ -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<ScalarValue>,
}

impl PartitionEvaluator for NthValueEvaluator {
Expand All @@ -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:
Expand Down Expand Up @@ -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()),

@jayzhan211 jayzhan211 Oct 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Returning last_valid for an empty frame looks right, since memoize only runs with an unbounded start. But no test reaches it: if this arm returns NULL for an empty range, window.slt still passes. The existing tests compute max(v) in the same operator, which keeps the early rows buffered, so the frame never becomes empty. Could you add a test like this one? It puts the reference in a separate operator and fails if empty frames return NULL.

statement ok
set datafusion.execution.batch_size = 2;

# The RANGE frame end stalls after each jump of `k`, so a frame can be empty
# once `memoize` has dropped the rows before it
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

statement ok
reset datafusion.execution.batch_size;

@comphead comphead Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it would make sense to include matrix in this file

@comphead comphead Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually it is good to have batchsize matrices everywhere in slt IMO.... 🤔

@comphead comphead Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, added as window_ignore_nulls_matrix.slt, which sweeps batch_size over 1, 2, 100 and 8192. At batch_size = 2 each call evaluates two frames, so the query reaches the take/zip path in evaluate_ranges but not this scalar arm. batch_size = 1 covers the arm. Each path fails its run if empty frames return NULL.

@comphead comphead Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

running slt for different batch sizes, it would be part of bigger thing

@comphead comphead Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#26700

(None, None) => ScalarValue::try_from(arr.data_type()),
}
}
}

fn evaluate_ranges(
&mut self,
values: &[ArrayRef],
ranges: &[Range<usize>],
) -> Result<ArrayRef> {
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::<UInt64Array>();
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
}
Expand All @@ -488,6 +536,10 @@ impl PartitionEvaluator for NthValueEvaluator {
impl NthValueEvaluator {
fn valid_index(&self, array: &ArrayRef, range: &Range<usize>) -> Option<usize> {
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()
{
Expand Down Expand Up @@ -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::<Vec<_>>();
// 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<dyn PhysicalExpr>];
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,
Expand Down
44 changes: 27 additions & 17 deletions datafusion/physical-expr/src/window/standard.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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)?);
Expand All @@ -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)?;
Expand Down Expand Up @@ -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![]
Expand All @@ -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<ScalarValue> =
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(|| {
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading