xudong963 commented on code in PR #23711:
URL: https://github.com/apache/datafusion/pull/23711#discussion_r3637195156
##########
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:
This is a current implementation bug/limitation in Arrow’s
zip/MutableArrayData dictionary path, rather than an inherent limitation of zip
semantics. It concatenates the complete backing dictionaries and can panic on
bounded-key overflow even when the selected output would fit. `interleave`
already has specialized dictionary merging/re-encoding logic. I’ll add a
bounded-dictionary regression test as well.
--
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]
