mbutrovich commented on code in PR #26453:
URL: https://github.com/apache/datafusion/pull/26453#discussion_r4317095151
##########
datafusion/functions-nested/src/sort.rs:
##########
@@ -338,28 +339,29 @@ fn sort_rows_with_nulls<N: ArrowNativeTypeOp, O:
OffsetSizeTrait>(
continue;
}
- let null_count = src_nulls.slice(start, row_len).null_count();
- let valid_count = row_len - null_count;
-
- // Compact valid values directly into the target region of the output
- // buffer: after nulls (if nulls_first) or at the start (if
nulls_last).
- let valid_offset = if nulls_first { null_count } else { 0 };
- let mut write_pos = out_start + valid_offset;
- for (i, &value) in src_values[start..end].iter().enumerate() {
- if src_nulls.is_valid(start + i) {
+ // Compact valid values to the start of the row's output region,
+ // counting them on the way, and sort them there.
+ let validity_bits =
+ BitIterator::new(src_nulls.validity(), src_nulls.offset() + start,
row_len);
+ let mut write_pos = out_start;
+ for (&value, is_valid) in
src_values[start..end].iter().zip(validity_bits) {
+ if is_valid {
out_values[write_pos] = value;
write_pos += 1;
}
}
+ let valid_count = write_pos - out_start;
+ let null_count = row_len - valid_count;
Review Comment:
What do you think about making this loop branchless? It currently branches
on every validity bit. With the benchmark's random NULLs (10% of elements),
that branch is hard to predict. If we write every value and advance `write_pos`
only past valid ones, the branch goes away.
Here is the loop body for `sort_rows_with_nulls::<i32, i32>` at this commit
and with the suggestion below, built with the workspace release profile (LTO,
one codegen unit).
aarch64, this commit:
```asm
lsr x13, x10, #3
ldrb w13, [x9, x13] ; load the validity byte
and x14, x10, #0x7
lsr w13, w13, w14
tbz w13, #0, LBB239_50 ; branch on the validity bit
cmp x27, x20 ; bounds check on out_values
b.hs LBB239_111
ldr w13, [x11]
ldr x14, [sp, #120]
str w13, [x14, x27, lsl #2]
add x27, x27, #1
b LBB239_50
```
aarch64, suggestion:
```asm
cmp x21, x20 ; bounds check on out_values
b.hs LBB239_113
ldr w13, [x12], #4
lsr x14, x10, #3
ldrb w14, [x9, x14] ; load the validity byte
and x15, x10, #0x7
lsr x14, x14, x15
and x14, x14, #0x1 ; validity bit as 0 or 1
add x10, x10, #1
ldr x15, [sp, #136]
str w13, [x15, x21, lsl #2]
add x21, x21, x14 ; write_pos += is_valid
```
On x86_64 (`--target x86_64-apple-darwin`), the branch at this commit is `bt
edi, r8d` followed by `jae LBB1005_49`. With the suggestion it becomes `bt edi,
r8d` followed by `adc rbx, 0`.
I ran the `int32_null_elements` cases on an Apple M5 Max, comparing this
commit with the suggestion back to back:
| case | this PR | suggestion | change |
|---|---|---|---|
| `int32_null_elements/5` | 257.1 us | 253.1 us | -1.8% |
| `int32_null_elements_desc/5` | 264.4 us | 258.8 us | -2.2% |
| `int32_null_elements_nulls_last/5` | 237.9 us | 234.7 us | -0.3% |
| `int32_null_elements/20` | 796.4 us | 711.9 us | -12.5% |
| `int32_null_elements_desc/20` | 834.3 us | 739.7 us | -13.0% |
| `int32_null_elements_nulls_last/20` | 755.3 us | 659.8 us | -12.2% |
| `int32_null_elements/100` | 3.762 ms | 3.459 ms | -8.0% |
| `int32_null_elements_desc/100` | 3.800 ms | 3.407 ms | -10.3% |
| `int32_null_elements_nulls_last/100` | 3.674 ms | 3.247 ms | -11.6% |
| `int32_null_elements/1000` | 44.60 ms | 40.34 ms | -9.6% |
| `int32_null_elements_desc/1000` | 45.93 ms | 40.65 ms | -11.5% |
| `int32_null_elements_nulls_last/1000` | 44.39 ms | 40.99 ms | -7.7% |
I haven't measured x86_64. Could you run it on your EPYC machine?
The loop can leave the last NULL's value in the slot right after the valid
values. The suggestion resets that slot, so NULL slots stay zeroed the way the
`NULLS FIRST` branch keeps them.
```suggestion
// Write every value but advance only past valid ones, so the loop
// doesn't branch on the validity bit.
let mut write_pos = out_start;
for (&value, is_valid) in
src_values[start..end].iter().zip(validity_bits) {
out_values[write_pos] = value;
write_pos += is_valid as usize;
}
let valid_count = write_pos - out_start;
let null_count = row_len - valid_count;
// The slot after the valid values can hold a NULL's value. Reset it
// like the other NULL slots.
if null_count > 0 {
out_values[write_pos] = N::default();
}
```
--
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]
