fix(arrow-select): merge dictionaries in merge and merge_n instead of panicking on key overflow by ninglinLiu · Pull Request #11925 · apache/arrow-rs

ninglinLiu · 2026-10-26T20:34:34Z

Which issue does this PR close?

Closes merge and merge_n panic on dictionaries whose combined values overflow the key type #11834

Rationale for this change

merge and merge_n build their output with the infallible MutableArrayData::new, which concatenates the input dictionaries and panics with DictionaryKeyOverflowError when the combined values do not fit the key type, even when every input has the same values and a merged dictionary would fit. concat and interleave already merge dictionaries in that case.

What changes are included in this PR?

Both kernels now call MutableArrayData::try_new. On DictionaryKeyOverflowError they fall back to interleave, which merges the dictionaries (or returns an error if the merged dictionary still does not fit). Other errors are returned as before, and the non-dictionary fast path is unchanged.
merge builds the (array, row) pairs from the mask, treating null mask entries as false and reading row 0 for a scalar side, the same way the slice-based path does.
merge_n builds the pairs from the indices. interleave has no null index, so None rows read from an extra one-row null array of the same type.
merge with two scalars delegates to zip, which is handled by fix(arrow-select): merge dictionaries in zip instead of panicking on key overflow #11663, so it is not touched here.

Are these changes tested?

Yes. Two new tests use two Int8-keyed dictionaries with the same 100 values, so concatenation overflows but the merged dictionary fits:

test_merge_dictionaries_whose_combined_values_overflow_the_key_type: array/array, array/scalar and scalar/array, with a null in the mask and a null key in the input.
test_merge_n_dictionaries_whose_combined_values_overflow_the_key_type: Option<usize> indices with None rows and a null key.

Both panic on main and pass with this change. cargo test -p arrow-select and cargo clippy -p arrow-select --all-targets -- -D warnings pass.

Are there any user-facing changes?

merge and merge_n return a merged dictionary instead of panicking. No API changes.

This PR was prepared with AI assistance and reviewed/tested by the author.

… panicking on key overflow

Rich-T-kid

PR makes sense to me. Maybe we should add benchmarks to this as well. Im sure we could avoid the interleave call in the future

ninglinLiu · 2026-10-27T05:08:27Z

Thanks for taking a look, @Rich-T-kid!

On benchmarks: the common path does the same work as before. MutableArrayData::try_new is the fallible form of new (new just unwraps it), so non-dictionary inputs and dictionaries that don't overflow go through exactly the same code. The existing arrow/benches/merge_kernels.rs covers that path. The interleave fallback only runs where the kernels used to panic with DictionaryKeyOverflowError, so there is no previous number to compare it against.

Agreed that the fallback could avoid interleave later, e.g. by merging the dictionaries once and remapping keys directly in the merge loop. If you think it's useful, I can add a dictionary case to merge_kernels.rs in a follow-up so that optimization has a baseline. I'd rather keep this PR focused on the panic.

Rich-T-kid · 2026-10-27T05:46:48Z

Agreed that the fallback could avoid interleave later, e.g. by merging the dictionaries once and remapping keys directly in the merge loop. If you think it's useful, I can add a dictionary case to merge_kernels.rs in a follow-up so that optimization has a baseline. I'd rather keep this PR focused on the panic.

yes this is what I was referring to, we should add a benchmark case where it will trigger this merging. Id recommend including it in this PR as it should be straight forward & minimal.

this isn't a very strong opinion though, if you'd like to do it in a follow up PR thats also fine

…at merges dictionaries

ninglinLiu · 2026-10-27T06:04:29Z

Makes sense, added it here in 9f1e1c9. merge_kernels now has an int8_dictionaries group with two array-vs-array cases over the usual masks: two dictionaries of 50 shared values (the concatenated dictionary still fits Int8 keys, so this is the existing MutableArrayData path) and two of 100 shared values (concatenation overflows, so this goes through the new dictionary-merging fallback).

Locally, for 8192 rows, the fallback is ~2x slower at 50% true (~63µs vs ~121µs) and ~3x at 90% true (~19µs vs ~61µs), which gives a baseline for remapping the keys directly later.

Rich-T-kid · 2026-10-27T14:11:16Z

