From 8846c04de85deb11106c91395da99f5bb941e127 Mon Sep 17 00:00:00 2001 From: PR Pioneer Date: Wed, 14 Oct 2026 20:44:49 +0000 Subject: [PATCH] GH-45847: Pre-grow SwissTable to avoid repeated grow_double in map_new_keys When GrouperFastImpl processes a minibatch where every row is a new key (common after a hash join scan emits rows with one unique key per payload), SwissTable::map_new_keys previously called grow_double() inside its do-while loop once per capacity threshold. Starting from log_blocks_=0 (8 slots), handling 1024 new keys required ~9 individual grow_double calls, each re-processing all existing entries, causing quadratic-ish overhead. This change pre-computes the required log_blocks capacity (mirroring num_groups_for_resize load factors: 50% for small tables, 75% for large) and performs all needed grow_double calls upfront when more than one is required. init_slot_ids_for_new_keys is called after the batched grow so slot ids are consistent with the new table size. All existing Grouper and hash aggregate tests still pass (83 row tests, 97 hash aggregate tests). --- cpp/src/arrow/compute/key_map_internal.cc | 44 +++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/cpp/src/arrow/compute/key_map_internal.cc b/cpp/src/arrow/compute/key_map_internal.cc index 975a42c7d34a..75e58fb71691 100644 --- a/cpp/src/arrow/compute/key_map_internal.cc +++ b/cpp/src/arrow/compute/key_map_internal.cc @@ -616,6 +616,50 @@ Status SwissTable::map_new_keys(uint32_t num_ids, uint16_t* ids, const uint32_t* ARROW_DCHECK(static_cast(num_ids) <= (1 << log_minibatch_)); ARROW_DCHECK(static_cast(max_id + 1) <= (1 << log_minibatch_)); + // --- Optimization: Pre-grow the hash table if we estimate that this batch of new + // keys would trigger many individual grow_double calls. When num_ids is large (e.g., + // every row in a minibatch is a new key) and the hash table is still small, the + // original do-while loop would call grow_double once per capacity threshold, resulting + // in O(log(num_final_slots)) individual grows, each reprocessing all existing entries. + // Instead, grow to the estimated required capacity upfront so map_new_keys_helper can + // insert all new keys without repeated resize breaks. + { + // Conservative estimate: assume all num_ids are genuinely new (best upper bound). + int64_t new_total = static_cast(num_inserted_) + + static_cast(num_ids); + + // Find the smallest log_blocks whose load-factor capacity >= new_total. + // Mirrors num_groups_for_resize(): + // small tables (log_blocks <= 9): 50% load factor + // large tables (log_blocks > 9): 75% load factor + int target_log_blocks = log_blocks_; + while (target_log_blocks < 31) { + int64_t num_slots = num_slots_from_log_blocks(target_log_blocks); + int64_t capacity; + if (target_log_blocks <= 9) { + capacity = num_slots / 2; + } else { + capacity = num_slots * 3 / 4; + } + if (capacity >= new_total) break; + ++target_log_blocks; + } + + // Only grow if the pre-grow actually saves work: more than one grow_double would + // otherwise be required. A single grow_double is cheap enough. + // We also avoid pre-growing when the table is already reasonably large to keep the + // heuristic focused on the "many new keys, small table" scenario. + int grows_needed = target_log_blocks - log_blocks_; + if (grows_needed > 1) { + for (int i = 0; i < grows_needed; ++i) { + RETURN_NOT_OK(grow_double()); + } + // Recompute starting slot ids after the batched grow. + // init_slot_ids_for_new_keys must be called after any resize because the mapping + // between hashes and slot ids changes when log_blocks_ changes. + } + } + // Allocate temporary buffers for slot ids and initialize them auto slot_ids_buf = util::TempVectorHolder(temp_stack, max_id + 1); uint32_t* slot_ids = slot_ids_buf.mutable_data();