From 664862911acf578a6e6af61990f64ea67a87824f Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Sat, 26 Sep 2026 10:22:45 +0200 Subject: [PATCH 1/3] [C++] Fix 32-bit block offset overflow in SwissTable AVX2 early filter SwissTable::early_filter_imp_avx2_x8 computed the byte offset of each block with _mm256_mullo_epi32, that is modulo 2^32. With 32-bit group ids (log_blocks >= 14) a block takes 40 bytes, so the offset of every block with id >= 107,374,183 (ceil(2^32 / 40)) wrapped around and the kernel read the status bytes of an unrelated location. The early filter then reported no match, or a match in the wrong slot, and find() missed the key. A hash join whose build side has more than about 400M distinct keys (log_blocks >= 27) silently dropped about 20%, 60% or 80% of the matches at log_blocks 27, 28 and 29. The kernel only runs on CPUs with AVX2 and efficient BMI2 (Intel), so the same join is exact on AMD and ARM, and with ARROW_USER_SIMD_LEVEL=NONE. Compute the offsets in 64 bits with _mm256_mul_epu32, which multiplies the even (A) or, after a 32-bit shift, the odd (B) 32-bit block ids into 64-bit products, as extract_group_ids_avx2 has done since GH-44513. Add a large memory test that runs early_filter on the blocks on both sides of the 4GB boundary of a table with 2^27 blocks. Co-Authored-By: Claude Opus 5.5 (1M context) --- cpp/src/arrow/compute/CMakeLists.txt | 1 + .../arrow/compute/key_map_internal_avx2.cc | 8 +- cpp/src/arrow/compute/key_map_test.cc | 84 +++++++++++++++++++ 3 files changed, 89 insertions(+), 4 deletions(-) create mode 100644 cpp/src/arrow/compute/key_map_test.cc diff --git a/cpp/src/arrow/compute/CMakeLists.txt b/cpp/src/arrow/compute/CMakeLists.txt index 88f9ae900647..bfddb66358ab 100644 --- a/cpp/src/arrow/compute/CMakeLists.txt +++ b/cpp/src/arrow/compute/CMakeLists.txt @@ -175,6 +175,7 @@ add_arrow_compute_test(expression_test add_arrow_compute_test(row_test SOURCES key_hash_test.cc + key_map_test.cc light_array_test.cc row/compare_test.cc row/grouper_test.cc diff --git a/cpp/src/arrow/compute/key_map_internal_avx2.cc b/cpp/src/arrow/compute/key_map_internal_avx2.cc index 353d5a59e67d..6f5e34c27b9c 100644 --- a/cpp/src/arrow/compute/key_map_internal_avx2.cc +++ b/cpp/src/arrow/compute/key_map_internal_avx2.cc @@ -52,12 +52,12 @@ int SwissTable::early_filter_imp_avx2_x8(const int num_hashes, const uint32_t* h // We now split inputs and process 4 at a time, // in order to process 64-bit blocks + // Block offsets are computed in 64 bits, as they may not fit in 32 bits. // - __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 voffset_A = _mm256_mul_epu32(vblock_id, _mm256_set1_epi32(num_block_bytes)); __m256i vstamp_A = _mm256_and_si256(vstamp, _mm256_set1_epi64x(0xffffffff)); - __m256i voffset_B = _mm256_srli_epi64(vblock_offset, 32); + __m256i voffset_B = _mm256_mul_epu32(_mm256_srli_epi64(vblock_id, 32), + _mm256_set1_epi32(num_block_bytes)); __m256i vstamp_B = _mm256_srli_epi64(vstamp, 32); auto blocks_i64 = diff --git a/cpp/src/arrow/compute/key_map_test.cc b/cpp/src/arrow/compute/key_map_test.cc new file mode 100644 index 000000000000..928b50b0bc40 --- /dev/null +++ b/cpp/src/arrow/compute/key_map_test.cc @@ -0,0 +1,84 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include +#include + +#include "arrow/compute/key_map_internal.h" +#include "arrow/memory_pool.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/testing/util.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/cpu_info.h" + +namespace arrow { + +using internal::CpuInfo; + +namespace compute { + +// With 32-bit group ids a block takes 40 bytes, so the byte offsets of the blocks from +// id ceil(2^32 / 40) on do not fit in 32 bits. +TEST(SwissTable, LARGE_MEMORY_TEST(EarlyFilterOver4GB)) { + if constexpr (sizeof(void*) == 4) { + GTEST_SKIP() << "Test only works on 64-bit platforms"; + } + + // 2^27 blocks of 40 bytes take 5GB. + constexpr int kLogBlocks = 27; + constexpr uint32_t kFirstBlockOver4GB = 107374183; + constexpr int kNumHashes = 16; + + // One hash for each of the 8 blocks below and the 8 blocks above the 4GB boundary, + // all with a non-zero stamp. + std::vector hashes(kNumHashes); + for (int i = 0; i < kNumHashes; ++i) { + uint32_t block_id = kFirstBlockOver4GB - kNumHashes / 2 + i; + hashes[i] = (block_id << (SwissTable::bits_hash_ - kLogBlocks)) | 1; + } + + for (int64_t hardware_flags : GetSupportedHardwareFlags({CpuInfo::AVX2})) { + ARROW_SCOPED_TRACE("hardware_flags = ", hardware_flags); + SwissTable table; + ASSERT_OK(table.init(hardware_flags, default_memory_pool(), kLogBlocks, + /*no_hash_array=*/true)); + // Insert every other hash into the first slot of its block, leaving the other blocks + // empty. + for (int i = 0; i < kNumHashes; i += 2) { + uint32_t block_id = SwissTable::block_id_from_hash(hashes[i], kLogBlocks); + table.insert_into_empty_slot(SwissTable::global_slot_id(block_id, 0), hashes[i], + /*group_id=*/i); + } + + uint8_t match_bitvector[kNumHashes / 8]; + uint8_t local_slots[kNumHashes]; + table.early_filter(kNumHashes, hashes.data(), match_bitvector, local_slots); + for (int i = 0; i < kNumHashes; ++i) { + ARROW_SCOPED_TRACE("block_id = ", + SwissTable::block_id_from_hash(hashes[i], kLogBlocks)); + // Inserted hashes match in the first slot, the others hit an empty block whose + // first slot is empty. + ASSERT_EQ(bit_util::GetBit(match_bitvector, i), i % 2 == 0); + ASSERT_EQ(local_slots[i], 0); + } + } +} + +} // namespace compute +} // namespace arrow From 6c3ad10e0d438a59ac698ce4eb3dba638b8be4e2 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Sat, 26 Sep 2026 10:24:35 +0200 Subject: [PATCH 2/3] [C++] Fix 32-bit block offset overflow in SwissTable::grow_double When grow_double() reinserts an entry that had overflowed its home block and the entry's new home block is full, it probes the following blocks. The address of those blocks was computed as block_id_new * block_size_after in 32-bit unsigned arithmetic, which wraps for block ids >= 107,374,183 once the table grows to 2^27 blocks (more than about 400M keys in one table). The entry was then written to an unrelated location of the new table, corrupting other entries, and could not be found afterwards. Use mutable_block_data(), which computes the offset in 64 bits, like the other block accesses since GH-45506. Add a large memory test that grows a table from 2^26 blocks with an entry that has to move past a full block across the 4GB boundary. Co-Authored-By: Claude Opus 5.5 (1M context) --- cpp/src/arrow/compute/key_map_internal.cc | 3 +- cpp/src/arrow/compute/key_map_test.cc | 74 +++++++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/key_map_internal.cc b/cpp/src/arrow/compute/key_map_internal.cc index 353449cf1646..975a42c7d34a 100644 --- a/cpp/src/arrow/compute/key_map_internal.cc +++ b/cpp/src/arrow/compute/key_map_internal.cc @@ -744,7 +744,8 @@ Status SwissTable::grow_double() { static_cast(std::countl_zero(block_new & kHighBitOfEachByte) >> 3); while (full_slots_new == kSlotsPerBlock) { block_id_new = (block_id_new + 1) & ((1 << log_blocks_after) - 1); - block_base_new = blocks_new->mutable_data() + block_id_new * block_size_after; + block_base_new = mutable_block_data(blocks_new->mutable_data(), block_id_new, + block_size_after); block_new = util::SafeLoadAs(block_base_new); full_slots_new = static_cast(std::countl_zero(block_new & kHighBitOfEachByte) >> 3); diff --git a/cpp/src/arrow/compute/key_map_test.cc b/cpp/src/arrow/compute/key_map_test.cc index 928b50b0bc40..402d43f1b0b2 100644 --- a/cpp/src/arrow/compute/key_map_test.cc +++ b/cpp/src/arrow/compute/key_map_test.cc @@ -80,5 +80,79 @@ TEST(SwissTable, LARGE_MEMORY_TEST(EarlyFilterOver4GB)) { } } +// When growing to a table over 4GB, entries that have to move past a full block must +// still land in the right block. +TEST(SwissTable, LARGE_MEMORY_TEST(GrowOver4GB)) { + if constexpr (sizeof(void*) == 4) { + GTEST_SKIP() << "Test only works on 64-bit platforms"; + } + + // Grow from 2^26 to 2^27 blocks. + constexpr int kLogBlocks = 26; + constexpr uint32_t kBlockId = (1u << kLogBlocks) - 2; + constexpr int kNumHashes = SwissTable::kSlotsPerBlock + 1; + + // All these hashes map to block kBlockId before growing and to block 2 * kBlockId + // after. The first 8 fill these blocks, so the last one is in block kBlockId + 1 before + // growing and has to move to block 2 * kBlockId + 1, which is over 4GB. + std::vector hashes(kNumHashes); + for (int i = 0; i < kNumHashes; ++i) { + hashes[i] = (kBlockId << (SwissTable::bits_hash_ - kLogBlocks)) | i; + } + + // Key i is equal to group id i. + SwissTable::EqualImpl equal_impl = + [](int num_keys, const uint16_t* selection, const uint32_t* group_ids, + uint32_t* out_num_keys_mismatch, uint16_t* out_selection_mismatch, void*) { + *out_num_keys_mismatch = 0; + for (int i = 0; i < num_keys; ++i) { + uint16_t id = selection ? selection[i] : static_cast(i); + if (group_ids[id] != id) { + out_selection_mismatch[(*out_num_keys_mismatch)++] = id; + } + } + }; + SwissTable::AppendImpl append_impl = [](int, const uint16_t*, void*) { + return Status::OK(); + }; + + for (int64_t hardware_flags : GetSupportedHardwareFlags({CpuInfo::AVX2})) { + ARROW_SCOPED_TRACE("hardware_flags = ", hardware_flags); + SwissTable table; + ASSERT_OK(table.init(hardware_flags, default_memory_pool(), kLogBlocks)); + for (int i = 0; i < kNumHashes; ++i) { + uint32_t slot_id = SwissTable::global_slot_id(kBlockId, 0) + i; + table.insert_into_empty_slot(slot_id, hashes[i], /*group_id=*/i); + table.hashes()[slot_id] = hashes[i]; + } + + // The table grows when 75% of its slots are used. Pretend that it is one key short + // of that and insert one more key. + table.num_inserted( + static_cast((int64_t{1} << (kLogBlocks + 3)) * 3 / 4 - 1)); + util::TempVectorStack temp_stack; + ASSERT_OK(temp_stack.Init(default_memory_pool(), 64 * table.minibatch_size())); + uint16_t new_key_id = 0; + uint32_t new_key_hash = 0; + uint32_t new_group_id; + ASSERT_OK(table.map_new_keys(/*num_ids=*/1, &new_key_id, &new_key_hash, &new_group_id, + &temp_stack, equal_impl, append_impl, + /*callback_ctx=*/nullptr)); + ASSERT_EQ(table.log_blocks(), kLogBlocks + 1); + + uint8_t match_bitvector[(kNumHashes + 7) / 8]; + uint8_t local_slots[kNumHashes]; + uint32_t group_ids[kNumHashes]; + table.early_filter(kNumHashes, hashes.data(), match_bitvector, local_slots); + table.find(kNumHashes, hashes.data(), match_bitvector, local_slots, group_ids, + &temp_stack, equal_impl, /*callback_ctx=*/nullptr); + for (int i = 0; i < kNumHashes; ++i) { + ARROW_SCOPED_TRACE("key = ", i); + ASSERT_TRUE(bit_util::GetBit(match_bitvector, i)); + ASSERT_EQ(group_ids[i], i); + } + } +} + } // namespace compute } // namespace arrow From af5e6567b3282dff32f322d36797b9bec79e3408 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Sat, 26 Sep 2026 11:19:21 +0200 Subject: [PATCH 3/3] [C++] Add key_map_test.cc to the Meson build The new SwissTable tests were only added to the CMake sources of arrow-compute-row-test. Add them to the Meson test sources as well. Co-Authored-By: Claude Opus 5.5 (1M context) --- cpp/src/arrow/compute/meson.build | 1 + 1 file changed, 1 insertion(+) diff --git a/cpp/src/arrow/compute/meson.build b/cpp/src/arrow/compute/meson.build index fed699e1ca41..9641658a252b 100644 --- a/cpp/src/arrow/compute/meson.build +++ b/cpp/src/arrow/compute/meson.build @@ -102,6 +102,7 @@ compute_tests = { 'arrow-compute-row-test': { 'sources': [ 'key_hash_test.cc', + 'key_map_test.cc', 'light_array_test.cc', 'row/compare_test.cc', 'row/grouper_test.cc',