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

[comphead]

Which issue does this PR close?

• Closes LAST_VALUE ... IGNORE NULLS is quadratic in the partition size #26532.

Rationale for this change

LAST_VALUE(expr) IGNORE NULLS over a frame that starts at the first row of the partition, such as the default RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, took time quadratic in the partition size. A 1M-row partition took 89 s, against 0.1 s with RESPECT NULLS. The issue has the reproducer.

What changes are included in this PR?

NthValueEvaluator::valid_index no longer slices the argument array for every row. It reads the frame's validity bits in place, since BooleanBuffer::slice is O(1), and for LAST_VALUE it finds the last valid row with a new last_set_bit. That function reads the bitmap a 64-bit word at a time from the end of the frame, so the cost per row is the distance back to the last non-NULL row instead of the size of the frame. FIRST_VALUE and NTH_VALUE keep their forward scans, but they also skip the per-row slice and null count. The results are unchanged: a frame without NULLs still resolves to its first, last or nth row.

Two related costs are not changed here:

• memoize still returns early for LAST_VALUE with IGNORE NULLS, so the frame start never moves and BoundedWindowAggExec keeps the whole partition buffered. That is LAST_VALUE and NTH_VALUE with IGNORE NULLS keep the whole partition buffered #26539.
• NTH_VALUE(expr, -n) IGNORE NULLS still counts the frame's valid rows for every row. That is NTH_VALUE(expr, -n) IGNORE NULLS is quadratic in the partition size #26538.

What is the testing strategy for this PR?

New window.slt cases cover LAST_VALUE, FIRST_VALUE and NTH_VALUE with IGNORE NULLS over partitions that span many 100-row batches. The partitions hold alternating NULLs, NULL runs across several bitmap words, only NULLs, or no NULLs. The cases use growing, sliding, whole-partition and peer-group frames. v increases within each partition, so every case is checked against max or min over the same frame, which computes the same value through the aggregate path. A unit test checks last_set_bit against a scan from the end for every bit offset up to 70 and every length up to 200. Both kinds of test fail if last_set_bit ignores the lead padding.

cargo bench -p datafusion-functions-window --bench nth_value, main against this branch on an Apple M5 Max:

Benchmark	main	This PR
last_value_expanding/30%_nulls	13.15 ms	111 µs
last_value_expanding/50%_nulls	8.96 ms	109 µs
last_value_expanding/80%_nulls	2.51 ms	109 µs
last_value_sliding_100/50%_nulls	602 µs	113 µs
first_value_expanding/50%_nulls	364 µs	119 µs
nth_value_10_expanding/50%_nulls	379 µs	137 µs
nth_value_neg10_expanding/50%_nulls	7.95 ms	7.65 ms
first_value/respect_nulls	47.8 µs	46.5 µs

The issue's query in datafusion-cli (release-nonlto profile), 1M rows:

Rows per partition	main	This PR
10,000	0.13 s	0.03 s
100,000	1.90 s	0.04 s
1,000,000	89.2 s	0.16 s

Are there any user-facing changes?

No, apart from the speed-up.

[andygrove]

Thanks @comphead, this is a nice fix with a very clear description. Two things in particular made it easy to review:

• The SLT cases use max/min over the same frame as an independent oracle, so all 30,000 rows are checked across four NULL patterns and five frame shapes without hand-written expected output.
• last_set_bit_matches_a_scan_from_the_end sweeps bit offsets 0–69 and lengths 0–200, which reaches every branch of UnalignedBitChunk::new (≤ 8 bytes, ≤ 16 bytes, and align_to with and without a prefix or suffix).

What I checked:

• last_set_bit: I traced UnalignedBitChunk::new in arrow-buffer 60. In every branch, lead_padding + len + trailing_padding is 64 × the number of words, and the prefix and suffix masks clear the bits outside the slice. So word_start + bit - lead_padding always lands inside bits.

