jachymb opened a new issue, #51496: URL: https://github.com/apache/arrow/issues/51496
### Describe the bug, including details regarding any error messages, version, and platform. **Disclaimer**: This report is AI-generated by an Opus5.5 agent, I honestly don't really understand what's going on, I just know there is a bug that wrecking my tables ### Summary An inner hash join (`pyarrow.Table.join`, Acero `HashJoinNode`) returns too few rows when the build side (the right input) has more than about 403 million distinct keys. It raises no error or warning. Beyond ~403M build rows about 20% of the matches are lost, beyond ~805M about 60%, and beyond ~1.61B about 80%. This only happens when the AVX2 early-filter kernel is used, that is on Intel x86-64 CPUs with AVX2 and BMI2 (see "Affected hardware"). With `ARROW_USER_SIMD_LEVEL=NONE` the same join is exact. If the sides are swapped so that the small table is the build side, the join is also exact. It reproduces with pyarrow 20.0.0 and 25.0.1, and the code in question is the same on `main` today (3ad410b7b1a2d3712e2d6a55d905a7035d026cae). Expected: every probe row whose key is in the build table produces one output row. In the reproducer below, that is exactly 2,000,000 rows. ### Affected hardware The bug is hardware-specific. It needs **all** of the following: - an x86-64 build with runtime AVX2 and BMI2 dispatch (`ARROW_HAVE_RUNTIME_AVX2` and `ARROW_HAVE_RUNTIME_BMI2`), as in the standard x86-64 pyarrow wheels; - a CPU that Arrow detects as **Intel** and that supports **AVX2** and **BMI2** (Haswell, 2013, or later); - a SIMD level that is not lowered below AVX2 with `ARROW_USER_SIMD_LEVEL`. This is the kernel's dispatch condition, `(hardware_flags_ & CpuInfo::AVX2) && CpuInfo::GetInstance()->HasEfficientBmi2()` ([key_map_internal.cc#L304-L313](https://github.com/apache/arrow/blob/3ad410b7b1a2d3712e2d6a55d905a7035d026cae/cpp/src/arrow/compute/key_map_internal.cc#L304-L313)). `HasEfficientBmi2()` is `vendor() == Vendor::Intel && IsSupported(BMI2)` ([cpu_info.h#L104-L107](https://github.com/apache/arrow/blob/3ad410b7b1a2d3712e2d6a55d905a7035d026cae/cpp/src/arrow/util/cpu_info.h#L104-L107)). | Platform | Status | Basis | |---|---|---| | Intel Xeon with AVX2 + BMI2 (tested: Google Cloud `n1-highmem-32` / `n1-highmem-64`, "Intel(R) Xeon(R) CPU @ 2.00GHz", AVX-512) | **Affected** | Verified: the reproducer below, and the sweep further down | | Same Intel machine with `ARROW_USER_SIMD_LEVEL=NONE` | Not affected (exact) | Verified at 900M and 1.7B build rows | | AMD x86-64, even with AVX2 + BMI2 (e.g. Zen) | Not affected | Source only: `HasEfficientBmi2()` is false for AMD, so the scalar `early_filter_imp` runs. Not tested above the threshold. | | ARM / aarch64 (e.g. Apple Silicon, Graviton) | Not affected | Source only: the kernel is x86 AVX2 code and is not compiled or dispatched there. Not tested. | | x86-64 builds without runtime AVX2 dispatch | Not affected | Source only: the `#if defined(ARROW_HAVE_RUNTIME_AVX2) && defined(ARROW_HAVE_RUNTIME_BMI2)` block is compiled out | So the same data and the same pyarrow version give a correct join on AMD or ARM and a wrong one on Intel. Results also differ between machines in a heterogeneous cluster. ### Reproducer The script needs only numpy and pyarrow. It builds a right-hand table of N unique int64 keys (default 450M) and a left-hand table of 2M keys that are all present in it. It runs `left.join(right, "k", join_type="inner")`, then the same join with the sides swapped, and then runs itself again in a subprocess with `ARROW_USER_SIMD_LEVEL=NONE`. Peak memory is roughly 80 bytes per build row: 35.7 GB at the default size, 68.6 GB at 900M (measured). <details> <summary>repro.py</summary> ```python """Reproducer: pyarrow Table.join (inner) silently drops matches when the build side is large. Usage: python repro.py [BUILD_ROWS] (default 450_000_000; 900e6 style also works) Builds a right-hand (build side) table of BUILD_ROWS unique int64 keys and a left-hand (probe side) table of 2,000,000 keys that are all present in it, so an inner join must return exactly 2,000,000 rows. Then runs the same join with the sides swapped, and finally re-runs itself in a subprocess with ARROW_USER_SIMD_LEVEL=NONE. Peak memory is about 80 bytes per build row: ~36 GB at the default size, ~69 GB at 900e6. Only numpy and pyarrow are required. """ import os import platform import subprocess import sys import time import numpy as np import pyarrow as pa N = int(float(sys.argv[1])) if len(sys.argv) > 1 else 450_000_000 # build-side rows M = 2_000_000 # probe-side rows, every one has exactly one match def cpu_model(): try: with open("/proc/cpuinfo") as f: return next(l.split(":", 1)[1].strip() for l in f if l.startswith("model name")) except (OSError, StopIteration): return platform.processor() def peak_rss(): try: import resource # not on Windows except ImportError: return "n/a" rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss # KiB on Linux, bytes on macOS return f"{rss * (1 if sys.platform == 'darwin' else 1024) / 1e9:.1f} GB" def check(label, left, right): t0 = time.time() got = left.join(right, keys="k", join_type="inner").num_rows status = "OK" if got == M else "WRONG" print(f"{label}: expected {M:,} rows, got {got:,} ({got / M:.2%}) " f"[{time.time() - t0:.0f}s] {status}", flush=True) return got == M def main(): simd = pa.runtime_info() print(f"pyarrow {pa.__version__}, numpy {np.__version__}, Python {platform.python_version()}, " f"{platform.system()} {platform.machine()}") print(f"CPU: {cpu_model()}, {os.cpu_count()} logical cores, " f"simd_level={simd.simd_level} (detected {simd.detected_simd_level})") print(f"build side (right): {N:,} unique int64 keys; probe side (left): {M:,} keys, all present", flush=True) build = pa.table({"k": np.arange(N, dtype=np.int64)}) probe = pa.table({"k": np.arange(M, dtype=np.int64) * (N // M)}) # unique, all < N ok = check("probe.join(build) [build side = large table]", probe, build) ok &= check("build.join(probe) [sides swapped, build side = small table]", build, probe) print(f"peak RSS: {peak_rss()}", flush=True) del build, probe pa.default_memory_pool().release_unused() if "ARROW_USER_SIMD_LEVEL" not in os.environ: print("\n--- same script with ARROW_USER_SIMD_LEVEL=NONE ---", flush=True) env = dict(os.environ, ARROW_USER_SIMD_LEVEL="NONE") subprocess.run([sys.executable, os.path.abspath(__file__), str(N)], env=env) sys.exit(0 if ok else 1) if __name__ == "__main__": main() ``` </details> ### Reproducer output Google Cloud `n1-highmem-32` (Intel Xeon @ 2.00GHz, 32 vCPUs, AVX-512, 208 GB RAM), Linux, Python 3.10.12, pyarrow 25.0.1 from PyPI, numpy 1.25.1. `exit code` is the script's exit status (1 = wrong count). ```text $ python repro.py # default: 450M build rows pyarrow 25.0.1, numpy 1.25.1, Python 3.10.12, Linux x86_64 CPU: Intel(R) Xeon(R) CPU @ 2.00GHz, 32 logical cores, simd_level=avx512 (detected avx512) build side (right): 450,000,000 unique int64 keys; probe side (left): 2,000,000 keys, all present probe.join(build) [build side = large table]: expected 2,000,000 rows, got 1,603,163 (80.16%) [12s] WRONG build.join(probe) [sides swapped, build side = small table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [1s] OK peak RSS: 35.7 GB --- same script with ARROW_USER_SIMD_LEVEL=NONE --- pyarrow 25.0.1, numpy 1.25.1, Python 3.10.12, Linux x86_64 CPU: Intel(R) Xeon(R) CPU @ 2.00GHz, 32 logical cores, simd_level=none (detected avx512) build side (right): 450,000,000 unique int64 keys; probe side (left): 2,000,000 keys, all present probe.join(build) [build side = large table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [13s] OK build.join(probe) [sides swapped, build side = small table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [1s] OK peak RSS: 35.7 GB exit code 1 $ python repro.py 900e6 pyarrow 25.0.1, numpy 1.25.1, Python 3.10.12, Linux x86_64 CPU: Intel(R) Xeon(R) CPU @ 2.00GHz, 32 logical cores, simd_level=avx512 (detected avx512) build side (right): 900,000,000 unique int64 keys; probe side (left): 2,000,000 keys, all present probe.join(build) [build side = large table]: expected 2,000,000 rows, got 814,637 (40.73%) [24s] WRONG build.join(probe) [sides swapped, build side = small table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [1s] OK peak RSS: 68.6 GB --- same script with ARROW_USER_SIMD_LEVEL=NONE --- pyarrow 25.0.1, numpy 1.25.1, Python 3.10.12, Linux x86_64 CPU: Intel(R) Xeon(R) CPU @ 2.00GHz, 32 logical cores, simd_level=none (detected avx512) build side (right): 900,000,000 unique int64 keys; probe side (left): 2,000,000 keys, all present probe.join(build) [build side = large table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [24s] OK build.join(probe) [sides swapped, build side = small table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [1s] OK peak RSS: 68.6 GB exit code 1 $ python repro.py 390e6 # below the threshold: exact pyarrow 25.0.1, numpy 1.25.1, Python 3.10.12, Linux x86_64 CPU: Intel(R) Xeon(R) CPU @ 2.00GHz, 32 logical cores, simd_level=avx512 (detected avx512) build side (right): 390,000,000 unique int64 keys; probe side (left): 2,000,000 keys, all present probe.join(build) [build side = large table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [10s] OK build.join(probe) [sides swapped, build side = small table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [0s] OK peak RSS: 26.7 GB --- same script with ARROW_USER_SIMD_LEVEL=NONE --- pyarrow 25.0.1, numpy 1.25.1, Python 3.10.12, Linux x86_64 CPU: Intel(R) Xeon(R) CPU @ 2.00GHz, 32 logical cores, simd_level=none (detected avx512) build side (right): 390,000,000 unique int64 keys; probe side (left): 2,000,000 keys, all present probe.join(build) [build side = large table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [11s] OK build.join(probe) [sides swapped, build side = small table]: expected 2,000,000 rows, got 2,000,000 (100.00%) [1s] OK peak RSS: 26.8 GB exit code 0 ``` ### How the loss scales with the build side These numbers come from an earlier sweep with a very similar script. The build side had a unique two-column key (int64, int32) plus one float32 payload column. The probe side had about 10M rows, all present in the build side. The machine was a Linux x86_64 cloud VM with 64 vCPUs on an Intel CPU (`simd_level='avx512'`), running Python 3.10.12, numpy 1.25.1 and pyarrow 25.0.1 from PyPI, with `use_threads=True`. The `log_blocks` column is inferred, not instrumented: the join grows the hash table at 75% fill with 8 slots per block, so it goes past 2^26 blocks above 6 * 2^26 = 402,653,184 rows, past 2^27 above 805,306,368 rows and past 2^28 above 1,610,612,736 rows. The last column is the share of blocks whose byte offset still fits in 32 bits, 107,374,183 / 2^log_blocks (see the next section). | build rows | joined / probe rows | inferred merged `log_blocks` | predicted share not overflowing | |---:|---:|---:|---:| | 300,000,000 | 1.000000 | 26 | 1.00 | | 390,000,000 | 1.000000 | 26 | 1.00 | | 420,000,000 | 0.803329 | 27 | 0.80 | | 540,000,000 | 0.806764 | 27 | 0.80 | | 780,000,000 | 0.821954 | 27 | 0.80 | | 830,000,000 | 0.411051 | 28 | 0.40 | | 1,100,000,000 | 0.422885 | 28 | 0.40 | | 1,300,000,000 | 0.438211 | 28 | 0.40 | | 1,500,000,000 | 0.459756 | 28 | 0.40 | | 1,590,000,000 | 0.471690 | 28 | 0.40 | | 1,630,000,000 | 0.213826 | 29 | 0.20 | | 1,700,000,000 | 0.215192 | 29 | 0.20 | Other runs on the same machine: - `use_threads=False`, 900M build rows: 0.413007. This rules out a race. - `ARROW_USER_SIMD_LEVEL=NONE` (`simd_level='none'`): 1.000000 at 900M and at 1.7B build rows. - pyarrow 20.0.0: 0.413163 at 900M and 0.215248 at 1.7B, the same as 25.0.1. We first hit this in a real pipeline. An inner join of a 49.4M-row left table against a 1.67B-row right table, with unique keys on both sides, returned 10,293,196 rows. An independent numpy sort/searchsorted join gave 47,863,899 rows, so 21.5% of the matches survived, in line with `log_blocks` = 29. No output rows were spurious; only matches were missing. Three runs of that join (the third with one more payload column) returned 10,293,196, 10,292,836 and 10,292,768 rows, so the count varies slightly from run to run. Swapping the sides (`right.join(left)`) returned exactly 47,863,899, and so did the original join after first reducing the right table with `pc.is_in` to the matching keys. ### Suspected cause `SwissTable::early_filter_imp_avx2_x8` computes each block's byte offset with a 32-bit multiply, [key_map_internal_avx2.cc#L56-L61](https://github.com/apache/arrow/blob/3ad410b7b1a2d3712e2d6a55d905a7035d026cae/cpp/src/arrow/compute/key_map_internal_avx2.cc#L56-L61): ```cpp __m256i vblock_offset = _mm256_mullo_epi32(vblock_id, _mm256_set1_epi32(num_block_bytes)); __m256i voffset_A = _mm256_and_si256(vblock_offset, _mm256_set1_epi64x(0xffffffff)); __m256i vstamp_A = _mm256_and_si256(vstamp, _mm256_set1_epi64x(0xffffffff)); __m256i voffset_B = _mm256_srli_epi64(vblock_offset, 32); __m256i vstamp_B = _mm256_srli_epi64(vstamp, 32); ``` These offsets are then used by `_mm256_i64gather_epi64(blocks_i64, voffset_A/B, 1)` ([L63-L66](https://github.com/apache/arrow/blob/3ad410b7b1a2d3712e2d6a55d905a7035d026cae/cpp/src/arrow/compute/key_map_internal_avx2.cc#L63-L66)). When `log_blocks >= 14`, group ids are 32 bits wide, so `num_block_bytes` = 8 status bytes + 8 * 4 bytes = 40. 2^32 / 40 = 107,374,182.4, so for `block_id >= 107,374,183` the product wraps. The kernel then reads its 8 "status bytes" from an unrelated location, usually finds no stamp match, and reports "no match". The share of blocks below that id is 107,374,183 / 2^log_blocks: 0.80 for `log_blocks` = 27, 0.40 for 28 and 0.20 for 29. Those shares are the three plateaus in the table above. The observed ratios sit slightly above these values and rise with fill within each plateau. My guess, not verified: sometimes the wrongly read status bytes happen to produce a candidate (a stamp match, or a "full block"), and the following `find()` step computes its addresses correctly, so it can still locate the key. The scalar path computes the same address with 64-bit arithmetic, `blocks + static_cast<int64_t>(block_id) * num_block_bytes` ([key_map_internal.h#L146-L149](https://github.com/apache/arrow/blob/3ad410b7b1a2d3712e2d6a55d905a7035d026cae/cpp/src/arrow/compute/key_map_internal.h#L146-L149)), which is why `ARROW_USER_SIMD_LEVEL=NONE` gives correct results. The AVX2 kernel is only chosen when `(hardware_flags_ & CpuInfo::AVX2) && CpuInfo::GetInstance()->HasEfficientBmi2()` ([key_map_internal.cc#L304-L313](https://github.com/apache/arrow/blob/3ad410b7b1a2d3712e2d6a55d905a7035d026cae/cpp/src/arrow/compute/key_map_internal.cc#L304-L313)), and `HasEfficientBmi2()` is true only for Intel ([cpu_info.h#L104-L107](https://github.com/apache/arrow/blob/3ad410b7b1a2d3712e2d6a55d905a7035d026cae/cpp/src/arrow/util/cpu_info.h#L104-L107)). So Intel x86 machines are affected, and AMD and ARM machines should not be. The merged build table is sized in `SwissTableMerge::PrepareForMerge` as `log_num_sources + max(partition log_blocks)`. With one partition (`use_threads=False`) the table simply grows to the same size, which matches the single-threaded result above. ### Suggested fix Compute 64-bit offsets, in the same way `extract_group_ids_avx2` has done since #45108. `_mm256_mul_epu32` multiplies the low 32 bits of each 64-bit lane into a 64-bit product, which also gives the A/B (even/odd lane) split directly. This is an untested sketch: ```cpp __m256i vblock_bytes = _mm256_set1_epi64x(num_block_bytes); __m256i voffset_A = _mm256_mul_epu32(vblock_id, vblock_bytes); __m256i voffset_B = _mm256_mul_epu32(_mm256_srli_epi64(vblock_id, 32), vblock_bytes); __m256i vstamp_A = _mm256_and_si256(vstamp, _mm256_set1_epi64x(0xffffffff)); __m256i vstamp_B = _mm256_srli_epi64(vstamp, 32); ``` For a regression test, a table with `log_blocks >= 27` needs 2^27 * 40 B = 5.4 GB of blocks. Two options: - A unit test that creates a `SwissTable` with `log_blocks = 27` directly, inserts a few keys that hash into blocks with id >= 107,374,183, and checks that `early_filter` agrees between the AVX2 and scalar paths. - The existing, skipped `HashJoin.BuildSideLargeRowIds` test (720,898,048 build rows, with the match key chosen to hash to 0xFFFFFFFE, that is, into the last block) looks like it should fail on an Intel AVX2 machine for this reason. I have not run it. ### Related - GH-44513 / #45108: overflow fixes for a large build side; widened `extract_group_ids_avx2` to 64-bit offsets. - GH-45334 / #45336: row offset overflow in the swiss join. - GH-45506 / #45515: "More overflow-safe Swiss table". The issue describes this `num_block_bytes * block_id` overflow risk in general (its example is a scalar site). The PR introduced `block_data()` with 64-bit promotion for the scalar paths, but in `early_filter_imp_avx2_x8` it only replaced the multiplier constant (`num_group_id_bits + 8` became `num_block_bytes`) and kept the 32-bit `_mm256_mullo_epi32`. ### Component(s) C++ -- 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]
