Use aligned read to extract group id

Author: zanmato1984

cpp/src/arrow/compute/key_map_internal.cc

@@ -94,11 +94,12 @@ inline void SwissTable::search_block(uint64_t block, int stamp, int start_slot,
   *out_slot = static_cast<int>(CountLeadingZeros(matches | block_high_bits) >> 3);
 }
 
-template <bool use_selection>
+template <typename T, bool use_selection>
 void SwissTable::extract_group_ids_imp(const int num_keys, const uint16_t* selection,
                                        const uint32_t* hashes, const uint8_t* local_slots,
                                        uint32_t* out_group_ids) const {
   if (log_blocks_ == 0) {
+    DCHECK_EQ(sizeof(T), sizeof(uint8_t));
     for (int i = 0; i < num_keys; ++i) {
       uint32_t id = use_selection ? selection[i] : i;
       uint32_t group_id =
@@ -108,18 +109,16 @@ void SwissTable::extract_group_ids_imp(const int num_keys, const uint16_t* selec
     }
   } else {
     int num_groupid_bits = num_groupid_bits_from_log_blocks(log_blocks_);
-    int num_groupid_bytes = num_groupid_bits / 8;
-    uint32_t group_id_mask = group_id_mask_from_num_groupid_bits(num_groupid_bits);
+    DCHECK_EQ(sizeof(T) * 8, num_groupid_bits);
     int num_block_bytes = num_block_bytes_from_num_groupid_bits(num_groupid_bits);
 
     for (int i = 0; i < num_keys; ++i) {
       uint32_t id = use_selection ? selection[i] : i;
       uint32_t hash = hashes[id];
       uint32_t block_id = block_id_from_hash(hash, log_blocks_);
-      uint32_t group_id = *reinterpret_cast<const uint32_t*>(
-          block_data(block_id, num_block_bytes) + local_slots[id] * num_groupid_bytes +
-          bytes_status_in_block_);
-      group_id &= group_id_mask;
+      const T* slots_base = reinterpret_cast<const T*>(
+          block_data(block_id, num_block_bytes) + bytes_status_in_block_);
+      uint32_t group_id = static_cast<uint32_t>(slots_base[local_slots[id]]);
       out_group_ids[id] = group_id;
     }
   }
@@ -137,13 +136,40 @@ void SwissTable::extract_group_ids(const int num_keys, const uint16_t* optional_
     num_processed = extract_group_ids_avx2(num_keys, hashes, local_slots, out_group_ids);
   }
 #endif
-  if (optional_selection) {
-    extract_group_ids_imp<true>(num_keys, optional_selection, hashes, local_slots,
-                                out_group_ids);
-  } else {
-    extract_group_ids_imp<false>(num_keys - num_processed, nullptr,
-                                 hashes + num_processed, local_slots + num_processed,
-                                 out_group_ids + num_processed);
+  int num_groupid_bits = num_groupid_bits_from_log_blocks(log_blocks_);
+  switch (num_groupid_bits) {
+    case 8:
+      if (optional_selection) {
+        extract_group_ids_imp<uint8_t, true>(num_keys, optional_selection, hashes,
+                                             local_slots, out_group_ids);
+      } else {
+        extract_group_ids_imp<uint8_t, false>(
+            num_keys - num_processed, nullptr, hashes + num_processed,
+            local_slots + num_processed, out_group_ids + num_processed);
+      }
+      break;
+    case 16:
+      if (optional_selection) {
+        extract_group_ids_imp<uint16_t, true>(num_keys, optional_selection, hashes,
+                                              local_slots, out_group_ids);
+      } else {
+        extract_group_ids_imp<uint16_t, false>(
+            num_keys - num_processed, nullptr, hashes + num_processed,
+            local_slots + num_processed, out_group_ids + num_processed);
+      }
+      break;
+    case 32:
+      if (optional_selection) {
+        extract_group_ids_imp<uint32_t, true>(num_keys, optional_selection, hashes,
+                                              local_slots, out_group_ids);
+      } else {
+        extract_group_ids_imp<uint32_t, false>(
+            num_keys - num_processed, nullptr, hashes + num_processed,
+            local_slots + num_processed, out_group_ids + num_processed);
+      }
+      break;
+    default:
+      DCHECK(false);
   }
 }
 

cpp/src/arrow/compute/key_map_internal.h

@@ -88,11 +88,15 @@ class ARROW_EXPORT SwissTable {
   /// \brief Extract group id for a given slot in a given block.
   ///
   static uint32_t extract_group_id(const uint8_t* block_ptr, int local_slot,
-                                   int64_t num_group_id_bits) {
-    uint32_t group_id_mask = group_id_mask_from_num_groupid_bits(num_group_id_bits);
-    uint32_t group_id = *reinterpret_cast<const uint32_t*>(
-        block_ptr + bytes_status_in_block_ + local_slot * num_group_id_bits / 8);
-    return group_id & group_id_mask;
+                                   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);
+    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_) +
+        (slot_bit_offset >> 5);
+    uint32_t group_id = (*group_id_ptr32 >> (slot_bit_offset & 31)) & group_id_mask;
+    return group_id;
   }
 
   inline void insert_into_empty_slot(uint32_t slot_id, uint32_t hash, uint32_t group_id);
@@ -106,6 +110,7 @@ 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
@@ -178,7 +183,7 @@ class ARROW_EXPORT SwissTable {
                          const uint32_t* hashes, const uint8_t* local_slots,
                          uint32_t* out_group_ids) const;
 
-  template <bool use_selection>
+  template <typename T, bool use_selection>
   void extract_group_ids_imp(const int num_keys, const uint16_t* selection,
                              const uint32_t* hashes, const uint8_t* local_slots,
                              uint32_t* out_group_ids) const;
@@ -257,10 +262,6 @@ class ARROW_EXPORT SwissTable {
     return bits_stamp_;
   }
 
-  static uint32_t group_id_mask_from_num_groupid_bits(int64_t num_groupid_bits) {
-    return static_cast<uint32_t>((1ULL << num_groupid_bits) - 1);
-  }
-
   static constexpr int bytes_status_in_block_ = 8;
 
   // Number of hash bits stored in slots in a block.