• Results are unchanged: a throwaway differential test compared the new valid_index with the one from 3ee3e39 on 29.7M cases:

  • random arrays of up to 320 rows, sliced at bit offsets 0–70
  • NULL density from 0% to 100%, including arrays whose validity buffer has no NULLs
  • every frame for arrays of up to 90 rows, and random frames above that
  • FIRST_VALUE, LAST_VALUE and NTH_VALUE with n in {±1, ±2, ±3, ±7, ±10, ±64, ±65, 130}
  • both IGNORE NULLS and RESPECT NULLS

  They agreed on every case.

• Mutation: dropping - chunk.lead_padding() makes the unit test fail, and makes the new SLT block fail when run on its own (a panic in BooleanBuffer::value), as the description says. The new block passes on main too, which is expected for a perf-only change.

• Tests: window.slt passes locally.

• Benchmarks: cargo bench -p datafusion-functions-window --bench nth_value on an Apple M3 Max, median of 4 alternating runs of saved main and PR binaries:

  Benchmark	main	This PR
  last_value_expanding/30%_nulls	15.78 ms	144 µs
  last_value_expanding/50%_nulls	9.37 ms	144 µs
  last_value_expanding/80%_nulls	3.17 ms	145 µs
  last_value_sliding_100/50%_nulls	723 µs	143 µs
  first_value_expanding/50%_nulls	457 µs	151 µs
  nth_value_10_expanding/50%_nulls	481 µs	183 µs
  nth_value_neg10_expanding/50%_nulls	7.66 ms	8.58 ms
  first_value_expanding/0%_nulls	233 µs	85 µs
  first_value/respect_nulls	59.0 µs	59.0 µs

  The 0% rows get faster too, because valid_index no longer slices the array before checking for a validity buffer. nth_value_neg10_expanding/50%_nulls was slower in each of the 4 rounds here, unlike your M5 Max run. The suggestion inline brings it to about 0.2 ms either way.

Before approving, I'd like the NTH_VALUE(expr, -n) IGNORE NULLS regression on frames without NULLs fixed. It's narrow, but this PR introduces it. Details, numbers and a suggested generalization of last_set_bit are in the inline comment. The generalization also takes care of the second follow-up in your description. Everything else looks good to me.

Follow-up: could you file an issue for the memoize early return? It skips memoizing both LAST_VALUE and NTH_VALUE under IGNORE NULLS, so with an UNBOUNDED PRECEDING frame the frame start never moves and BoundedWindowAggExec buffers the whole partition. That way it won't get lost once this merges.

[andygrove]

This arm regresses NTH_VALUE(expr, -n) IGNORE NULLS on frames that have no NULLs while the argument still has a validity buffer. On main such a frame failed the null_count() > 0 check and fell through to the O(1) Some(range.end - reverse_index) path. Here every frame takes the bitmap path, so set_indices().nth(valid_indices_len - reverse_index) steps through every valid row of the frame. That's O(frame) bit iteration per row, where main did an O(frame / 64) popcount.

That's the common case with sparse NULLs: most frames have none, but nearly every batch carries a validity buffer. Negative n also comes from the planner, not just from users who write it. get_best_fitting_window reverses a window to reuse an existing sort order, and partition_evaluator then flips the sign of n (is_reversed()). So a sliding NTH_VALUE(x, 3) can run as -3.

From the nth_value bench on an Apple M3 Max. The first three rows are cases I added locally: "0.1% nulls" is create_primitive_array(8192, 0.001), and "null at end" is 8192 rows where only the last one is NULL. Medians of 4 alternating runs (3 for the last column):

Benchmark	main	This PR	With nth_set_bit_from_end below
nth_value_neg10_sliding_100/0.1%_nulls	336 µs	559 µs	179 µs
nth_value_neg10_sliding_1000/0.1%_nulls	2.13 ms	3.53 ms	194 µs
nth_value_neg10_expanding/null_at_end	395 µs	13.6 ms	176 µs
nth_value_neg10_expanding/30%_nulls	14.3 ms	14.2 ms	186 µs
nth_value_neg10_expanding/50%_nulls	7.66 ms	8.58 ms	196 µs
last_value_expanding/50%_nulls	9.37 ms	144 µs	156 µs
last_value_sliding_100/50%_nulls	723 µs	143 µs	154 µs

