Skip to content

Commit ee5b4e3

Browse files
compheadcomphead
committed
perf: make LAST_VALUE ... IGNORE NULLS linear in the partition size
For every row, `NthValueEvaluator::valid_index` sliced the argument to the frame and enumerated the frame's valid rows from its start to find the last one. A frame that starts at the partition's first row grows with every row, so a partition took time quadratic in its size. Read the frame's validity bits in place, and find the last valid row from the end of the frame, a 64-bit word at a time.
1 parent 3ee3e39 commit ee5b4e3

2 files changed

Lines changed: 157 additions & 18 deletions

File tree

‎datafusion/functions-window/src/nth_value.rs‎

Lines changed: 73 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@
1919
2020
use crate::utils::{get_scalar_value_from_args, get_signed_integer};
2121

22-
use arrow::buffer::NullBuffer;
22+
use arrow::buffer::BooleanBuffer;
2323
use arrow::datatypes::FieldRef;
24+
use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
2425
use datafusion_common::arrow::array::ArrayRef;
2526
use datafusion_common::arrow::datatypes::{DataType, Field};
2627
use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err};
@@ -487,14 +488,15 @@ impl PartitionEvaluator for NthValueEvaluator {
487488
impl NthValueEvaluator {
488489
fn valid_index(&self, array: &ArrayRef, range: &Range<usize>) -> Option<usize> {
489490
let n_range = range.end - range.start;
490-
if self.ignore_nulls {
491-
// Calculate valid indices, inside the window frame boundaries.
492-
let slice = array.slice(range.start, n_range);
493-
if let Some(nulls) = slice.nulls()
494-
&& nulls.null_count() > 0
495-
{
496-
return self.valid_index_with_nulls(nulls, range.start);
497-
}
491+
if self.ignore_nulls
492+
&& let Some(nulls) = array.nulls()
493+
{
494+
// Read the validity bits inside the window frame boundaries in place. Slicing
495+
// `array` instead would allocate and count the frame's nulls for every row.
496+
let valid = nulls.inner().slice(range.start, n_range);
497+
return self
498+
.valid_index_with_nulls(&valid)
499+
.map(|idx| idx + range.start);
498500
}
499501
// Either no nulls, or nulls are regarded as valid rows
500502
match self.state.kind {
@@ -525,27 +527,26 @@ impl NthValueEvaluator {
525527
}
526528
}
527529

528-
fn valid_index_with_nulls(&self, nulls: &NullBuffer, offset: usize) -> Option<usize> {
530+
/// Returns the index within the frame of the value to return, given the frame's validity
531+
/// bits.
532+
fn valid_index_with_nulls(&self, valid: &BooleanBuffer) -> Option<usize> {
529533
match self.state.kind {
530-
NthValueKind::First => nulls.valid_indices().next().map(|idx| idx + offset),
531-
NthValueKind::Last => nulls.valid_indices().last().map(|idx| idx + offset),
534+
NthValueKind::First => valid.set_indices().next(),
535+
NthValueKind::Last => last_set_bit(valid),
532536
NthValueKind::Nth => {
533537
match self.n.cmp(&0) {
534538
Ordering::Greater => {
535539
// SQL indices are not 0-based.
536540
let index = (self.n as usize) - 1;
537-
nulls.valid_indices().nth(index).map(|idx| idx + offset)
541+
valid.set_indices().nth(index)
538542
}
539543
Ordering::Less => {
540544
let reverse_index = (-self.n) as usize;
541-
let valid_indices_len = nulls.len() - nulls.null_count();
545+
let valid_indices_len = valid.count_set_bits();
542546
if reverse_index > valid_indices_len {
543547
return None;
544548
}
545-
nulls
546-
.valid_indices()
547-
.nth(valid_indices_len - reverse_index)
548-
.map(|idx| idx + offset)
549+
valid.set_indices().nth(valid_indices_len - reverse_index)
549550
}
550551
Ordering::Equal => None,
551552
}
@@ -554,6 +555,27 @@ impl NthValueEvaluator {
554555
}
555556
}
556557

558+
/// Returns the index of the last set bit in `bits`, or `None` if no bit is set.
559+
///
560+
/// The bits are read a 64-bit word at a time from the end, so the cost grows with the distance
561+
/// from the end to the last set bit rather than with the length of `bits`. `LAST_VALUE ...
562+
/// IGNORE NULLS` looks up the last valid row of every frame, and a frame that starts at the
563+
/// partition's first row grows with every row, so enumerating the frame's valid rows from its
564+
/// start makes the partition quadratic.
565+
fn last_set_bit(bits: &BooleanBuffer) -> Option<usize> {
566+
let chunk = UnalignedBitChunk::new(bits.values(), bits.offset(), bits.len());
567+
// The words hold `lead_padding` unset bits before `bits` and `trailing_padding` after it.
568+
let mut word_start = chunk.lead_padding() + bits.len() + chunk.trailing_padding();
569+
for word in chunk.iter().rev() {
570+
word_start -= 64;
571+
if word != 0 {
572+
let bit = 63 - word.leading_zeros() as usize;
573+
return Some(word_start + bit - chunk.lead_padding());
574+
}
575+
}
576+
None
577+
}
578+
557579
#[cfg(test)]
558580
mod tests {
559581
use super::*;
@@ -709,4 +731,37 @@ mod tests {
709731
)
710732
);
711733
}
734+
735+
#[test]
736+
fn last_set_bit_matches_a_scan_from_the_end() {
737+
// The slices cover every bit offset within a word, byte-aligned and unaligned starts,
738+
// and lengths from empty to several words, which are the cases `UnalignedBitChunk`
739+
// splits into a prefix, aligned words and a suffix.
740+
let mut state = 0x9E37_79B9_7F4A_7C15_u64;
741+
let random = BooleanBuffer::collect_bool(640, |_| {
742+
state ^= state << 13;
743+
state ^= state >> 7;
744+
state ^= state << 17;
745+
state.is_multiple_of(5)
746+
});
747+
let patterns = [
748+
random,
749+
BooleanBuffer::collect_bool(640, |i| i.is_multiple_of(97)),
750+
BooleanBuffer::new_set(640),
751+
BooleanBuffer::new_unset(640),
752+
];
753+
for bits in &patterns {
754+
for offset in 0..70 {
755+
for len in 0..=200 {
756+
let slice = bits.slice(offset, len);
757+
let expected = (0..len).rev().find(|&i| slice.value(i));
758+
assert_eq!(
759+
last_set_bit(&slice),
760+
expected,
761+
"offset {offset}, len {len}"
762+
);
763+
}
764+
}
765+
}
766+
}
712767
}

‎datafusion/sqllogictest/test_files/window.slt‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4821,6 +4821,90 @@ NULL 1
48214821
statement ok
48224822
DROP TABLE t;
48234823

4824+
# Test for ignore nulls in LAST_VALUE, FIRST_VALUE and NTH_VALUE over partitions that span many
4825+
# batches, with runs of NULLs that cross several words of the validity bitmap. Partition 0
4826+
# alternates NULLs, partition 1 has runs of 497 NULLs, partition 2 is all NULLs and partition 3
4827+
# has none. `v` increases with `i` in every partition, so the last (first) non-NULL value of a
4828+
# frame is also its maximum (minimum), which `max` (`min`) computes independently.
4829+
statement ok
4830+
set datafusion.execution.batch_size = 100;
4831+
4832+
statement ok
4833+
CREATE TABLE ignore_nulls_t AS
4834+
SELECT i % 4 AS p, i,
4835+
CASE i % 4
4836+
WHEN 0 THEN CASE WHEN (i / 4) % 2 = 0 THEN i END
4837+
WHEN 1 THEN CASE WHEN (i / 4) % 500 < 3 THEN i END
4838+
WHEN 3 THEN i
4839+
END AS v
4840+
FROM (SELECT value AS i FROM generate_series(1, 30000));
4841+
4842+
query IIII
4843+
SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END),
4844+
count(fv), sum(CASE WHEN fv IS DISTINCT FROM mn THEN 1 ELSE 0 END)
4845+
FROM (SELECT
4846+
LAST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS lv,
4847+
max(v) OVER (PARTITION BY p ORDER BY i) AS mx,
4848+
FIRST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS fv,
4849+
min(v) OVER (PARTITION BY p ORDER BY i) AS mn
4850+
FROM ignore_nulls_t);
4851+
----
4852+
22499 0 22499 0
4853+
4854+
query IIII
4855+
SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END),
4856+
count(fv), sum(CASE WHEN fv IS DISTINCT FROM mn THEN 1 ELSE 0 END)
4857+
FROM (SELECT
4858+
LAST_VALUE(v) IGNORE NULLS OVER (
4859+
PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS lv,
4860+
max(v) OVER (
4861+
PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS mx,
4862+
FIRST_VALUE(v) IGNORE NULLS OVER (
4863+
PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS fv,
4864+
min(v) OVER (
4865+
PARTITION BY p ORDER BY i ROWS BETWEEN 1000 PRECEDING AND CURRENT ROW) AS mn
4866+
FROM ignore_nulls_t);
4867+
----
4868+
22499 0 22499 0
4869+
4870+
query IIII
4871+
SELECT count(n1), sum(CASE WHEN n1 IS DISTINCT FROM mn THEN 1 ELSE 0 END),
4872+
count(nl), sum(CASE WHEN nl IS DISTINCT FROM mx THEN 1 ELSE 0 END)
4873+
FROM (SELECT
4874+
NTH_VALUE(v, 1) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS n1,
4875+
min(v) OVER (PARTITION BY p ORDER BY i) AS mn,
4876+
NTH_VALUE(v, -1) IGNORE NULLS OVER (PARTITION BY p ORDER BY i) AS nl,
4877+
max(v) OVER (PARTITION BY p ORDER BY i) AS mx
4878+
FROM ignore_nulls_t);
4879+
----
4880+
22499 0 22499 0
4881+
4882+
query II
4883+
SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END)
4884+
FROM (SELECT
4885+
LAST_VALUE(v) IGNORE NULLS OVER (PARTITION BY p ORDER BY i
4886+
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS lv,
4887+
max(v) OVER (PARTITION BY p ORDER BY i
4888+
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS mx
4889+
FROM ignore_nulls_t);
4890+
----
4891+
22500 0
4892+
4893+
# With peers in a RANGE frame, the last non-NULL value comes from the last peer group that has
4894+
# one. All non-NULL values of a peer group are equal here, so it is still the frame's maximum.
4895+
query II
4896+
SELECT count(lv), sum(CASE WHEN lv IS DISTINCT FROM mx THEN 1 ELSE 0 END)
4897+
FROM (SELECT
4898+
LAST_VALUE(CASE WHEN v IS NOT NULL THEN i / 28 END) IGNORE NULLS
4899+
OVER (PARTITION BY p ORDER BY i / 28) AS lv,
4900+
max(CASE WHEN v IS NOT NULL THEN i / 28 END) OVER (PARTITION BY p ORDER BY i / 28) AS mx
4901+
FROM ignore_nulls_t);
4902+
----
4903+
22500 0
4904+
4905+
statement ok
4906+
DROP TABLE ignore_nulls_t;
4907+
48244908

48254909
# Test for ignore nulls in nth_VALUE without null values
48264910
statement ok

0 commit comments

Comments
 (0)