perf: optimize LEAD/LAG IGNORE NULLS evaluation

datafusion/functions-window/src/lead_lag.rs

@@ -18,6 +18,9 @@
 //! `lead` and `lag` window function implementations
 
 use crate::utils::{get_scalar_value_from_args, get_signed_integer};
+use arrow::array::UInt64Builder;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{is_not_null, take};
 use arrow::datatypes::FieldRef;
 use datafusion_common::arrow::array::ArrayRef;
 use datafusion_common::arrow::datatypes::DataType;
@@ -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()?)
+        .map_err(|error| arrow_datafusion_err!(error))
 }
 // TODO: change the original arrow::compute::kernels::window::shift impl to support an optional default value
 fn shift_with_default_value(