Repository navigation
Perf: Improve dictionary slice concatenation - #12008
Conversation
github-actions
Bot
added
arrow-kernels
|
run benchmark concatenate_kernels |
|
🤖 Benchmark starting (GKE) | trigger Target: concatenate_kernels Sharding factor: 1 (1 workers). Comparing codex/dictionary-concat-slices (90cbf31) to bf0cde3 (merge-base) diff Run configurationrun benchmark concatenate_kernels
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark failed or incomplete (GKE) | trigger Target: concatenate_kernels Comparing codex/dictionary-concat-slices (90cbf31) to bf0cde3 (merge-base) diff Run configurationrun benchmark concatenate_kernels
shards: 1
Errors / missing shardsPer-runner informationconcatenate_kernels — shard 1/1Node: gk3-benchmark-cluster-nap-bch756b4-46b2a638-fvcz Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench concatenate_kernelsCPU Details (lscpu)Resource UsageNo completed measurement resource samples. File an issue against this benchmark runner |
|
The controlled DataFusion end-to-end run is complete: results. Using DataFusion #26816 with every Arrow crate pinned to either The queries expected to concatenate sliced dictionary group keys improved most:
No spill occurred on either side. |
|
should have marked this a draft |
|
run benchmark concatenate_kernel |
|
🤖 Benchmark starting (GKE) | trigger Target: concatenate_kernel Sharding factor: 1 (1 workers). Comparing codex/dictionary-concat-slices (90cbf31) to bf0cde3 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: dict
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: concatenate_kernel Comparing codex/dictionary-concat-slices (90cbf31) to bf0cde3 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: dict
shards: 1
DetailsPer-runner informationconcatenate_kernel — shard 1/1Node: gk3-benchmark-cluster-nap-1h5f9ztd-949bd8ae-9fwn Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench concatenate_kernelCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
90cbf31 to
32f0197
Compare
|
@CuteChuanChuan could you take a look |
Rich-T-kid
changed the title
|
I had ai run some benchmarks with this pr as the tip in datafusion and the results look good #12008 (comment). I think it may be worth adding targeted benchmarks here just so we can choose a correct threshold as 16x may be too conservative. |
|
run benchmark concatenate_kernel env: |
|
🤖 Benchmark starting (GKE) | trigger Target: concatenate_kernel Sharding factor: 1 (1 workers). Comparing codex/dictionary-concat-slices (6149e4e) to b65bab9 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat .*dict
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: concatenate_kernel Comparing codex/dictionary-concat-slices (6149e4e) to b65bab9 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat .*dict
shards: 1
DetailsPer-runner informationconcatenate_kernel — shard 1/1Node: gk3-benchmark-cluster-nap-xeijzk4q-715bf30c-v8g2 Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench concatenate_kernelCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
# Which issue does this PR close? Part of apache#11885. Adds benchmark coverage for apache#11901 and apache#12008, and follows up on the Boolean filter coverage requested in [apache#11918](apache#11918 (comment)). This consolidates and supersedes apache#11980 and apache#12041. # Why is this change needed? The existing suites do not cover Boolean arrays as filter data or the cost of operating on small key ranges backed by oversized dictionaries. They also do not exercise the 16-values-per-selected-key boundary between dense and sparse dictionary mappings. # What changes are included in this PR? - Adds nullable and non-null Boolean filter cases across the existing selectivities. - Adds prepared `RowConverter` cases for rotating 8,192-key ranges backed by unique string dictionaries with 1M and 4M values. - Measures the full-value encoding cost from apache#11901 without including `RowConverter` construction. - Adds concat cases for 12 distinct dictionaries sliced to 682 keys, with 12, 16, and 24 values per selected key. - Straddles apache#12008's 16-values-per-selected-key cutoff so we can assess the threshold. # Are these changes tested? Built and ran each filtered benchmark locally. Benchmark bot runs: - [`filter_kernels`](apache#12054 (comment)) - [`row_format`](apache#12054 (comment)) - [`concatenate_kernel`](apache#12054 (comment)) # Are there any significant user-facing changes? No. This only adds benchmark coverage.
|
run benchmark concatenate_kernel |
|
🤖 Benchmark starting (GKE) | trigger Target: concatenate_kernel Sharding factor: 1 (1 workers). Comparing codex/dictionary-concat-slices (5242170) to 258b0d2 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat .*dict
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark concatenate_kernel |
|
🤖 Benchmark starting (GKE) | trigger Target: concatenate_kernel Sharding factor: 1 (1 workers). Comparing codex/dictionary-concat-slices (5242170) to 258b0d2 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat .*dict
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark concatenate_kernel |
|
🤖 Benchmark starting (GKE) | trigger Target: concatenate_kernel Sharding factor: 1 (1 workers). Comparing codex/dictionary-concat-slices (48da00e) to 258b0d2 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat str_dict_slice
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: concatenate_kernel Comparing codex/dictionary-concat-slices (48da00e) to 258b0d2 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat str_dict_slice
shards: 1
DetailsPer-runner informationconcatenate_kernel — shard 1/1Node: gk3-benchmark-cluster-nap-1mwxpx69-f006e89d-qbnz Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench concatenate_kernelCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark concatenate_kernel env: |
|
🤖 Benchmark starting (GKE) | trigger Target: concatenate_kernel Sharding factor: 1 (1 workers). Comparing codex/dictionary-concat-slices (31c0d5c) to 258b0d2 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat str_dict_slice
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: concatenate_kernel Comparing codex/dictionary-concat-slices (31c0d5c) to 258b0d2 (merge-base) diff Run configurationrun benchmark concatenate_kernel
env:
BENCH_FILTER: ^concat str_dict_slice
shards: 1
DetailsPer-runner informationconcatenate_kernel — shard 1/1Node: gk3-benchmark-cluster-nap-rw3tiycv-83893b76-mq2f Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench concatenate_kernelCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
I ran the partitioned ClickBench suite twice with Parquet RLE-to-dictionary enabled: run 1, run 2. Both sides use DataFusion #26816 at Configuration:
These are the runner's minimum times, which it uses for the comparison verdict. Total minimum time fell by 22.3% and 22.5% in the two runs. The first run classified 9 queries faster and 6 slower; the repeat classified 12 faster and none slower. The smaller changes were noisy, but the gains on these four queries repeated. One caveat: the first changed-side Q33 run had a single spill-heavy 21.6 s iteration, so its five-run mean was worse despite the 1.95x minimum-time result. That did not recur: Q33 was 4.36-4.70 s on the changed side in the repeat versus 7.89-8.16 s at baseline. The repeat still spilled 400 MiB and reported higher pool peaks for Q33/Q34, so the speedup is clear but the memory behavior is worth keeping an eye on. |
|
okay this is good to go! @jayzhan211 @alamb I think you have the most context with the dictionary/datafusion work thats been going on so you may be interested. the comment above gives a decent breakdown on the downstream affects |
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
@Rich-T-kid thanks for pushing this through: the #11885 repro with 1M values drops from 502 µs to 164 µs here. The inline comments have two slowdowns to fix around the sparse threshold, and these two points don't anchor on the diff:
In arrow-select/src/coalesce/dictionary.rs on main, merge_sources still builds a source.len() mask and mapping per source: through BatchCoalescer, 12 x 682 rows takes 82 µs with 1,000 values and 470 µs with 1M, against 174 µs for this PR's concat on its old base. The repartition query in #11885 takes that path, so could merge_dictionary_values_with_masks share the referenced-key path? A follow-up PR is fine, because it's a different call path with its own tests, but the DataFusion gains probably won't show until it lands.
In arrow/benches/concatenate_kernel.rs around line 203, values_per_key of 12, 16 and 24 never reach the sparse path, which needs more than 87,296 values at 682 rows, so the 8.5-10% in the description comes from the dense path. Could you add ratios on both sides of the threshold, plus a scattered-key variant, and update the description to cite them?
- for values_per_key in [12, 16, 24] {
+ for values_per_key in [12, 16, 24, 128, 256, 400, 1024] {| let selected_rows = key_mask | ||
| .map(BooleanBuffer::count_set_bits) | ||
| .unwrap_or(keys.len()); | ||
| let use_sparse_mapping = | ||
| values.len() > selected_rows.saturating_mul(MAX_DENSE_MAPPING_VALUES_PER_ROW); |
There was a problem hiding this comment.
The sparse check counts selected rows, but the remap does one hash lookup per key, so mostly-null and repeated inputs get slower: concat of 2 x 8,192 rows with 100k values and 500 valid goes from 39 µs to 57 µs, and 2 x 1M rows with 1M values and 7,000 valid goes from 2.6 ms to 28.9 ms. Could the check count key lookups instead? With the change below those two take 38 µs and 2.4 ms, and the 1M-value repro stays at 170 µs.
--- a/arrow-select/src/dictionary.rs
+++ b/arrow-select/src/dictionary.rs
masks: Option<&[BooleanBuffer]>,
+ key_lookups: Option<&[usize]>,
) -> Result<MergedDictionaries<K>, ArrowError> {
- let selected_rows = key_mask
- .map(BooleanBuffer::count_set_bits)
- .unwrap_or(keys.len());
+ // The caller remaps once per lookup, which can far exceed the selected rows.
+ let lookups = key_lookups.map_or(keys.len(), |lookups| lookups[idx]);
let use_sparse_mapping =
- values.len() > selected_rows.saturating_mul(MAX_DENSE_MAPPING_VALUES_PER_ROW);
+ values.len() > lookups.saturating_mul(MAX_DENSE_MAPPING_VALUES_PER_ROW);
--- a/arrow-select/src/interleave.rs
+++ b/arrow-select/src/interleave.rs
- let masks: Vec<_> = dictionaries
+ let (masks, key_lookups): (Vec<_>, Vec<_>) = dictionaries
+ let mut lookups = 0;
for (_, key_idx) in indices.iter().filter(|(a, _)| *a == a_idx) {
key_mask.set_bit(*key_idx, true);
+ lookups += 1;
}
- key_mask.finish()
+ (key_mask.finish(), lookups)
})
- .collect();
+ .unzip();
- let merged = merge_dictionary_values(&dictionaries, Some(&masks))?;
+ let merged = merge_dictionary_values(&dictionaries, Some(&masks), Some(&key_lookups))?;
--- a/arrow-select/src/concat.rs
+++ b/arrow-select/src/concat.rs
- let merged = merge_dictionary_values(&dictionaries, None)?;
+ let merged = merge_dictionary_values(&dictionaries, None, None)?;Timings per input shape, and a test that separates lookups from selected rows
Median of 3 rounds, each the best of at least 9 runs, on an M4 Pro under load, so ignore differences under 10%. Null-slot keys are in range.
| Input | Base | This PR | With the change |
|---|---|---|---|
| concat 12 x 682 rows, 1M values | 502 µs | 164 µs | 170 µs |
| concat 2 x 8,192 rows, 100k values, 500 valid | 39 µs | 57 µs | 38 µs |
| concat 2 x 1M rows, 1M values, 7,000 valid | 2.6 ms | 28.9 ms | 2.4 ms |
| concat 2 x 1M rows, 1M values, 10 valid | 1.7 ms | 3.7 ms | 1.5 ms |
| interleave 1M indices over 100 rows, 2 x 1M values | 2.7 ms | 4.0 ms | 2.8 ms |
| interleave 8,192 indices over 100 rows, 2 x 1M values | 104 µs | 58 µs | 106 µs |
The last row is what the change gives up: heavily repeated rows go back to the dense path.
#[test]
fn test_sparse_gate_counts_key_lookups() {
let values: ArrayRef = Arc::new(StringArray::from_iter_values(
(0..8 * MAX_DENSE_MAPPING_VALUES_PER_ROW).map(|idx| idx.to_string()),
));
let is_sparse = |dictionary: &DictionaryArray<Int32Type>,
masks: Option<&[BooleanBuffer]>,
lookups: Option<&[usize]>| {
let merged = merge_dictionary_values(&[dictionary], masks, lookups).unwrap();
matches!(merged.key_mappings[0], KeyMapping::Sparse(_))
};
// 1 valid row of 10: over the gate for 1 row, but concat looks up all 10 rows.
let keys = Int32Array::from_iter((0..10).map(|idx| (idx == 0).then_some(999)));
let mostly_null = DictionaryArray::new(keys, Arc::clone(&values));
assert!(!is_sparse(&mostly_null, None, None));
// 7 lookups stay under the gate, keeping the sparse path.
let keys = Int32Array::from_iter_values([999; 7]);
let short = DictionaryArray::new(keys, Arc::clone(&values));
assert!(is_sparse(&short, None, None));
// Interleave selects 1 row but may look it up repeatedly.
let mask = BooleanBuffer::collect_bool(7, |idx| idx == 0);
assert!(!is_sparse(
&short,
Some(std::slice::from_ref(&mask)),
Some(&[10])
));
assert!(is_sparse(&short, Some(&[mask]), Some(&[7])));
}There was a problem hiding this comment.
Changed the gate to use the actual remap count. Concat defaults to the key-array length, interleave passes per-input lookup counts, and the new gate test covers the mostly-null and repeated-selection cases. This is in bdd44b6.
| } | ||
|
|
||
| // Sparse hashing pays off only when retained values greatly outnumber selected rows. | ||
| const MAX_DENSE_MAPPING_VALUES_PER_ROW: usize = 128; |
There was a problem hiding this comment.
With scattered keys, concat is 18-23% slower just above 128:1: 12 x 682 rows at 129:1 goes from 216 µs to 256 µs, and 2 x 6,900 rows goes from 412 µs to 487 µs. The slice benchmark and the #11885 repro use contiguous keys, where sparse wins from 129:1, so scattered keys cross over later, near 256:1. Could we use 256? That gives up as much as 23% on contiguous keys between 129:1 and 256:1, and the DataFusion shape in #11885 is about 880:1, so it isn't affected.
| const MAX_DENSE_MAPPING_VALUES_PER_ROW: usize = 128; | |
| const MAX_DENSE_MAPPING_VALUES_PER_ROW: usize = 256; |
Head / base ratio by values per selected row and key order
Median of 5 rounds of best-of-9 on an M4 Pro, spread mostly under 12%. All keys valid. Each cell is contiguous / scattered.
| Values per row | 12 x 682 rows | 3 x 2,730 rows | 2 x 6,900 rows |
|---|---|---|---|
| 128 (dense) | 0.97 / 1.02 | 0.96 / 0.98 | 0.95 / 0.97 |
| 129 | 0.97 / 1.19 | 0.92 / 1.23 | 0.93 / 1.18 |
| 200 | 0.85 / 1.01 | 0.82 / 1.10 | 0.78 / 1.09 |
| 256 | 0.77 / 0.95 | 0.79 / 1.05 | 0.80 / 0.96 |
| 400 | 0.64 / 0.76 | 0.57 / 0.75 | 0.63 / 0.84 |
| 1024 | 0.40 / 0.53 | 0.35 / 0.50 | 0.44 / 0.56 |
There was a problem hiding this comment.
Raised the cutoff to 256. I also added contiguous and scattered benchmark cases at 128, 256, 400, and 1024, and updated the description with the crossover numbers. bdd44b6
| /// Provides `key_mappings[`array_idx`][`old_key`] -> new_key` | ||
| pub key_mappings: Vec<Vec<K::Native>>, | ||
| /// Maps each input dictionary key to its merged key. | ||
| pub key_mappings: Vec<KeyMapping<K>>, |
There was a problem hiding this comment.
This needs more than a text-conflict rebase: since #11743, coalesce/dictionary.rs indexes mapping[key.as_usize()] and merge_dictionary_values_with_masks calls get_masked_values, and this PR changes the mapping's type and deletes that helper. Could you rebase and re-run the ClickBench comparison? Those numbers pinned 258b0d2031, before BatchCoalescer stopped calling concat, so the Q16-Q34 gains may not carry over.
There was a problem hiding this comment.
Rebased onto current main. The coalescer now reads through KeyMapping::get, and merge_dictionary_values_with_masks uses the selected-index helper that replaced get_masked_values. I kept the separate BatchCoalescer optimization out of this PR and made that scope explicit in the description. bdd44b6
| fn assert_mapping(mapping: &KeyMapping<Int32Type>, expected: &[i32]) { | ||
| let actual = (0..expected.len()) | ||
| .map(|idx| mapping.get(idx).unwrap_or_default()) | ||
| .collect::<Vec<_>>(); | ||
| assert_eq!(actual, expected); | ||
| } |
There was a problem hiding this comment.
assert_mapping no longer checks length or variant and reads a missing entry as 0, so the &[] call at line 595 can't fail and these tests would still pass if an input moved to the sparse path. Comparing the whole dense Vec keeps the old assertions:
| fn assert_mapping(mapping: &KeyMapping<Int32Type>, expected: &[i32]) { | |
| let actual = (0..expected.len()) | |
| .map(|idx| mapping.get(idx).unwrap_or_default()) | |
| .collect::<Vec<_>>(); | |
| assert_eq!(actual, expected); | |
| } | |
| fn assert_mapping(mapping: &KeyMapping<Int32Type>, expected: &[i32]) { | |
| let KeyMapping::Dense(actual) = mapping else { | |
| panic!("expected a dense mapping"); | |
| }; | |
| assert_eq!(actual, expected); | |
| } |
There was a problem hiding this comment.
Fixed this as suggested: assert_mapping now requires the Dense variant and compares the whole Vec, so a missing or truncated entry cannot read back as zero. bdd44b6
| } | ||
|
|
||
| #[test] | ||
| fn test_concat_oversized_dictionary_slices() { |
There was a problem hiding this comment.
This test and test_interleave_oversized_dictionaries hardcode 1,000 values and never assert the sparse path ran, so raising the threshold can move them to the dense path with no failure. Could they derive sizes from MAX_DENSE_MAPPING_VALUES_PER_ROW, and add sparse cases for primitive values, null dictionary values, and null keys through both concat and interleave? A randomized base-vs-head run over those shapes found no differences in 60,000 operations, so this is only about keeping them covered.
There was a problem hiding this comment.
Updated both tests to derive their sizes from the cutoff and assert that the sparse mapping ran. They now cover primitive values, null dictionary values, and null keys through both concat and interleave. bdd44b6
Sorting selected keys both enables deduplication and preserves the dense path's dictionary-value order. Add a focused test that keeps this ordering invariant explicit.
The sparse path only wins once retained values substantially outnumber selected rows. Raise the strict cutoff to 128x and keep dense iteration monomorphic so ordinary dictionaries avoid mapping dispatch overhead.
The sparse mapping path must preserve the existing overflow contract. Build 257 UInt8 dictionaries with one unique referenced value each and verify the merge returns DictionaryKeyOverflowError.
f61c8a2 to
bdd44b6
Compare
github-actions
Bot
added
the
arrow-core
|
Rebased this onto current main and left BatchCoalescer as a separate follow-up, as suggested. The concat benchmark now covers 128, 256, 400, and 1024 with both contiguous and scattered keys, and the description cites the crossover data. Commit bdd44b6. |
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
@Rich-T-kid thanks for the quick update: the gate and threshold changes hold up. At 257 values per key the slice benchmark is now 0.78x of main with contiguous keys and 0.84x with scattered keys, and 23,505 randomized concat and interleave results over sliced inputs, null keys and repeated rows match main. Could you open an issue for the BatchCoalescer follow-up and link it from the description? merge_sources still builds a mask and mapping per source sized to the whole dictionary, and the repartition example in #11885 goes through that path, so it would lose its tracking when this PR closes the issue.
Which issue does this PR close?
Closes #11885
Why is this change needed?
Slicing a dictionary array only slices its keys; it still retains the full values array. When concat or interleave merges short slices from large dictionaries, the current code builds a mask and key mapping for every dictionary value, including values that no selected row references. The work therefore scales with the full dictionary instead of the number of output rows.
What changes are included in this PR?
When a dictionary has more than 256 values per key lookup, collect only the referenced keys and store the old-to-new key mappings sparsely. The gate uses the number of remaps performed by concat or interleave, rather than the number of distinct selected rows, so mostly-null and repeatedly selected inputs stay on the dense path when hashing would be slower.
Let
rbe the number of selected rows,lthe number of key lookups,uthe number of unique referenced keys, andvthe number of retained dictionary values. The sparse path collects keys inO(r), sorts them inO(r log r), deduplicates them inO(r), builds the hash mapping inO(u)expected time, and remaps the output inO(l)expected time.This replaces the dense path's
O(v)mask and mapping work. Sorting also makes repeated keys adjacent and preserves the ascending dictionary-key order produced by the dense mask path, so both paths build the merged dictionary in the same order. Both concat and interleave use this shared dictionary merge path. BatchCoalescer has a separate dictionary merge path and is intentionally left for a follow-up.Are these changes tested?
Tests cover concat and interleave of short slices backed by large dictionaries, the key-lookup gate, sparse primitive dictionaries, null dictionary values, null keys,
UInt8key overflow, and the existing dense path.The benchmark now covers ratios of 12, 16, 24, 128, 256, 400, and 1024 values per key with both contiguous and scattered keys. Review measurements placed the scattered-key crossover near 256:1: at 400:1 the measured head/base ratios were 0.57-0.84, and at 1024:1 they were 0.35-0.56. The earlier hosted run improved the 12:1, 16:1, and 24:1 dense-path cases by 10.0%, 8.5%, and 8.7%. The benchmark harness was introduced in #12054.
Are there any significant user-facing changes?
No. This only changes how dictionary values are merged internally.