This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 7d42518fb8 perf: optimize LEAD/LAG IGNORE NULLS evaluation (#23711)
7d42518fb8 is described below
commit 7d42518fb8e8a1bdb1f7b542cbedc55b9cbff73a
Author: xudong.w <[email protected]>
AuthorDate: Fri Jul 24 10:01:36 2026 +0800
perf: optimize LEAD/LAG IGNORE NULLS evaluation (#23711)
## Which issue does this PR close?
- N/A
## Rationale for this change
The whole-partition evaluation path for `LEAD` and `LAG` with `IGNORE
NULLS` currently collects every valid row index, performs a binary
search for every output row, converts every selected value to a
`ScalarValue`, and finally rebuilds an Arrow array from those scalars.
This makes index selection `O(n log m)` for `n` rows and `m` non-null
rows, and the per-row scalar materialization is particularly expensive
for strings and nested values.
A local Criterion microbenchmark with 100,000 rows and 50% nulls
measured the following speedups across the tested offsets:
| Data type | Speedup |
| --- | ---: |
| `Int64` | 6.7–6.9x |
| `Utf8View` | 10.5–11.6x |
| `List` | 21.7–21.8x |
## What changes are included in this PR?
- Generate nullable Arrow gather indices with a single linear scan and
bounded `VecDeque` state instead of collecting all valid indices and
binary-searching for every row.
- Materialize the result with Arrow's `take` kernel.
- Use Arrow's `zip` kernel to fill non-null default values without
constructing one `ScalarValue` per row.
- Keep dictionary arrays with non-null defaults on scalar
materialization because `zip` concatenates dictionaries and can overflow
bounded dictionary key types.
The index generation and output materialization are now `O(n)`. There
are no public API changes.
## Are these changes tested?
Yes.
## Are there any user-facing changes?
No. This changes the implementation of whole-partition `LEAD`/`LAG ...
IGNORE NULLS` evaluation without changing the SQL behavior or public
API.
---
datafusion/functions-window/src/lead_lag.rs | 242 +++++++++++++++++++++++-----
1 file changed, 203 insertions(+), 39 deletions(-)
diff --git a/datafusion/functions-window/src/lead_lag.rs
b/datafusion/functions-window/src/lead_lag.rs
index 2363a6beed..fea4a1a4aa 100644
--- a/datafusion/functions-window/src/lead_lag.rs
+++ b/datafusion/functions-window/src/lead_lag.rs
@@ -18,6 +18,8 @@
//! `lead` and `lag` window function implementations
use crate::utils::{get_scalar_value_from_args, get_signed_integer};
+use arrow::array::UInt64Builder;
+use arrow::compute::{interleave, take};
use arrow::datatypes::FieldRef;
use datafusion_common::arrow::array::ArrayRef;
use datafusion_common::arrow::datatypes::DataType;
@@ -419,6 +421,52 @@ fn offset_magnitude(offset: i64) -> usize {
}
}
+enum ShiftIndexBuilder {
+ Take(UInt64Builder),
+ Interleave(Vec<(usize, usize)>),
+}
+
+impl ShiftIndexBuilder {
+ fn new(capacity: usize, default_is_null: bool) -> Self {
+ if default_is_null {
+ Self::Take(UInt64Builder::with_capacity(capacity))
+ } else {
+ Self::Interleave(Vec::with_capacity(capacity))
+ }
+ }
+
+ fn append_option(&mut self, index: Option<usize>) {
+ match self {
+ Self::Take(indices) => {
+ indices.append_option(index.map(|index| index as u64));
+ }
+ Self::Interleave(indices) => {
+ // `interleave` receives `[array, default]`.
+ indices.push(index.map_or((1, 0), |index| (0, index)));
+ }
+ }
+ }
+
+ fn finish(
+ self,
+ array: &ArrayRef,
+ default_value: &ScalarValue,
+ ) -> Result<ArrayRef, DataFusionError> {
+ match self {
+ Self::Take(mut indices) => {
+ let indices = indices.finish();
+ take(array.as_ref(), &indices, None)
+ .map_err(|error| arrow_datafusion_err!(error))
+ }
+ Self::Interleave(indices) => {
+ let default = default_value.to_array_of_size(1)?;
+ interleave(&[array.as_ref(), default.as_ref()], &indices)
+ .map_err(|error| arrow_datafusion_err!(error))
+ }
+ }
+ }
+}
+
impl WindowShiftEvaluator {
fn is_lag(&self) -> bool {
// Mode is LAG, when shift_offset is positive
@@ -433,53 +481,57 @@ fn evaluate_all_with_ignore_null(
default_value: &ScalarValue,
is_lag: bool,
) -> Result<ArrayRef, DataFusionError> {
+ if offset == 0 {
+ return Ok(Arc::clone(array));
+ }
+
// Arrays without NULLs do not necessarily have a null bitmap.
let Some(nulls) = array.nulls() else {
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)
- }
- .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
- }
- .and_then(|new_pos| {
- if new_pos < valid_indices.len() {
- Some(valid_indices[new_pos])
- } else {
- None
- }
- }),
+ let shift = offset_magnitude(offset);
+ if shift >= array.len() {
+ return default_value.to_array_of_size(array.len());
+ }
+
+ let mut indices = ShiftIndexBuilder::new(array.len(),
default_value.is_null());
+ if is_lag {
+ let mut preceding = VecDeque::new();
+ for index in 0..array.len() {
+ let result_index = if preceding.len() == shift {
+ preceding.front().copied()
+ } else {
+ None
};
+ indices.append_option(result_index);
- match result_index {
- Some(index) => ScalarValue::try_from_array(array, index),
- None => Ok(default_value.clone()),
+ if nulls.is_valid(index) {
+ if preceding.len() == shift {
+ preceding.pop_front();
+ }
+ preceding.push_back(index);
+ }
+ }
+ } else {
+ let mut following = VecDeque::new();
+ let mut next_index = 0;
+ for index in 0..array.len() {
+ while following.front().is_some_and(|next| *next <= index) {
+ following.pop_front();
+ }
+ next_index = next_index.max(index.saturating_add(1));
+ while following.len() < shift && next_index < array.len() {
+ if nulls.is_valid(next_index) {
+ following.push_back(next_index);
+ }
+ next_index += 1;
}
- })
- .collect();
+ indices.append_option(following.get(shift - 1).copied());
+ }
+ }
- let new_array = new_array_results?;
- ScalarValue::iter_to_array(new_array)
+ indices.finish(array, default_value)
}
// TODO: change the original arrow::compute::kernels::window::shift impl to
support an optional default value
fn shift_with_default_value(
@@ -692,7 +744,8 @@ impl PartitionEvaluator for WindowShiftEvaluator {
mod tests {
use super::*;
use arrow::array::*;
- use datafusion_common::cast::as_int32_array;
+ use arrow::datatypes::Int8Type;
+ use datafusion_common::cast::{as_dictionary_array, as_int32_array,
as_string_array};
use datafusion_physical_expr::expressions::{Column, Literal};
fn test_i32_result(
@@ -843,6 +896,117 @@ mod tests {
)
}
+ #[test]
+ fn test_evaluate_all_with_ignore_null() -> Result<()> {
+ let input: ArrayRef = Arc::new(Int32Array::from(vec![
+ None,
+ Some(10),
+ None,
+ Some(20),
+ Some(30),
+ None,
+ ]));
+
+ let cases = [
+ (
+ 1,
+ ScalarValue::Int32(None),
+ Int32Array::from(vec![
+ None,
+ None,
+ Some(10),
+ Some(10),
+ Some(20),
+ Some(30),
+ ]),
+ ),
+ (
+ -1,
+ ScalarValue::Int32(None),
+ Int32Array::from(vec![
+ Some(10),
+ Some(20),
+ Some(20),
+ Some(30),
+ None,
+ None,
+ ]),
+ ),
+ (
+ 2,
+ ScalarValue::Int32(Some(-1)),
+ Int32Array::from(vec![
+ Some(-1),
+ Some(-1),
+ Some(-1),
+ Some(-1),
+ Some(10),
+ Some(20),
+ ]),
+ ),
+ (
+ -2,
+ ScalarValue::Int32(Some(-1)),
+ Int32Array::from(vec![
+ Some(20),
+ Some(30),
+ Some(30),
+ Some(-1),
+ Some(-1),
+ Some(-1),
+ ]),
+ ),
+ (
+ 0,
+ ScalarValue::Int32(Some(-1)),
+ Int32Array::from(vec![None, Some(10), None, Some(20),
Some(30), None]),
+ ),
+ ];
+
+ for (offset, default_value, expected) in cases {
+ let actual = evaluate_all_with_ignore_null(
+ &input,
+ offset,
+ &default_value,
+ offset > 0,
+ )?;
+ assert_eq!(expected, *as_int32_array(&actual)?);
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn test_ignore_nulls_dictionary_with_bounded_keys() -> Result<()> {
+ let keys =
+
Int8Array::from_iter(std::iter::once(None).chain((0_i8..=127).map(Some)));
+ let values =
+ StringArray::from_iter_values((0..128).map(|index|
format!("value-{index}")));
+ let input: ArrayRef = Arc::new(DictionaryArray::<Int8Type>::try_new(
+ keys,
+ Arc::new(values),
+ )?);
+ let default_value = ScalarValue::Dictionary(
+ Box::new(DataType::Int8),
+ Box::new(ScalarValue::Utf8(Some("default".to_string()))),
+ );
+
+ let actual = evaluate_all_with_ignore_null(&input, 1, &default_value,
true)?;
+ let actual = as_dictionary_array::<Int8Type>(actual.as_ref())?;
+ let values = as_string_array(actual.values().as_ref())?;
+
+ assert_eq!(actual.len(), 129);
+ assert_eq!(values.len(), 128);
+ for index in 0..2 {
+ let key = actual.key(index).expect("non-null default");
+ assert_eq!(values.value(key), "default");
+ }
+ for index in 2..actual.len() {
+ let key = actual.key(index).expect("selected value");
+ assert_eq!(values.value(key), format!("value-{}", index - 2));
+ }
+ Ok(())
+ }
+
#[test]
fn test_ignore_nulls_without_null_bitmap() -> Result<()> {
let input = Int32Array::from(vec![1, 2, 3]);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