Generalizing last_set_bit to the n-th set bit from the end fixes this, and it also covers the second follow-up in the description:

/// 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
}

Call it as NthValueKind::Last => nth_set_bit_from_end(valid, 1) and Ordering::Less => nth_set_bit_from_end(valid, (-self.n) as usize). I checked it three ways:

• your unit test extended to n in [1, 2, 3, 63, 64, 65, 130] against (0..len).rev().filter(|&i| slice.value(i)).nth(n - 1)
• the differential test described in the review summary
• window.slt

For LAST_VALUE the extra popcount costs about 1.4 ns per row (the last two rows above). If you'd rather not pay that, Last can keep calling last_set_bit. I used ilog2 because clippy's manual_ilog2 fires on 63 - word.leading_zeros() inside the shift. In last_set_bit the as usize cast only hides the lint.

If you'd rather keep this PR to LAST_VALUE, the smallest fix is to restore main's shortcut in this arm. This also passes the differential test:

if valid_indices_len == valid.len() {
    return Some(valid.len() - reverse_index);
}

Bench cases I added, with the updated criterion_group!

fn bench_sparse_nulls(c: &mut Criterion) {
    let mut group = c.benchmark_group("nth_value_sparse_nulls");
    // ~8 NULLs in 8192 rows: a validity buffer exists, but most sliding frames have no NULLs
    let sparse =
        Arc::new(create_primitive_array::<Int64Type>(ARRAY_SIZE, 0.001)) as ArrayRef;
    // Only the last row is NULL: every expanding frame but the last has no NULLs
    let null_at_end = Arc::new(
        (0..ARRAY_SIZE as i64)
            .map(|i| (i != ARRAY_SIZE as i64 - 1).then_some(i))
            .collect::<arrow::array::Int64Array>(),
    ) as ArrayRef;
    for window_size in [100usize, 1000] {
        group.bench_function(
            BenchmarkId::new(
                format!("nth_value_neg10_sliding_{window_size}"),
                "0.1%_nulls",
            ),
            |b| {
                b.iter(|| {
                    let mut evaluator =
                        create_evaluator(NthValueKind::Nth, true, Some(-10));
                    let values_slice = slice::from_ref(&sparse);
                    for i in 0..sparse.len() {
                        let start = i.saturating_sub(window_size - 1);
                        let range = Range { start, end: i + 1 };
                        black_box(evaluator.evaluate(values_slice, &range).unwrap());
                    }
                })
            },
        );
    }
    group.bench_function(
        BenchmarkId::new("nth_value_neg10_expanding", "null_at_end"),
        |b| {
            b.iter(|| {
                let mut evaluator = create_evaluator(NthValueKind::Nth, true, Some(-10));
                let values_slice = slice::from_ref(&null_at_end);
                for i in 0..null_at_end.len() {
                    let range = Range {
                        start: 0,
                        end: i + 1,
                    };
                    black_box(evaluator.evaluate(values_slice, &range).unwrap());
                }
            })
        },
    );
    group.finish();
}

criterion_group!(benches, bench_nth_value_ignore_nulls, bench_sparse_nulls);

[comphead]

Thanks @andygrove, good catch. I took your nth_set_bit_from_end as written in 70fa4d9. It now covers both LAST_VALUE and NTH_VALUE(expr, -n), and the unit test checks n in [1, 2, 3, 63, 64, 65, 130]. I kept a single helper. The extra popcount costs LAST_VALUE about 1 ns per row here too.

nth_value bench on an Apple M5 Max, the previous head against this commit:

Benchmark	a717eb4	70fa4d9
nth_value_neg10_expanding/30%_nulls	13.40 ms	147 µs
nth_value_neg10_expanding/50%_nulls	7.64 ms	156 µs
nth_value_neg10_expanding/80%_nulls	2.41 ms	156 µs
last_value_expanding/50%_nulls	113 µs	120 µs
last_value_sliding_100/50%_nulls	116 µs	126 µs

window.slt and the extended test suite pass. I filed #26537 for the memoize early return and updated the description.

[comphead]

nice improvement for the PR, thanks @andygrove