diff --git a/datafusion/functions-window/src/lead_lag.rs b/datafusion/functions-window/src/lead_lag.rs index 2363a6beed63..781563d33292 100644 --- a/datafusion/functions-window/src/lead_lag.rs +++ b/datafusion/functions-window/src/lead_lag.rs @@ -18,6 +18,9 @@ //! `lead` and `lag` window function implementations use crate::utils::{get_scalar_value_from_args, get_signed_integer}; +use arrow::array::UInt64Builder; +use arrow::compute::kernels::zip::zip; +use arrow::compute::{is_not_null, take}; use arrow::datatypes::FieldRef; use datafusion_common::arrow::array::ArrayRef; use datafusion_common::arrow::datatypes::DataType; @@ -438,48 +441,70 @@ fn evaluate_all_with_ignore_null( 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) + let shift = if is_lag { + offset_magnitude(offset) + } else { + offset as usize + }; + if shift >= array.len() { + return default_value.to_array_of_size(array.len()); + } + + let mut indices = UInt64Builder::with_capacity(array.len()); + if is_lag { + let mut preceding = VecDeque::new(); + for index in 0..array.len() { + indices + .append_option((preceding.len() == shift).then(|| preceding[0] as u64)); + if nulls.is_valid(index) { + if preceding.len() == shift { + preceding.pop_front(); } - .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 + preceding.push_back(index); + } + } + } else { + let mut following = VecDeque::new(); + let mut next_index = 0; + for index in 0..array.len() { + if following.front().is_some_and(|next| *next < index) { + following.pop_front(); + } + while following.len() <= shift && next_index < array.len() { + if nulls.is_valid(next_index) { + following.push_back(next_index); } - .and_then(|new_pos| { - if new_pos < valid_indices.len() { - Some(valid_indices[new_pos]) - } else { - None - } - }), - }; + next_index += 1; + } + indices.append_option(following.get(shift).map(|index| *index as u64)); + } + } - match result_index { - Some(index) => ScalarValue::try_from_array(array, index), + let indices = indices.finish(); + // `zip` concatenates dictionaries, which can overflow bounded key types. + if matches!(array.data_type(), DataType::Dictionary(_, _)) && !default_value.is_null() + { + let values = indices + .iter() + .map(|index| match index { + Some(index) => ScalarValue::try_from_array(array, index as usize), None => Ok(default_value.clone()), - } - }) - .collect(); + }) + .collect::, DataFusionError>>()?; + return ScalarValue::iter_to_array(values); + } + + let shifted = take(array.as_ref(), &indices, None) + .map_err(|error| arrow_datafusion_err!(error))?; + if default_value.is_null() { + return Ok(shifted); + } - let new_array = new_array_results?; - ScalarValue::iter_to_array(new_array) + let has_value = + is_not_null(&indices).map_err(|error| arrow_datafusion_err!(error))?; + let shifted: &dyn arrow::array::Array = shifted.as_ref(); + zip(&has_value, &shifted, &default_value.to_scalar()?) + .map_err(|error| arrow_datafusion_err!(error)) } // TODO: change the original arrow::compute::kernels::window::shift impl to support an optional default value fn shift_with_default_value(