Repository navigation
feat(arrow-row): add MSD radix sort kernel for row-encoded keys - #9683
mbutrovich wants to merge 42 commits into
Conversation
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
|
I updated numbers in the PR description and full table in the comment after it to reflect current PR status. My attempt to bring it to DataFusion has so far been disappointing: apache/datafusion#21525 I think we need a deeper redesign of that stream rather than just replacing the existing logic. apache/datafusion#21525 (comment) |
|
@Dandandan I also made a couple of API additions to |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark lexsort env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6 |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark lexsort resources:
arch: amd64 |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
resources:
arch: amd64
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark lexsort env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64 |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
resources:
arch: amd64
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-1ev1lxub-b60917df-zn9x Instance: c4-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-1afkvotn-5d03c368-fcjq Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-1ev1lxub-3df3c234-xmx6 Instance: c4-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (d3824b8) to cd91f89 (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-mmrqzil7-d086d2ea-4wb4 Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark lexsort |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark lexsort env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6 |
|
run benchmark lexsort resources:
arch: amd64 |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
resources:
arch: amd64
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark lexsort env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64 |
|
🤖 Benchmark starting (GKE) | trigger Target: lexsort Sharding factor: 1 (1 workers). Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
resources:
arch: amd64
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-1kxm755f-e4e6be63-5zv5 Instance: c4-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-lxiegy2m-6da527fb-j6lg Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-lxiegy2m-b02d82b1-wtxz Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: lexsort Comparing radix_row_sort (03ee81c) to 4e04fea (merge-base) diff Run configurationrun benchmark lexsort
env:
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64
shards: 1
DetailsPer-runner informationlexsort — shard 1/1Node: gk3-benchmark-cluster-nap-gag2u7qv-61f675e5-xbbk Instance: c4-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench lexsortCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
I added An unlimited sort passes a zero-sized The benchmark bot at 03ee81c against d3824b8, which has no With a |
Which issue does this PR close?
N/A.
Rationale for this change
Sorting on several columns is a common step in table rewrites, sort-merge joins, and
ORDER BYwith a composite key. Todaylexsort_to_indicescalls a dynamically dispatched comparator per column on every comparison. Row-encoded keys compare as bytes, soRowscan instead be sorted by MSD radix sort, which reads a fixed-size key from each row once per pass. DuckDB's sort redesign takes the same approach.Why we want it
On the benchmark bot, multi-column keys sort 2x to 6x faster than
lexsort_to_indiceson most cases, including the cost of encoding toRows. It also beatslexsort_rows, which sorts the sameRowsby comparison, in every case.Engines built on arrow-rs could use it in sort-heavy paths. Two examples are a Comet or DataFusion Iceberg table rewrite with a sort order and a sort-merge join on a composite key. I haven't measured either.
Trade-offs
lexsort_to_indiceson a single nullable column (0.36x to 0.81x), a single string column (0.64x to 0.86x), and keys with a list column such as[i32, str_list(4)](0.20x to 0.75x), because the row format encodes every column of every row.limitit stops ordering rows early, but it still encodes and makes one pass over every row. On[i64(0..1000), i32]at 32768 rows,lexsort_to_indicesis faster atlimit100 (radix sort at 0.28x to 0.36x) and radix sort is faster fromlimit1000 (1.28x to 1.68x). At 1M rows,lexsort_to_indicesis faster atlimit100 and 1000 (0.09x to 0.18x).arrow_row::radixis new public API, so the module is permanent once released.The module docs list when to use each sort.
What changes are included in this PR?
arrow_row::radix::radix_sort_to_indices, a radix sort ofRowswith an optionallimit. The module docs say when to use it and when to preferlexsort_to_indices.lexsort_radixtoarrow/benches/lexsort.rsnext tolexsort_to_indicesandlexsort_rows. It encodes the columns and sorts them.encodeto the same benchmark. It runs only the encoding, so the difference between it andlexsort_radixis the sort.lexsort_to_indicesandlexsort_radixcases with alimitof 100, 1000, and n/10 on three keys at 32768 and 1M rows.Benchmarks
These numbers are from the benchmark bot at 03ee81c.
lexsort_radixincludes encoding the rows. Radix sort is faster thanlexsort_rows, which sorts the sameRowsby comparison, in every case on both CPUs. Againstlexsort_to_indices, it is faster in both the default and aligned runs in 43 of 64 cases on the Xeon and 50 on Neoverse-V2, and slower in 19 and 14. It loses on a single nullable or string column, such as[i32_opt](0.49x to 0.59x on the Xeon), and on keys with a list column, such as[i32, str_list(4)](0.20x to 0.29x). At 1M rows, the two integer keys are 2.22x to 3.35x faster on the Xeon and 4.26x to 5.33x on Neoverse-V2.The chart shows
[i64(0..1000), i32]with alimit. The other two keys in the benchmark show the same pattern, and the tables below list all three. Atlimitn/10, radix sort is 2.89x to 4.60x faster thanlexsort_to_indicesacross the three keys and both CPUs. The bot prints milliseconds to one decimal place, so the 1M-row ratios are coarse.The 76 cases already on
mainare within 0.93x to 1.13x of the merge base in all four runs, with medians of 0.98x to 1.00x.All cases
Speedup is the baseline's time divided by
lexsort_radix's time. The aligned run sets-C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6on both sides. The bot prints milliseconds to one decimal place.Xeon Platinum 8581C:
lexsort_to_indiceslexsort_rows[decimal128(18,2), i32][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str(16)][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str_opt(50)][dict(100,str_opt(50)), dict(100,str_opt(50))][f64][f64_opt][i32(0..10), i32(0..10), i32(0..10)][i32(1..204K), i64(0..n/10)][i32, i32_list, str(16)][i32, i32_opt][i32, str(16)][i32, str_list(4)][i32, str_list_opt(4)][i32, str_opt(16)][i32][i32_list_opt, i32_opt][i32_opt, dict(100,str_opt(50))][i32_opt, i32_list][i32_opt, i32_list_opt, str_opt(50)][i32_opt, i32_list_opt][i32_opt][i64(0..1000), i32][i64][i64_opt][str(16)][str_list(4), i32][str_list_opt(4), i32][str_opt(16), str(16), str_opt(16), str_opt(16), str_opt(16)][str_opt(16), str(16)][str_opt(16), str_opt(50), str(16)][str_opt(16)][decimal128(18,2), i32][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str(16)][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str_opt(50)][dict(100,str_opt(50)), dict(100,str_opt(50))][f64][f64_opt][i32(0..10), i32(0..10), i32(0..10)][i32(1..204K), i64(0..n/10)][i32, i32_list, str(16)][i32, i32_opt][i32, str(16)][i32, str_list(4)][i32, str_list_opt(4)][i32, str_opt(16)][i32][i32_list_opt, i32_opt][i32_opt, dict(100,str_opt(50))][i32_opt, i32_list][i32_opt, i32_list_opt, str_opt(50)][i32_opt, i32_list_opt][i32_opt][i64(0..1000), i32][i64][i64_opt][str(16)][str_list(4), i32][str_list_opt(4), i32][str_opt(16), str(16), str_opt(16), str_opt(16), str_opt(16)][str_opt(16), str(16)][str_opt(16), str_opt(50), str(16)][str_opt(16)][i32(0..10), i32(0..10), i32(0..10)][i32(1..204K), i64(0..n/10)]Neoverse-V2:
lexsort_to_indiceslexsort_rows[decimal128(18,2), i32][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str(16)][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str_opt(50)][dict(100,str_opt(50)), dict(100,str_opt(50))][f64][f64_opt][i32(0..10), i32(0..10), i32(0..10)][i32(1..204K), i64(0..n/10)][i32, i32_list, str(16)][i32, i32_opt][i32, str(16)][i32, str_list(4)][i32, str_list_opt(4)][i32, str_opt(16)][i32][i32_list_opt, i32_opt][i32_opt, dict(100,str_opt(50))][i32_opt, i32_list][i32_opt, i32_list_opt, str_opt(50)][i32_opt, i32_list_opt][i32_opt][i64(0..1000), i32][i64][i64_opt][str(16)][str_list(4), i32][str_list_opt(4), i32][str_opt(16), str(16), str_opt(16), str_opt(16), str_opt(16)][str_opt(16), str(16)][str_opt(16), str_opt(50), str(16)][str_opt(16)][decimal128(18,2), i32][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str(16)][dict(100,str_opt(50)), dict(100,str_opt(50)), dict(100,str_opt(50)), str_opt(50)][dict(100,str_opt(50)), dict(100,str_opt(50))][f64][f64_opt][i32(0..10), i32(0..10), i32(0..10)][i32(1..204K), i64(0..n/10)][i32, i32_list, str(16)][i32, i32_opt][i32, str(16)][i32, str_list(4)][i32, str_list_opt(4)][i32, str_opt(16)][i32][i32_list_opt, i32_opt][i32_opt, dict(100,str_opt(50))][i32_opt, i32_list][i32_opt, i32_list_opt, str_opt(50)][i32_opt, i32_list_opt][i32_opt][i64(0..1000), i32][i64][i64_opt][str(16)][str_list(4), i32][str_list_opt(4), i32][str_opt(16), str(16), str_opt(16), str_opt(16), str_opt(16)][str_opt(16), str(16)][str_opt(16), str_opt(50), str(16)][str_opt(16)][i32(0..10), i32(0..10), i32(0..10)][i32(1..204K), i64(0..n/10)]Cases with a
limit, on Xeon Platinum 8581C:lexsort_to_indices[i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i64(0..1000), i32][i64(0..1000), i32][i64(0..1000), i32][str_opt(16), str(16)][str_opt(16), str(16)][str_opt(16), str(16)][i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i64(0..1000), i32][i64(0..1000), i32][i64(0..1000), i32][str_opt(16), str(16)][str_opt(16), str(16)][str_opt(16), str(16)]Cases with a
limit, on Neoverse-V2:lexsort_to_indices[i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i64(0..1000), i32][i64(0..1000), i32][i64(0..1000), i32][str_opt(16), str(16)][str_opt(16), str(16)][str_opt(16), str(16)][i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i32(1..204K), i64(0..n/10)][i64(0..1000), i32][i64(0..1000), i32][i64(0..1000), i32][str_opt(16), str(16)][str_opt(16), str(16)][str_opt(16), str(16)]Are these changes tested?
Unit tests are in
arrow-row/src/radix.rs, including a fuzz test over random column types and sort options that checks the output is a sorted permutation of the rows, and tests for rows whose keys tie.Are there any user-facing changes?
Yes. The
arrow_row::radixmodule andradix_sort_to_indicesare new public API.AI Disclosure
I wrote this PR with help from Claude but understand and support the core approach, implementation, and test coverage.