diff --git a/datafusion/functions-window/src/nth_value.rs b/datafusion/functions-window/src/nth_value.rs index b3678e80f2273..6724dbf9c6a60 100644 --- a/datafusion/functions-window/src/nth_value.rs +++ b/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) -> Option { 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,28 +527,20 @@ impl NthValueEvaluator { } } - fn valid_index_with_nulls(&self, nulls: &NullBuffer, offset: usize) -> Option { + /// 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 { 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, } } @@ -554,6 +548,34 @@ impl NthValueEvaluator { } } +/// 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 { + 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}" + ); + } + } + } + } + } } diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 7c92514e723bd..81da6d91fcfd7 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/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