xudong963 commented on code in PR #23711:
URL: https://github.com/apache/datafusion/pull/23711#discussion_r3637197004
##########
datafusion/functions-window/src/lead_lag.rs:
##########
@@ -438,48 +441,70 @@ fn evaluate_all_with_ignore_null(
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)
+ 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::<Result<Vec<_>, 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()?)
Review Comment:
Yes, direct interleave looks preferable here. Including index construction,
it was about 1.55x faster for Int64, 1.41x for Utf8View, and 3.81x for List in
my tested workload. I’ll switch the non-null-default path to build the
interleave indices directly, avoiding the intermediate take result and
subsequent zip.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
