This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23711-abc226e6c637cdb00012e8e6cd5884a60ae23f3e in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7d42518fb8e8a1bdb1f7b542cbedc55b9cbff73a Author: xudong.w <[email protected]> AuthorDate: Fri Jul 24 10:01:36 2026 +0800 perf: optimize LEAD/LAG IGNORE NULLS evaluation (#23711) ## Which issue does this PR close? - N/A ## Rationale for this change The whole-partition evaluation path for `LEAD` and `LAG` with `IGNORE NULLS` currently collects every valid row index, performs a binary search for every output row, converts every selected value to a `ScalarValue`, and finally rebuilds an Arrow array from those scalars. This makes index selection `O(n log m)` for `n` rows and `m` non-null rows, and the per-row scalar materialization is particularly expensive for strings and nested values. A local Criterion microbenchmark with 100,000 rows and 50% nulls measured the following speedups across the tested offsets: | Data type | Speedup | | --- | ---: | | `Int64` | 6.7–6.9x | | `Utf8View` | 10.5–11.6x | | `List` | 21.7–21.8x | ## What changes are included in this PR? - Generate nullable Arrow gather indices with a single linear scan and bounded `VecDeque` state instead of collecting all valid indices and binary-searching for every row. - Materialize the result with Arrow's `take` kernel. - Use Arrow's `zip` kernel to fill non-null default values without constructing one `ScalarValue` per row. - Keep dictionary arrays with non-null defaults on scalar materialization because `zip` concatenates dictionaries and can overflow bounded dictionary key types. The index generation and output materialization are now `O(n)`. There are no public API changes. ## Are these changes tested? Yes. ## Are there any user-facing changes? No. This changes the implementation of whole-partition `LEAD`/`LAG ... IGNORE NULLS` evaluation without changing the SQL behavior or public API. --- datafusion/functions-window/src/lead_lag.rs | 242 +++++++++++++++++++++++----- 1 file changed, 203 insertions(+), 39 deletions(-) diff --git a/datafusion/functions-window/src/lead_lag.rs b/datafusion/functions-window/src/lead_lag.rs index 2363a6beed..fea4a1a4aa 100644 --- a/datafusion/functions-window/src/lead_lag.rs +++ b/datafusion/functions-window/src/lead_lag.rs @@ -18,6 +18,8 @@ //! `lead` and `lag` window function implementations use crate::utils::{get_scalar_value_from_args, get_signed_integer}; +use arrow::array::UInt64Builder; +use arrow::compute::{interleave, take}; use arrow::datatypes::FieldRef; use datafusion_common::arrow::array::ArrayRef; use datafusion_common::arrow::datatypes::DataType; @@ -419,6 +421,52 @@ fn offset_magnitude(offset: i64) -> usize { } } +enum ShiftIndexBuilder { + Take(UInt64Builder), + Interleave(Vec<(usize, usize)>), +} + +impl ShiftIndexBuilder { + fn new(capacity: usize, default_is_null: bool) -> Self { + if default_is_null { + Self::Take(UInt64Builder::with_capacity(capacity)) + } else { + Self::Interleave(Vec::with_capacity(capacity)) + } + } + + fn append_option(&mut self, index: Option<usize>) { + match self { + Self::Take(indices) => { + indices.append_option(index.map(|index| index as u64)); + } + Self::Interleave(indices) => { + // `interleave` receives `[array, default]`. + indices.push(index.map_or((1, 0), |index| (0, index))); + } + } + } + + fn finish( + self, + array: &ArrayRef, + default_value: &ScalarValue, + ) -> Result<ArrayRef, DataFusionError> { + match self { + Self::Take(mut indices) => { + let indices = indices.finish(); + take(array.as_ref(), &indices, None) + .map_err(|error| arrow_datafusion_err!(error)) + } + Self::Interleave(indices) => { + let default = default_value.to_array_of_size(1)?; + interleave(&[array.as_ref(), default.as_ref()], &indices) + .map_err(|error| arrow_datafusion_err!(error)) + } + } + } +} + impl WindowShiftEvaluator { fn is_lag(&self) -> bool { // Mode is LAG, when shift_offset is positive @@ -433,53 +481,57 @@ fn evaluate_all_with_ignore_null( default_value: &ScalarValue, is_lag: bool, ) -> Result<ArrayRef, DataFusionError> { + if offset == 0 { + return Ok(Arc::clone(array)); + } + // Arrays without NULLs do not necessarily have a null bitmap. let Some(nulls) = array.nulls() else { return shift_with_default_value(array, offset, default_value); }; - let valid_indices: Vec<usize> = nulls.valid_indices().collect::<Vec<_>>(); - let direction = !is_lag; - let new_array_results: Result<Vec<_>, DataFusionError> = (0..array.len()) - .map(|id| { - let result_index = match valid_indices.binary_search(&id) { - Ok(pos) => if direction { - pos.checked_add(offset as usize) - } else { - pos.checked_sub(offset.unsigned_abs() as usize) - } - .and_then(|new_pos| { - if new_pos < valid_indices.len() { - Some(valid_indices[new_pos]) - } else { - None - } - }), - Err(pos) => if direction { - pos.checked_add(offset as usize) - } else if pos > 0 { - pos.checked_sub(offset.unsigned_abs() as usize) - } else { - None - } - .and_then(|new_pos| { - if new_pos < valid_indices.len() { - Some(valid_indices[new_pos]) - } else { - None - } - }), + let shift = offset_magnitude(offset); + if shift >= array.len() { + return default_value.to_array_of_size(array.len()); + } + + let mut indices = ShiftIndexBuilder::new(array.len(), default_value.is_null()); + if is_lag { + let mut preceding = VecDeque::new(); + for index in 0..array.len() { + let result_index = if preceding.len() == shift { + preceding.front().copied() + } else { + None }; + indices.append_option(result_index); - match result_index { - Some(index) => ScalarValue::try_from_array(array, index), - None => Ok(default_value.clone()), + if nulls.is_valid(index) { + if preceding.len() == shift { + preceding.pop_front(); + } + preceding.push_back(index); + } + } + } else { + let mut following = VecDeque::new(); + let mut next_index = 0; + for index in 0..array.len() { + while following.front().is_some_and(|next| *next <= index) { + following.pop_front(); + } + next_index = next_index.max(index.saturating_add(1)); + while following.len() < shift && next_index < array.len() { + if nulls.is_valid(next_index) { + following.push_back(next_index); + } + next_index += 1; } - }) - .collect(); + indices.append_option(following.get(shift - 1).copied()); + } + } - let new_array = new_array_results?; - ScalarValue::iter_to_array(new_array) + indices.finish(array, default_value) } // TODO: change the original arrow::compute::kernels::window::shift impl to support an optional default value fn shift_with_default_value( @@ -692,7 +744,8 @@ impl PartitionEvaluator for WindowShiftEvaluator { mod tests { use super::*; use arrow::array::*; - use datafusion_common::cast::as_int32_array; + use arrow::datatypes::Int8Type; + use datafusion_common::cast::{as_dictionary_array, as_int32_array, as_string_array}; use datafusion_physical_expr::expressions::{Column, Literal}; fn test_i32_result( @@ -843,6 +896,117 @@ mod tests { ) } + #[test] + fn test_evaluate_all_with_ignore_null() -> Result<()> { + let input: ArrayRef = Arc::new(Int32Array::from(vec![ + None, + Some(10), + None, + Some(20), + Some(30), + None, + ])); + + let cases = [ + ( + 1, + ScalarValue::Int32(None), + Int32Array::from(vec![ + None, + None, + Some(10), + Some(10), + Some(20), + Some(30), + ]), + ), + ( + -1, + ScalarValue::Int32(None), + Int32Array::from(vec![ + Some(10), + Some(20), + Some(20), + Some(30), + None, + None, + ]), + ), + ( + 2, + ScalarValue::Int32(Some(-1)), + Int32Array::from(vec![ + Some(-1), + Some(-1), + Some(-1), + Some(-1), + Some(10), + Some(20), + ]), + ), + ( + -2, + ScalarValue::Int32(Some(-1)), + Int32Array::from(vec![ + Some(20), + Some(30), + Some(30), + Some(-1), + Some(-1), + Some(-1), + ]), + ), + ( + 0, + ScalarValue::Int32(Some(-1)), + Int32Array::from(vec![None, Some(10), None, Some(20), Some(30), None]), + ), + ]; + + for (offset, default_value, expected) in cases { + let actual = evaluate_all_with_ignore_null( + &input, + offset, + &default_value, + offset > 0, + )?; + assert_eq!(expected, *as_int32_array(&actual)?); + } + Ok(()) + } + + #[test] + fn test_ignore_nulls_dictionary_with_bounded_keys() -> Result<()> { + let keys = + Int8Array::from_iter(std::iter::once(None).chain((0_i8..=127).map(Some))); + let values = + StringArray::from_iter_values((0..128).map(|index| format!("value-{index}"))); + let input: ArrayRef = Arc::new(DictionaryArray::<Int8Type>::try_new( + keys, + Arc::new(values), + )?); + let default_value = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Utf8(Some("default".to_string()))), + ); + + let actual = evaluate_all_with_ignore_null(&input, 1, &default_value, true)?; + let actual = as_dictionary_array::<Int8Type>(actual.as_ref())?; + let values = as_string_array(actual.values().as_ref())?; + + assert_eq!(actual.len(), 129); + assert_eq!(values.len(), 128); + for index in 0..2 { + let key = actual.key(index).expect("non-null default"); + assert_eq!(values.value(key), "default"); + } + for index in 2..actual.len() { + let key = actual.key(index).expect("selected value"); + assert_eq!(values.value(key), format!("value-{}", index - 2)); + } + Ok(()) + } + #[test] fn test_ignore_nulls_without_null_bitmap() -> Result<()> { let input = Int32Array::from(vec![1, 2, 3]); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
