[C++][Acero] Hash join silently drops matches when the build side has more than ~400M distinct keys (32-bit overflow in SwissTable::early_filter_imp_avx2_x8)

[jachymb]

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 (3ad410b).

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). HasEfficientBmi2() is vendor() == Vendor::Intel && IsSupported(BMI2) (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).

repro.py

"""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()

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).

$ 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:

    __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).

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), 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), and HasEfficientBmi2() is true only for Intel (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:

    __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

• [C++][Python] Pyarrow.Table.join() breaks on large tables v.18.0.0.dev486 #44513 / GH-44513: [C++] Fix overflow issues for large build side in swiss join #45108: overflow fixes for a large build side; widened extract_group_ids_avx2 to 64-bit offsets.
• [C++][Acero] Overflow issues in swiss join when the number of build side rows is big enough #45334 / GH-45334: [C++][Acero] Fix swiss join overflow issues in row offset calculation for fixed length and null masks #45336: row offset overflow in the swiss join.
• [C++][Acero] Swiss table still has risks of overflow #45506 / GH-45506: [C++][Acero] More overflow-safe Swiss table #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++

[jachymb]

PR: #51497

[zanmato1984]

Issue resolved by pull request 51497
#51497