Jefffrey commented on code in PR #23711:
URL: https://github.com/apache/datafusion/pull/23711#discussion_r3636051976
##########
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:
i wonder if doing interleave instead of take + zip for the non-null default
value path would be more efficient? or take+zip is plenty efficient enough 🤔
##########
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()
Review Comment:
is this a limitation of zip or a bug with it?
--
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]