+    // 2 x 50 values fit in Int8 keys when the dictionaries are concatenated,
+    // 2 x 100 values do not, so merging them has to merge the dictionaries
+    for (description, distinct) in [
+        ("array_vs_array_concatenated_values_fit", 50),


if there are existing benchmarks that cover the case where dictionarys fit within the key space I think we should remove this.

There isn't one at the moment: before this PR, merge_kernels.rs only benchmarked i32, strings and binary, so no merge benchmark used dictionary inputs at all. This case is the only baseline for dictionaries whose concatenated values still fit the key type (the MutableArrayData path), and it lets the overflow case below be compared against it. I'd like to keep it for that reason, but I'm happy to drop it if you'd rather keep only the overflow case.

okay thanks for checking

Rich-T-kid

thanks for the iteration @ninglinLiu

Rich-T-kid · 2026-10-27T15:21:07Z

run benchmark merge_kernels
env:
BENCH_FILTER: "merge_8192_from_int8_dictionaries"

adriangbot · 2026-10-27T15:21:13Z

🤖 Benchmark starting (GKE) | trigger

Target: merge_kernels

Sharding factor: 1 (1 workers).

Comparing fix/merge-dict-overflow (9f1e1c9) to 69be95c (merge-base) diff

Run configuration

run benchmark merge_kernels
env:
  BENCH_FILTER: merge_8192_from_int8_dictionaries
shards: 1

Results will be posted when all workers finish.

File an issue against this benchmark runner

adriangbot · 2026-10-27T15:27:35Z

🤖 Benchmark completed (GKE) | trigger

Target: merge_kernels

Comparing fix/merge-dict-overflow (9f1e1c9) to 69be95c (merge-base) diff

Run configuration

run benchmark merge_kernels
env:
  BENCH_FILTER: merge_8192_from_int8_dictionaries
shards: 1

Details

group                                                                                        changed
-----                                                                                        -------
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/10pct_true          1.00     16.5±0.02µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/1pct_true           1.00      4.0±0.01µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/50pct_nulls         1.00     29.7±0.10µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/50pct_true          1.00     43.3±0.17µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/90pct_true          1.00     16.9±0.02µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/99pct_true          1.00      4.5±0.01µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/all_false           1.00      2.2±0.01µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_fit/all_true            1.00      2.2±0.01µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/10pct_true     1.00     60.6±0.10µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/1pct_true      1.00     61.3±0.13µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/50pct_nulls    1.00     86.9±0.57µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/50pct_true     1.00     99.0±0.49µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/90pct_true     1.00     60.6±0.10µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/99pct_true     1.00     61.4±0.08µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/all_false      1.00     60.2±0.10µs        ? ?/sec
merge_8192_from_int8_dictionaries/array_vs_array_concatenated_values_overflow/all_true       1.00     60.4±0.12µs        ? ?/sec

Per-runner information

merge_kernels — shard 1/1

Node: gk3-benchmark-cluster-nap-gqz8246l-d890b5d1-9gw5

Instance: c4a-highmem-16 (12 vCPU / 65 GiB)

uname:

Linux bench-c6445921439-3928-h85jp 6.12.105+ #1 SMP Sat Sep 19 07:55:25 UTC 2026 aarch64 GNU/Linux

BENCH_COMMAND:

cargo bench --features=arrow,async,test_common,experimental,object_store --bench merge_kernels

CPU Details (lscpu)

Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Resource Usage

base (merge-base)

Metric	Value
Wall time	5.0s
Peak memory	4.9 MiB
Avg memory	831.3 KiB
CPU user	0.0s
CPU sys	0.0s
Peak spill	0 B

branch

Metric	Value
Wall time	160.0s
Peak memory	13.5 MiB
Avg memory	12.5 MiB
CPU user	154.2s
CPU sys	0.0s
Peak spill	0 B

File an issue against this benchmark runner

fix(arrow-select): merge dictionaries in merge and merge_n instead of…

4253beb

… panicking on key overflow

github-actions Bot added the arrow-kernels Crates that transform arrow data (e.g. compute, cast) label Oct 26, 2026

Rich-T-kid reviewed Oct 27, 2026

View reviewed changes

bench(arrow): add dictionary cases to merge_kernels, including one th…

9f1e1c9

…at merges dictionaries

github-actions Bot added the arrow-core Core crates related to representing arrow data label Oct 27, 2026

Rich-T-kid mentioned this pull request Oct 27, 2026

fix(arrow-select): merge dictionaries in zip instead of panicking on key overflow #11663

Merged

Rich-T-kid reviewed Oct 27, 2026

View reviewed changes

Rich-T-kid approved these changes Oct 27, 2026

View reviewed changes

Rich-T-kid added this pull request to the merge queue Oct 28, 2026

Merged via the queue into apache:main with commit 14c2fef Oct 28, 2026
21 checks passed

Jefffrey added the bug label Nov 1, 2026

Navigation Menu

Repository navigation

fix(arrow-select): merge dictionaries in merge and merge_n instead of panicking on key overflow - #11925

fix(arrow-select): merge dictionaries in merge and merge_n instead of panicking on key overflow#11925
Rich-T-kid merged 2 commits into
apache:mainfrom
ninglinLiu:fix/merge-dict-overflow

ninglinLiu commented Oct 26, 2026

Uh oh!

Rich-T-kid left a comment •

edited

Loading

Uh oh!

ninglinLiu commented Oct 27, 2026

Uh oh!

Rich-T-kid commented Oct 27, 2026

Uh oh!

ninglinLiu commented Oct 27, 2026 •

edited

Loading

Uh oh!

Rich-T-kid Oct 27, 2026

Uh oh!

ninglinLiu Oct 27, 2026

Uh oh!

Rich-T-kid Oct 27, 2026

Uh oh!

Rich-T-kid left a comment

Uh oh!

Rich-T-kid commented Oct 27, 2026

Uh oh!

adriangbot commented Oct 27, 2026

Uh oh!

adriangbot commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

ninglinLiu commented Oct 26, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Rich-T-kid left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ninglinLiu commented Oct 27, 2026

Uh oh!

Rich-T-kid commented Oct 27, 2026

Uh oh!

ninglinLiu commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Rich-T-kid Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

ninglinLiu Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

Rich-T-kid Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

Rich-T-kid left a comment

Choose a reason for hiding this comment

Uh oh!

Rich-T-kid commented Oct 27, 2026

Uh oh!

adriangbot commented Oct 27, 2026

Uh oh!

adriangbot commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Rich-T-kid left a comment •

edited

Loading

ninglinLiu commented Oct 27, 2026 •

edited

Loading