perf: make LAST_VALUE ... IGNORE NULLS linear in the partition size

datafusion/functions-window/src/nth_value.rs

@@ -19,8 +19,9 @@
 
 use crate::utils::{get_scalar_value_from_args, get_signed_integer};
 
-use arrow::buffer::NullBuffer;
+use arrow::buffer::BooleanBuffer;
 use arrow::datatypes::FieldRef;
+use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
 use datafusion_common::arrow::array::ArrayRef;
 use datafusion_common::arrow::datatypes::{DataType, Field};
 use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err};
@@ -487,14 +488,15 @@ impl PartitionEvaluator for NthValueEvaluator {
 impl NthValueEvaluator {
     fn valid_index(&self, array: &ArrayRef, range: &Range<usize>) -> Option<usize> {
         let n_range = range.end - range.start;
-        if self.ignore_nulls {
-            // Calculate valid indices, inside the window frame boundaries.
-            let slice = array.slice(range.start, n_range);
-            if let Some(nulls) = slice.nulls()
-                && nulls.null_count() > 0
-            {
-                return self.valid_index_with_nulls(nulls, range.start);
-            }
+        if self.ignore_nulls
+            && let Some(nulls) = array.nulls()
+        {
+            // Read the validity bits inside the window frame boundaries in place. Slicing
+            // `array` instead would allocate and count the frame's nulls for every row.
+            let valid = nulls.inner().slice(range.start, n_range);
+            return self
+                .valid_index_with_nulls(&valid)
+                .map(|idx| idx + range.start);
         }
         // Either no nulls, or nulls are regarded as valid rows
         match self.state.kind {
@@ -525,35 +527,55 @@ impl NthValueEvaluator {
         }
     }
 
-    fn valid_index_with_nulls(&self, nulls: &NullBuffer, offset: usize) -> Option<usize> {
+    /// Returns the index within the frame of the value to return, given the frame's validity
+    /// bits.
+    fn valid_index_with_nulls(&self, valid: &BooleanBuffer) -> Option<usize> {
         match self.state.kind {
-            NthValueKind::First => nulls.valid_indices().next().map(|idx| idx + offset),
-            NthValueKind::Last => nulls.valid_indices().last().map(|idx| idx + offset),
+            NthValueKind::First => valid.set_indices().next(),
+            NthValueKind::Last => nth_set_bit_from_end(valid, 1),
             NthValueKind::Nth => {
                 match self.n.cmp(&0) {
                     Ordering::Greater => {
                         // SQL indices are not 0-based.
                         let index = (self.n as usize) - 1;
-                        nulls.valid_indices().nth(index).map(|idx| idx + offset)
-                    }
-                    Ordering::Less => {
-                        let reverse_index = (-self.n) as usize;
-                        let valid_indices_len = nulls.len() - nulls.null_count();
-                        if reverse_index > valid_indices_len {
-                            return None;
-                        }
-                        nulls
-                            .valid_indices()
-                            .nth(valid_indices_len - reverse_index)
-                            .map(|idx| idx + offset)
+                        valid.set_indices().nth(index)
                     }
+                    Ordering::Less => nth_set_bit_from_end(valid, (-self.n) as usize),
                     Ordering::Equal => None,
                 }
             }
         }
     }
 }
 
+/// Returns the index of the `n`th set bit from the end of `bits`, where `n = 1` is the last
+/// set bit, or `None` if fewer than `n` bits are set.
+///
+/// The bits are read a 64-bit word at a time from the end, so the cost grows with the distance
+/// from the end to that bit rather than with the length of `bits`. Enumerating the set bits
+/// from the start instead is quadratic for a frame that grows with every row.
+fn nth_set_bit_from_end(bits: &BooleanBuffer, n: usize) -> Option<usize> {
+    debug_assert!(n > 0, "n is 1-based");
+    let chunk = UnalignedBitChunk::new(bits.values(), bits.offset(), bits.len());
+    // The words hold `lead_padding` unset bits before `bits` and `trailing_padding` after it.
+    let mut word_start = chunk.lead_padding() + bits.len() + chunk.trailing_padding();
+    let mut remaining = n;
+    for mut word in chunk.iter().rev() {
+        word_start -= 64;
+        let ones = word.count_ones() as usize;
+        if ones < remaining {
+            remaining -= ones;
+            continue;
+        }
+        // Clear the `remaining - 1` highest set bits, leaving the one we want on top
+        for _ in 1..remaining {
+            word ^= 1 << word.ilog2();
+        }
+        return Some(word_start + word.ilog2() as usize - chunk.lead_padding());
+    }
+    None
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -709,4 +731,40 @@ mod tests {
             )
         );
     }
+
+    #[test]
+    fn nth_set_bit_from_end_matches_a_scan_from_the_end() {
+        // The slices cover every bit offset within a word, byte-aligned and unaligned starts,
+        // and lengths from empty to several words, which are the cases `UnalignedBitChunk`
+        // splits into a prefix, aligned words and a suffix.
+        let mut state = 0x9E37_79B9_7F4A_7C15_u64;
+        let random = BooleanBuffer::collect_bool(640, |_| {
+            state ^= state << 13;
+            state ^= state >> 7;
+            state ^= state << 17;
+            state.is_multiple_of(5)
+        });
+        let patterns = [
+            random,
+            BooleanBuffer::collect_bool(640, |i| i.is_multiple_of(97)),
+            BooleanBuffer::new_set(640),
+            BooleanBuffer::new_unset(640),
+        ];
+        for bits in &patterns {
+            for offset in 0..70 {
+                for len in 0..=200 {
+                    let slice = bits.slice(offset, len);
+                    for n in [1, 2, 3, 63, 64, 65, 130] {
+                        let expected =
+                            (0..len).rev().filter(|&i| slice.value(i)).nth(n - 1);
+                        assert_eq!(
+                            nth_set_bit_from_end(&slice, n),
+                            expected,
+                            "offset {offset}, len {len}, n {n}"
+                        );
+                    }
+                }
+            }
+        }
+    }
 }

datafusion/sqllogictest/test_files/window.slt

@@ -4821,6 +4821,107 @@ NULL 1
 statement ok
 DROP TABLE t;
 
+# Test for ignore nulls in LAST_VALUE, FIRST_VALUE and NTH_VALUE over partitions that span many
+# batches, with runs of NULLs that cross several words of the validity bitmap. Partition 0
+# alternates NULLs, partition 1 has runs of 497 NULLs, partition 2 is all NULLs and partition 3
+# has none. `v` increases with `i` in every partition, so the last (first) non-NULL value of a
+# frame is also its maximum (minimum), which `max` (`min`) computes independently.
+statement ok
+set datafusion.execution.batch_size = 100;
+
+statement ok
+CREATE TABLE ignore_nulls_t AS
+SELECT i % 4 AS p, i,
+  CASE i % 4
+    WHEN 0 THEN CASE WHEN (i / 4) % 2 = 0 THEN i END
+    WHEN 1 THEN CASE WHEN (i / 4) % 500 < 3 THEN i END
+    WHEN 3 THEN i
+  END AS v
+FROM (SELECT value AS i FROM generate_series(1, 30000));
+
+query IIII
+SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END),
+  count(fv), sum(CASE WHEN fv IS DISTINCT FROM mn THEN 1 ELSE 0 END)
+FROM (SELECT
+  LAST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS lv,
+  max(v) OVER (PARTITION BY p ORDER BY i) AS mx,
+  FIRST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS fv,
+  min(v) OVER (PARTITION BY p ORDER BY i) AS mn
+FROM ignore_nulls_t);
+----
+22499 0 22499 0
+
+query IIII
+SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END),
+  count(fv), sum(CASE WHEN fv IS DISTINCT FROM mn THEN 1 ELSE 0 END)
+FROM (SELECT
+  LAST_VALUE(v) IGNORE NULLS OVER (
+    PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS lv,
+  max(v) OVER (
+    PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS mx,
+  FIRST_VALUE(v) IGNORE NULLS OVER (
+    PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS fv,
+  min(v) OVER (
+    PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS mn
+FROM ignore_nulls_t);
+----
+22499 0 22499 0
+
+query IIII
+SELECT count(n1), sum(CASE WHEN n1 IS DISTINCT FROM mn THEN 1 ELSE 0 END),
+  count(nl), sum(CASE WHEN nl IS DISTINCT FROM mx THEN 1 ELSE 0 END)
+FROM (SELECT
+  NTH_VALUE(v, 1) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS n1,
+  min(v) OVER (PARTITION BY p ORDER BY i) AS mn,
+  NTH_VALUE(v, -1) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS nl,
+  max(v) OVER (PARTITION BY p ORDER BY i) AS mx
+FROM ignore_nulls_t);
+----
+22499 0 22499 0
+
+# NTH_VALUE(v, -65) over a sliding frame whose valid rows span two bitmap words. Once full, the
+# frame holds exactly 65 valid rows in partition 0, more in partition 3 and fewer in partitions 1
+# and 2. The expected value is the valid row whose running count is 64 less than the current
+# row's, if the frame holds at least 65 valid rows.
+query III
+WITH r AS (SELECT p, v,
+  count(v) OVER (PARTITION BY p ORDER BY i) AS c,
+  count(v) OVER (PARTITION BY p ORDER BY i ROWS BETWEEN 129 PRECEDING AND CURRENT ROW) AS fc,
+  NTH_VALUE(v, -65) IGNORE NULLS
+    OVER (PARTITION BY p ORDER BY i ROWS BETWEEN 129 PRECEDING AND CURRENT ROW) AS nv
+FROM ignore_nulls_t)
+SELECT count(r.nv), sum(CASE WHEN r.fc = 65 THEN 1 ELSE 0 END),
+  sum(CASE WHEN r.nv IS DISTINCT FROM CASE WHEN r.fc >= 65 THEN e.v END THEN 1 ELSE 0 END)
+FROM r LEFT JOIN r e ON e.p = r.p AND e.v IS NOT NULL AND e.c = r.c - 64;
+----
+14807 7372 0
+
+query II
+SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END)
+FROM (SELECT
+  LAST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i
+    ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS lv,
+  max(v) OVER (PARTITION BY p ORDER BY i
+    ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS mx
+FROM ignore_nulls_t);
+----
+22500 0
+
+# With peers in a RANGE frame, the last non-NULL value comes from the last peer group that has
+# one. All non-NULL values of a peer group are equal here, so it is still the frame's maximum.
+query II
+SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END)
+FROM (SELECT
+  LAST_VALUE(CASE WHEN v IS NOT NULL THEN i / 28 END) IGNORE NULLS
+    OVER (PARTITION BY p ORDER BY i / 28) AS lv,
+  max(CASE WHEN v IS NOT NULL THEN i / 28 END) OVER (PARTITION BY p ORDER BY i / 28) AS mx
+FROM ignore_nulls_t);
+----
+22500 0
+
+statement ok
+DROP TABLE ignore_nulls_t;
+
 
 # Test for ignore nulls in nth_VALUE without null values
 statement ok