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.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_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..402d43f1b0b2 --- /dev/null +++ b/cpp/src/arrow/compute/key_map_test.cc @@ -0,0 +1,158 @@ +// 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); + } + } +} + +// 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 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',