TODO: change the original arrow::compute::kernels::window::shift impl to support an optional default value
(
array: &ArrayRef,
offset: i64,
default_value: &ScalarValue,
)
| 479 | } |
| 480 | // TODO: change the original arrow::compute::kernels::window::shift impl to support an optional default value |
| 481 | fn shift_with_default_value( |
| 482 | array: &ArrayRef, |
| 483 | offset: i64, |
| 484 | default_value: &ScalarValue, |
| 485 | ) -> Result<ArrayRef> { |
| 486 | use datafusion_common::arrow::compute::concat; |
| 487 | |
| 488 | let value_len = array.len() as i64; |
| 489 | if offset == 0 { |
| 490 | Ok(Arc::clone(array)) |
| 491 | } else if offset == i64::MIN || offset.abs() >= value_len { |
| 492 | default_value.to_array_of_size(value_len as usize) |
| 493 | } else { |
| 494 | let slice_offset = (-offset).clamp(0, value_len) as usize; |
| 495 | let length = array.len() - offset.unsigned_abs() as usize; |
| 496 | let slice = array.slice(slice_offset, length); |
| 497 | |
| 498 | // Generate array with remaining `null` items |
| 499 | let nulls = offset.unsigned_abs() as usize; |
| 500 | let default_values = default_value.to_array_of_size(nulls)?; |
| 501 | |
| 502 | // Concatenate both arrays, add nulls after if shift > 0 else before |
| 503 | if offset > 0 { |
| 504 | concat(&[default_values.as_ref(), slice.as_ref()]) |
| 505 | .map_err(|e| arrow_datafusion_err!(e)) |
| 506 | } else { |
| 507 | concat(&[slice.as_ref(), default_values.as_ref()]) |
| 508 | .map_err(|e| arrow_datafusion_err!(e)) |
| 509 | } |
| 510 | } |
| 511 | } |
| 512 | |
| 513 | impl PartitionEvaluator for WindowShiftEvaluator { |
| 514 | fn get_range(&self, idx: usize, n_rows: usize) -> Result<Range<usize>> { |
no test coverage detected
searching dependent graphs…