Some more cleanup found in the last commit

Author: zanmato1984

cpp/src/arrow/compute/key_map_internal.cc

@@ -272,14 +272,14 @@ void SwissTable::early_filter_imp(const int num_keys, const uint32_t* hashes,
 // How many groups we can keep in the hash table without the need for resizing.
 // When we reach this limit, we need to break processing of any further rows and resize.
 //
-uint64_t SwissTable::num_groups_for_resize() const {
+int64_t SwissTable::num_groups_for_resize() const {
   // Consider N = 9 (aka 2 ^ 9 = 512 blocks) as small.
   // When N = 9, a slot id takes N + 3 = 12 bits, rounded up to 16 bits. This is also the
   // number of bits needed for a key id. Since each slot stores a status byte and a key
   // id, then a slot takes 1 byte + 16 bits = 3 bytes. Therefore a block of 8 slots takes
   // 24 bytes. The threshold of a small hash table ends up being 24 bytes * 512 = 12 KB.
   constexpr int log_blocks_small_ = 9;
-  uint64_t num_slots = 1ULL << (log_blocks_ + 3);
+  int64_t num_slots = num_slots_from_log_blocks(log_blocks_);
   if (log_blocks_ <= log_blocks_small_) {
     // Resize small hash tables when 50% full.
     return num_slots / 2;
@@ -290,7 +290,8 @@ 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<uint32_t>((1ULL << (log_blocks_ + 3)) - 1);
+  uint32_t global_slot_id_mask =
+      static_cast<uint32_t>((1ULL << (log_blocks_ + kLogSlotsPerBlock)) - 1ULL);
   return global_slot_id & global_slot_id_mask;
 }
 
@@ -398,18 +399,20 @@ bool SwissTable::find_next_stamp_match(const uint32_t hash, const uint32_t in_sl
   int local_slot;
   const uint8_t* blockbase;
   for (;;) {
-    blockbase = block_data(start_slot_id >> 3, num_block_bytes);
+    blockbase = block_data(start_slot_id >> kLogSlotsPerBlock, num_block_bytes);
     uint64_t block = *reinterpret_cast<const uint64_t*>(blockbase);
 
-    search_block<true>(block, stamp, start_slot_id & 7, &local_slot, &match_found);
+    search_block<true>(block, stamp, start_slot_id & kLocalSlotMask, &local_slot,
+                       &match_found);
 
-    start_slot_id = wrap_global_slot_id((start_slot_id & ~7U) + local_slot + match_found);
+    start_slot_id =
+        wrap_global_slot_id((start_slot_id & ~kLocalSlotMask) + local_slot + match_found);
 
     // Match found can be 1 in two cases:
     // - match was found
     // - match was not found in a full block
     // In the second case search needs to continue in the next block.
-    if (match_found == 0 || blockbase[7 - local_slot] == stamp) {
+    if (match_found == 0 || blockbase[kMaxLocalSlot - local_slot] == stamp) {
       break;
     }
   }
@@ -635,7 +638,7 @@ Status SwissTable::map_new_keys(uint32_t num_ids, uint16_t* ids, const uint32_t*
         // First slot in the new starting block
         const int16_t id = ids[i];
         uint32_t block_id = block_id_from_hash(hashes[id], log_blocks_);
-        slot_ids[id] = global_slot_id(block_id, 0);
+        slot_ids[id] = global_slot_id(block_id, /*local_slot_id=*/0);
       }
     }
   } while (num_ids > 0);
@@ -647,7 +650,8 @@ Status SwissTable::grow_double() {
   // Before and after metadata
   int num_group_id_bits_before = num_groupid_bits_from_log_blocks(log_blocks_);
   int num_group_id_bits_after = num_groupid_bits_from_log_blocks(log_blocks_ + 1);
-  uint64_t group_id_mask_before = ~0ULL >> (64 - num_group_id_bits_before);
+  uint32_t group_id_mask_before =
+      group_id_mask_from_num_groupid_bits(num_group_id_bits_before);
   int log_blocks_after = log_blocks_ + 1;
   int bits_shift_for_block_and_stamp_after =
       ComputeBitsShiftForBlockAndStamp(log_blocks_after);
@@ -657,7 +661,7 @@ Status SwissTable::grow_double() {
   int64_t block_size_total_after =
       num_bytes_total_blocks(block_size_after, log_blocks_after);
   int64_t hashes_size_total_after =
-      (bits_hash_ / 8 * (1 << (log_blocks_after + 3))) + padding_;
+      (bits_hash_ / 8 * num_slots_from_log_blocks(log_blocks_after)) + padding_;
   constexpr uint32_t stamp_mask = (1 << bits_stamp_) - 1;
 
   // Allocate new buffers
@@ -685,7 +689,7 @@ Status SwissTable::grow_double() {
     util::SafeStore(double_block_base_new + block_size_after, kHighBitOfEachByte);
 
     for (uint32_t j = 0; j < full_slots; ++j) {
-      uint64_t slot_id = global_slot_id(i, j);
+      uint32_t slot_id = global_slot_id(i, j);
       uint32_t hash = hashes()[slot_id];
       uint32_t block_id_new = block_id_from_hash(hash, log_blocks_after);
       bool is_overflow_entry = ((block_id_new >> 1) != static_cast<uint64_t>(i));
@@ -695,22 +699,22 @@ Status SwissTable::grow_double() {
 
       uint32_t ihalf = block_id_new & 1;
       uint8_t stamp_new = (hash >> bits_shift_for_block_and_stamp_after) & stamp_mask;
-      uint64_t group_id_bit_offs = j * num_group_id_bits_before;
-      uint64_t group_id =
-          (util::SafeLoadAs<uint64_t>(block_base + bytes_status_in_block_ +
+      int group_id_bit_offs = j * num_group_id_bits_before;
+      uint32_t group_id =
+          (util::SafeLoadAs<uint32_t>(block_base + bytes_status_in_block_ +
                                       (group_id_bit_offs >> 3)) >>
            (group_id_bit_offs & 7)) &
           group_id_mask_before;
 
-      uint64_t slot_id_new = global_slot_id(i * 2 + ihalf, full_slots_new[ihalf]);
+      uint32_t slot_id_new = global_slot_id(i * 2 + ihalf, full_slots_new[ihalf]);
       hashes_new[slot_id_new] = hash;
       uint8_t* block_base_new = double_block_base_new + ihalf * block_size_after;
-      block_base_new[7 - full_slots_new[ihalf]] = stamp_new;
-      int64_t group_id_bit_offs_new = full_slots_new[ihalf] * num_group_id_bits_after;
+      block_base_new[kMaxLocalSlot - full_slots_new[ihalf]] = stamp_new;
+      int group_id_bit_offs_new = full_slots_new[ihalf] * num_group_id_bits_after;
       uint64_t* ptr = reinterpret_cast<uint64_t*>(
           block_base_new + bytes_status_in_block_ + (group_id_bit_offs_new >> 3));
-      util::SafeStore(ptr,
-                      util::SafeLoad(ptr) | (group_id << (group_id_bit_offs_new & 7)));
+      util::SafeStore(ptr, util::SafeLoad(ptr) | (static_cast<uint64_t>(group_id)
+                                                  << (group_id_bit_offs_new & 7)));
       full_slots_new[ihalf]++;
     }
   }
@@ -724,17 +728,17 @@ Status SwissTable::grow_double() {
     uint32_t full_slots = CountLeadingZeros(block & kHighBitOfEachByte) >> 3;
 
     for (uint32_t j = 0; j < full_slots; ++j) {
-      uint64_t slot_id = global_slot_id(i, j);
+      uint32_t slot_id = global_slot_id(i, j);
       uint32_t hash = hashes()[slot_id];
       uint32_t block_id_new = block_id_from_hash(hash, log_blocks_after);
       bool is_overflow_entry = ((block_id_new >> 1) != static_cast<uint64_t>(i));
       if (!is_overflow_entry) {
         continue;
       }
 
-      uint64_t group_id_bit_offs = j * num_group_id_bits_before;
-      uint64_t group_id =
-          (util::SafeLoadAs<uint64_t>(block_base + bytes_status_in_block_ +
+      int group_id_bit_offs = j * num_group_id_bits_before;
+      uint32_t group_id =
+          (util::SafeLoadAs<uint32_t>(block_base + bytes_status_in_block_ +
                                       (group_id_bit_offs >> 3)) >>
            (group_id_bit_offs & 7)) &
           group_id_mask_before;
@@ -753,13 +757,13 @@ Status SwissTable::grow_double() {
             static_cast<int>(CountLeadingZeros(block_new & kHighBitOfEachByte) >> 3);
       }
 
-      hashes_new[block_id_new * 8 + full_slots_new] = hash;
-      block_base_new[7 - full_slots_new] = stamp_new;
-      int64_t group_id_bit_offs_new = full_slots_new * num_group_id_bits_after;
+      hashes_new[block_id_new * kSlotsPerBlock + full_slots_new] = hash;
+      block_base_new[kMaxLocalSlot - full_slots_new] = stamp_new;
+      int group_id_bit_offs_new = full_slots_new * num_group_id_bits_after;
       uint64_t* ptr = reinterpret_cast<uint64_t*>(
           block_base_new + bytes_status_in_block_ + (group_id_bit_offs_new >> 3));
-      util::SafeStore(ptr,
-                      util::SafeLoad(ptr) | (group_id << (group_id_bit_offs_new & 7)));
+      util::SafeStore(ptr, util::SafeLoad(ptr) | (static_cast<uint64_t>(group_id)
+                                                  << (group_id_bit_offs_new & 7)));
     }
   }
 
@@ -800,9 +804,9 @@ Status SwissTable::init(int64_t hardware_flags, MemoryPool* pool, int log_blocks
   if (no_hash_array) {
     hashes_ = nullptr;
   } else {
-    uint64_t num_slots = 1ULL << (log_blocks_ + 3);
-    const uint64_t hash_size = sizeof(uint32_t);
-    const uint64_t hash_bytes = hash_size * num_slots + padding_;
+    int64_t num_slots = num_slots_from_log_blocks(log_blocks);
+    const int hash_size = bits_hash_ >> 3;
+    const int64_t hash_bytes = hash_size * num_slots + padding_;
     ARROW_ASSIGN_OR_RAISE(hashes_, AllocateBuffer(hash_bytes, pool_));
   }
 

cpp/src/arrow/compute/key_map_internal.h

@@ -90,7 +90,7 @@ class ARROW_EXPORT SwissTable {
   static uint32_t extract_group_id(const uint8_t* block_ptr, int local_slot,
                                    int num_group_id_bits) {
     // Extract group id using aligned 32-bit read.
-    uint32_t group_id_mask = static_cast<uint32_t>((1ULL << num_group_id_bits) - 1);
+    uint32_t group_id_mask = group_id_mask_from_num_groupid_bits(num_group_id_bits);
     int slot_bit_offset = local_slot * num_group_id_bits;
     const uint32_t* group_id_ptr32 =
         reinterpret_cast<const uint32_t*>(block_ptr + bytes_status_in_block_) +
@@ -99,7 +99,8 @@ class ARROW_EXPORT SwissTable {
     return group_id;
   }
 
-  inline void insert_into_empty_slot(uint32_t slot_id, uint32_t hash, uint32_t group_id);
+  inline void insert_into_empty_slot(uint32_t global_slot_id, uint32_t hash,
+                                     uint32_t group_id);
 
   static uint32_t block_id_from_hash(uint32_t hash, int log_blocks) {
     return hash >> (bits_hash_ - log_blocks);
@@ -110,22 +111,16 @@ class ARROW_EXPORT SwissTable {
   }
 
   static int num_groupid_bits_from_log_blocks(int log_blocks) {
-    assert(log_blocks >= 0 && log_blocks <= 32 - 3);
-    int required_bits = log_blocks + 3;
-    return required_bits <= 8    ? 8
-           : required_bits <= 16 ? 16
-           : required_bits <= 32 ? 32
-                                 : 64;
+    assert(log_blocks >= 0);
+    int required_bits = log_blocks + kLogSlotsPerBlock;
+    assert(required_bits <= 32);
+    return required_bits <= 8 ? 8 : required_bits <= 16 ? 16 : 32;
   }
 
   static int num_block_bytes_from_num_groupid_bits(int num_groupid_bits) {
     return num_groupid_bits + bytes_status_in_block_;
   }
 
-  static int64_t num_bytes_total_blocks(int num_block_bytes, int log_blocks) {
-    return (static_cast<int64_t>(num_block_bytes) << log_blocks) + padding_;
-  }
-
   const uint8_t* block_data(uint32_t block_id, int num_block_bytes) const {
     return block_data(blocks_->data(), block_id, num_block_bytes);
   }
@@ -188,10 +183,24 @@ class ARROW_EXPORT SwissTable {
                              const uint32_t* hashes, const uint8_t* local_slots,
                              uint32_t* out_group_ids) const;
 
-  inline uint64_t next_slot_to_visit(uint64_t block_index, int slot,
-                                     int match_found) const;
+  static constexpr int kLogSlotsPerBlock = 3;
+  static constexpr int kMaxLocalSlot = kSlotsPerBlock - 1;
+  static constexpr uint32_t kLocalSlotMask = (1U << kLogSlotsPerBlock) - 1U;
+
+  static int64_t num_slots_from_log_blocks(int log_blocks) {
+    return 1LL << (log_blocks + kLogSlotsPerBlock);
+  }
+
+  static int64_t num_bytes_total_blocks(int num_block_bytes, int log_blocks) {
+    return (static_cast<int64_t>(num_block_bytes) << log_blocks) + padding_;
+  }
+
+  inline int64_t num_groups_for_resize() const;
 
-  inline uint64_t num_groups_for_resize() const;
+  static uint32_t group_id_mask_from_num_groupid_bits(int num_groupid_bits) {
+    // num_groupid_bits could be 32, so using 64-bit shifting.
+    return static_cast<uint32_t>((1ULL << num_groupid_bits) - 1ULL);
+  }
 
   inline uint32_t wrap_global_slot_id(uint32_t global_slot_id) const;
 
@@ -307,7 +316,7 @@ class ARROW_EXPORT SwissTable {
   MemoryPool* pool_;
 };
 
-void SwissTable::insert_into_empty_slot(uint32_t slot_id, uint32_t hash,
+void SwissTable::insert_into_empty_slot(uint32_t global_slot_id, uint32_t hash,
                                         uint32_t group_id) {
   const int num_groupid_bits = num_groupid_bits_from_log_blocks(log_blocks_);
 
@@ -317,19 +326,20 @@ void SwissTable::insert_into_empty_slot(uint32_t slot_id, uint32_t hash,
          num_groupid_bits == 64);
 
   const int num_block_bytes = num_block_bytes_from_num_groupid_bits(num_groupid_bits);
-  constexpr uint64_t stamp_mask = 0x7f;
+  constexpr uint32_t stamp_mask = 0x7f;
 
-  int start_slot = (slot_id & 7);
-  int stamp = static_cast<int>((hash >> bits_shift_for_block_and_stamp_) & stamp_mask);
-  uint32_t block_id = slot_id >> 3;
+  int start_slot = (global_slot_id & kLocalSlotMask);
+  int stamp = (hash >> bits_shift_for_block_and_stamp_) & stamp_mask;
+  uint32_t block_id = global_slot_id >> kLogSlotsPerBlock;
   uint8_t* blockbase = mutable_block_data(block_id, num_block_bytes);
 
-  blockbase[7 - start_slot] = static_cast<uint8_t>(stamp);
-  int groupid_bit_offset = static_cast<int>(start_slot * num_groupid_bits);
+  blockbase[kMaxLocalSlot - start_slot] = static_cast<uint8_t>(stamp);
+  int groupid_bit_offset = start_slot * num_groupid_bits;
 
   // Block status bytes should start at an address aligned to 8 bytes
   assert((reinterpret_cast<uint64_t>(blockbase) & 7) == 0);
-  uint64_t* ptr = reinterpret_cast<uint64_t*>(blockbase) + 1 + (groupid_bit_offset >> 6);
+  uint64_t* ptr = reinterpret_cast<uint64_t*>(blockbase + bytes_status_in_block_) +
+                  (groupid_bit_offset >> 6);
   *ptr |= (static_cast<uint64_t>(group_id) << (groupid_bit_offset & 63));
 }