jachymb opened a new pull request, #51497:
URL: https://github.com/apache/arrow/pull/51497
### Rationale for this change
`SwissTable::early_filter_imp_avx2_x8` computes each block's byte offset
with `_mm256_mullo_epi32`, so modulo 2^32. With 32-bit group ids a block takes
40 bytes, and every block with id >= 107,374,183 (ceil(2^32 / 40)) gets a
wrapped offset. The kernel then reads the status bytes from the wrong place,
reports no match (or the wrong slot), and `find()` misses the key. There is no
error.
A hash join whose build side has more than ~403M distinct keys (`log_blocks`
>= 27) loses about 20%, 60% and 80% of its matches at `log_blocks` 27, 28 and
29. The kernel is only dispatched on CPUs with AVX2 and efficient BMI2
(`HasEfficientBmi2()`, that is Intel), so the same join is correct on AMD, on
ARM and with `ARROW_USER_SIMD_LEVEL=NONE`. See GH-XXXXX for the full analysis
and a pyarrow reproducer.
While checking the other block offset computations I found the same overflow
in scalar code. When `grow_double()` reinserts an overflow entry and the
entry's new home block is full, it probes the next blocks at
`blocks_new->mutable_data() + block_id_new * block_size_after`, a 32-bit
multiply. GH-45506 fixed the line just above but not this one. Once a single
table grows to 2^27 blocks, such an entry is written to an unrelated location
and can't be found. This happens on any CPU. A multi-threaded hash join builds
small per-partition tables and merges them into a table that is sized up front,
so the join numbers below come from the early filter bug alone. A single table
that grows past ~403M keys can hit this one, for example in a group-by or in a
join built as one partition.
### What changes are included in this PR?
- `key_map_internal_avx2.cc`: compute `voffset_A`/`voffset_B` in 64 bits
with `_mm256_mul_epu32` on the even and odd 32-bit block ids. This is the same
even/odd split as before, now without wraparound, and matches what
`extract_group_ids_avx2` does since GH-44513.
- `key_map_internal.cc`: use `mutable_block_data()`, which computes the
offset in 64 bits, for the probe in `grow_double()`.
- New `key_map_test.cc` (added to `arrow-compute-row-test` in CMake and
Meson) with two large-memory tests:
- `SwissTable.EarlyFilterOver4GB`: a table with 2^27 blocks and no hash
array (5.4 GB). It inserts keys into the blocks on both sides of the 4GB
boundary and checks `early_filter` for every supported hardware flag set
(scalar, and AVX2 where available).
- `SwissTable.GrowOver4GB` (about 14.5 GB peak): it grows a table from
2^26 to 2^27 blocks with an entry that has to move past a full block beyond
4GB, then checks `early_filter` and `find` for all keys. It triggers the growth
through `num_inserted()` and the 75% fill threshold, and asserts that the table
grew.
### Are these changes tested?
Yes. Both new tests are `LARGE_MEMORY_TEST`s, so they only run with
`ARROW_LARGE_MEMORY_TESTS=ON`. The AVX2 part of `EarlyFilterOver4GB` only
exercises the bug on Intel CPUs with AVX2 and BMI2.
The results below are from Google Cloud `n1-highmem-32` (Intel Xeon @
2.00GHz, AVX2, BMI2 and AVX-512), Ubuntu 24.04, gcc 13.3, with
`ARROW_RUNTIME_SIMD_LEVEL=MAX`, in Release and in Debug with
`BUILD_WARNING_LEVEL=CHECKIN`:
- `main` with only the new tests added:
- `EarlyFilterOver4GB` fails at `key_map_test.cc:78` with `local_slots[i]`
7, expected 0, at `block_id = 107374183`, `hardware_flags = 32` (AVX2). It
passes with `ARROW_USER_SIMD_LEVEL=NONE`.
- `GrowOver4GB` fails at `key_map_test.cc:151` (key 8 not found,
`hardware_flags = 0`), with any SIMD level.
- The early filter fix alone: `EarlyFilterOver4GB` passes, and `GrowOver4GB`
still fails at `:151`.
- This PR: both tests pass, also with `ARROW_USER_SIMD_LEVEL=NONE` (about 10
s and 16 s in Release).
- `arrow-compute-row-test`: 90 tests pass (the 88 existing ones and the 2
new ones).
- `arrow-acero-hash-join-node-test`: 36 pass and 1 is skipped
(`BuildSideLargeRowIds`, which is skipped in its body), the same as on `main`.
- No new compiler warnings in the Debug `-Werror` build.
Before the fix, the pyarrow reproducer from the issue gave these results on
the same machine type with pyarrow 25.0.1 (an inner join where all 2,000,000
probe keys have a match):
| build rows | rows returned | with `ARROW_USER_SIMD_LEVEL=NONE` |
|---:|---:|---:|
| 390,000,000 | 2,000,000 (100.00%) | 2,000,000 |
| 450,000,000 | 1,603,163 (80.16%) | 2,000,000 |
| 900,000,000 | 814,637 (40.73%) | 2,000,000 |
I did not rerun the pyarrow reproducer on a patched build. The C++ tests
above cover the same code path.
### Are there any user-facing changes?
No API changes. Hash joins with more than ~403M distinct build keys on Intel
AVX2 CPUs, and single hash tables that grow past ~403M keys, now return correct
results.
**This PR contains a "Critical Fix".** Both bugs silently produce incorrect
data. The AVX2 early filter makes hash joins with more than ~403M distinct
build-side keys drop 20% to 80% of the matching rows on Intel CPUs with AVX2
and BMI2. The `grow_double()` overflow can misplace entries, and so lose keys,
when a hash table grows to 2^27 blocks or more (past ~403M keys), on any CPU.
### Was AI used for this PR?
In accordance to the [AI generation
guidelines](https://arrow.apache.org/docs/dev/developers/overview.html#ai-generated-code),
please disclose below whether and how AI was used in this PR.
**PR code and description written by:**
- [ ] Human
- [x] AI
**Reviewed before submission by:**
- [ ] Human
- [x] AI
- [ ] Not reviewed
Claude Code (Claude Opus) did the root-cause analysis and wrote the fix, the
tests and this description, under my direction. I found the bug in a production
pipeline. A separate AI pass reviewed the diff and rebuilt and reran the tests
on an Intel machine.
--
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]
