mbutrovich commented on PR #9683: URL: https://github.com/apache/arrow-rs/pull/9683#issuecomment-6776870310
I picked this back up and reworked the kernel. Instead of radix sorting `u32` indices and reading every row through `Rows` on each pass, it now skips the bytes that every row shares, copies the next 8 bytes of each row into a `u128` next to the row's index, radix sorts those, and compares the rest of a row only to break ties. It no longer uses `unsafe`. The changes are [a12caca...4f9d3e9](https://github.com/apache/arrow-rs/compare/a12caca2534723779ec8261838f842a7beea292d...4f9d3e901d62220383cccf9420de91db280ed9e1). a12caca has the same algorithm as the PR before this round of work (d9018f5), after the API changes to return `Result<UInt32Array, ArrowError>` and drop the `Row` helpers. The benchmark bot compared the old and new kernels in the same run on a Xeon Platinum 8581C, in a default build ([results](https://github.com/apache/arrow-rs/pull/9683#issuecomment-6774927982)) and with function and block alignment set to 64 bytes on both sides ([results](https://github.com/apache/arrow-rs/pull/9683#issuecomment-6776042189)). Each time includes encoding the rows. Across the 58 `lexsort_radix` cases, the new kernel is 1.21x to 1.22x faster at the median, 1.36x to 1.39x as a geometric mean, and 0.96x to 4.66x per case. The untouched `lexsort_to_indices` and `lexsort_rows` cases are within 0.94x to 1.06x. <img width="1600" height="1764" alt="radix-progress-bot" src="https://github.com/user-attachments/assets/947c206c-3ae0-4a89-8c21-4f7b106760e1"; /> <details><summary>Times</summary> | Case | Rows | a12caca | c0ef5c0 | Speedup (default / aligned) | |---|---:|---:|---:|---:| | `[i64(0..1000), i32]` | 4096 | 263 us | 59 us | 4.49x / 4.66x | | `[decimal128(18,2), i32]` | 32768 | 3.00 ms | 931 us | 3.22x / 3.44x | | `[i32_opt]` | 4096 | 184 us | 59 us | 3.13x / 3.31x | | `[i32]` | 4096 | 116 us | 42 us | 2.73x / 3.04x | | `[i32, i32_opt]` | 4096 | 125 us | 54 us | 2.32x / 2.61x | | `[i32, str(16)]` | 4096 | 143 us | 68 us | 2.10x / 2.31x | | `[i32_opt, dict(100,str_opt(50))]` | 4096 | 184 us | 109 us | 1.69x / 1.61x | | `[dict(100,str_opt(50)), dict(100,str_opt(50))]` | 4096 | 248 us | 191 us | 1.30x / 1.31x | | `[str_opt(16), str(16)]` | 32768 | 1.62 ms | 1.28 ms | 1.27x / 1.30x | | `[i32, str_list(4)]` | 4096 | 790 us | 646 us | 1.22x / 1.20x | | `[i32_opt, i32_list_opt]` | 32768 | 2.40 ms | 2.50 ms | 0.96x / 1.00x | </details> The largest gains come from skipping shared leading bytes (small decimals and small-range integers) and from replacing `memcmp` on rows with integer comparisons when sorting small buckets. Keys dominated by list encoding gain the least. Against `lexsort_to_indices`, measured within each run against the merge base, radix sort went from 24 faster and 31 slower cases on the Xeon to 37 faster and 18 slower, and from 37 and 19 to 44 and 14 on Neoverse-V2, counting cases that are faster or slower in both the default and aligned runs. The earlier counts come from the runs on a12caca above ([Xeon](https://github.com/apache/arrow-rs/pull/9683#issuecomment-6757338133), [Neoverse-V2](https://github.com/apache/arrow-rs/pull/9683#issuecomment-6757323355)), and the later ones from the runs on c0ef5c0, which has the same code as 4f9d3e9 ([Xeon](https://github.com/apache/arrow-rs/pull/9683#issuecomment-6774550808), [Neoverse-V2](https://github.com/apache/arrow-rs/pull/9683#issuecomment-6774535276)). The PR description has the current numbers in full. -- 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]
