From 4bc99671a1866588762cbb9e855b8eab7cef73b4 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 27 Dec 2024 18:23:58 +0800 Subject: [PATCH] Fix --- cpp/src/arrow/compute/key_map_internal.cc | 9 +++++++-- cpp/src/arrow/compute/key_map_internal.h | 2 ++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/key_map_internal.cc b/cpp/src/arrow/compute/key_map_internal.cc index bbaf21ec4e25e..b4d30469c8cb1 100644 --- a/cpp/src/arrow/compute/key_map_internal.cc +++ b/cpp/src/arrow/compute/key_map_internal.cc @@ -258,6 +258,11 @@ uint64_t SwissTable::num_groups_for_resize() const { } } +uint32_t SwissTable::wrap_global_slot_id(uint32_t global_slot_id) const { + uint32_t global_slot_id_mask = static_cast((1ULL << (log_blocks_ + 3)) - 1); + return global_slot_id & global_slot_id_mask; +} + void SwissTable::early_filter(const int num_keys, const uint32_t* hashes, uint8_t* out_match_bitvector, uint8_t* out_local_slots) const { @@ -357,7 +362,7 @@ bool SwissTable::find_next_stamp_match(const uint32_t hash, const uint32_t in_sl constexpr uint64_t stamp_mask = 0x7f; const int stamp = static_cast((hash >> bits_shift_for_block_and_stamp_) & stamp_mask); - uint32_t start_slot_id = in_slot_id; + uint32_t start_slot_id = wrap_global_slot_id(in_slot_id); int match_found; int local_slot; uint8_t* blockbase; @@ -367,7 +372,7 @@ bool SwissTable::find_next_stamp_match(const uint32_t hash, const uint32_t in_sl search_block(block, stamp, start_slot_id & 7, &local_slot, &match_found); - start_slot_id = (start_slot_id & ~7U) + local_slot + match_found; + start_slot_id = wrap_global_slot_id((start_slot_id & ~7U) + local_slot + match_found); // Match found can be 1 in two cases: // - match was found diff --git a/cpp/src/arrow/compute/key_map_internal.h b/cpp/src/arrow/compute/key_map_internal.h index 3c8b9bf47c983..79ba7e23ad6b1 100644 --- a/cpp/src/arrow/compute/key_map_internal.h +++ b/cpp/src/arrow/compute/key_map_internal.h @@ -168,6 +168,8 @@ class ARROW_EXPORT SwissTable { inline uint64_t num_groups_for_resize() const; + inline uint32_t wrap_global_slot_id(uint32_t global_slot_id) const; + void init_slot_ids(const int num_keys, const uint16_t* selection, const uint32_t* hashes, const uint8_t* local_slots, const uint8_t* match_bitvector, uint32_t* out_slot_ids) const;