From 6711e7c40ee8f7f0a41f67c4e25400fc92cee53b Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Thu, 23 Jul 2026 18:18:11 +0800 Subject: [PATCH] fix: avoid dictionary overflow in LEAD/LAG IGNORE NULLS --- 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 2363a6beed63f..fea4a1a4aadda 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) { + 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 { + 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 { + 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 = nulls.valid_indices().collect::>(); - let direction = !is_lag; - let new_array_results: Result, 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::::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::(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]);