andygrove commented on code in PR #26533:
URL: https://github.com/apache/datafusion/pull/26533#discussion_r4329483242
##########
datafusion/functions-window/src/nth_value.rs:
##########
@@ -525,27 +527,26 @@ 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 => last_set_bit(valid),
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)
+ valid.set_indices().nth(index)
}
Ordering::Less => {
let reverse_index = (-self.n) as usize;
- let valid_indices_len = nulls.len() -
nulls.null_count();
+ let valid_indices_len = valid.count_set_bits();
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(valid_indices_len -
reverse_index)
Review Comment:
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:
```rust
/// 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:
```rust
if valid_indices_len == valid.len() {
return Some(valid.len() - reverse_index);
}
```
<details><summary>Bench cases I added, with the updated
<code>criterion_group!</code></summary>
```rust
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);
```
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
