BUG: translate number of elements into number of needed buckets, avoiding rehashing for some cases

Author: realead

pandas/_libs/hashtable.pyx

@@ -19,6 +19,7 @@ from pandas._libs.khash cimport (
     are_equivalent_float64_t,
     are_equivalent_khcomplex64_t,
     are_equivalent_khcomplex128_t,
+    kh_get_bucket_n,
     kh_str_t,
     khcomplex64_t,
     khcomplex128_t,
@@ -152,7 +153,7 @@ def unique_label_indices(const int64_t[:] labels):
         ndarray[int64_t, ndim=1] arr
         Int64VectorData *ud = idx.data
 
-    kh_resize_int64(table, min(n, SIZE_HINT_LIMIT))
+    kh_resize_int64(table, min(kh_get_bucket_n(n), SIZE_HINT_LIMIT))
 
     with nogil:
         for i in range(n):

pandas/_libs/hashtable_class_helper.pxi.in

@@ -393,7 +393,7 @@ cdef class {{name}}HashTable(HashTable):
     def __cinit__(self, int64_t size_hint=1):
         self.table = kh_init_{{dtype}}()
         if size_hint is not None:
-            size_hint = min(size_hint, SIZE_HINT_LIMIT)
+            size_hint = min(kh_get_bucket_n(size_hint), SIZE_HINT_LIMIT)
             kh_resize_{{dtype}}(self.table, size_hint)
 
     def __len__(self) -> int:
@@ -741,7 +741,7 @@ cdef class StringHashTable(HashTable):
     def __init__(self, int64_t size_hint=1):
         self.table = kh_init_str()
         if size_hint is not None:
-            size_hint = min(size_hint, SIZE_HINT_LIMIT)
+            size_hint = min(kh_get_bucket_n(size_hint), SIZE_HINT_LIMIT)
             kh_resize_str(self.table, size_hint)
 
     def __dealloc__(self):
@@ -1063,7 +1063,7 @@ cdef class PyObjectHashTable(HashTable):
     def __init__(self, int64_t size_hint=1):
         self.table = kh_init_pymap()
         if size_hint is not None:
-            size_hint = min(size_hint, SIZE_HINT_LIMIT)
+            size_hint = min(kh_get_bucket_n(size_hint), SIZE_HINT_LIMIT)
             kh_resize_pymap(self.table, size_hint)
 
     def __dealloc__(self):

pandas/_libs/hashtable_func_helper.pxi.in

@@ -121,7 +121,7 @@ def duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first'):
         kh_{{ttype}}_t *table = kh_init_{{ttype}}()
         ndarray[uint8_t, ndim=1, cast=True] out = np.empty(n, dtype='bool')
 
-    kh_resize_{{ttype}}(table, min(n, SIZE_HINT_LIMIT))
+    kh_resize_{{ttype}}(table, min(kh_get_bucket_n(n), SIZE_HINT_LIMIT))
 
     if keep not in ('last', 'first', False):
         raise ValueError('keep must be either "first", "last" or False')

pandas/_libs/khash.pxd

@@ -120,4 +120,7 @@ cdef extern from "khash_python.h":
 
     bint kh_exist_strbox(kh_strbox_t*, khiter_t) nogil
 
+    khuint_t kh_get_bucket_n(khuint_t element_n) nogil
+
+
 include "khash_for_primitive_helper.pxi"

pandas/_libs/src/klib/khash_python.h

@@ -244,3 +244,13 @@ void PANDAS_INLINE kh_destroy_str_starts(kh_str_starts_t* table) {
 void PANDAS_INLINE kh_resize_str_starts(kh_str_starts_t* table, khuint_t val) {
 	kh_resize_str(table->table, val);
 }
+
+// utility function: given the number of elements
+// returns number of necessary buckets
+khuint_t PANDAS_INLINE kh_get_bucket_n(khuint_t element_n){
+    khuint_t candidate = element_n;
+    kroundup32(candidate);
+    khuint_t upper_bound = (khuint_t)(candidate * __ac_HASH_UPPER + 0.5);
+    return (upper_bound < element_n) ? 2*candidate : candidate;
+
+}

pandas/tests/libs/test_hashtable.py

@@ -163,6 +163,16 @@ def test_get_state(self, table_type, dtype):
         assert "n_buckets" in state
         assert "upper_bound" in state
 
+    def test_no_reallocation(self, table_type, dtype):
+        N = 110
+        keys = np.arange(N).astype(dtype)
+        table = table_type(N)
+        n_buckets_start = table.get_state()["n_buckets"]
+        table.map_locations(keys)
+        n_buckets_end = table.get_state()["n_buckets"]
+        # orgininal number of buckets was enough:
+        assert n_buckets_start == n_buckets_end
+
 
 def test_get_labels_groupby_for_Int64(writable):
     table = ht.Int64HashTable()
@@ -198,6 +208,17 @@ def test_tracemalloc_for_empty_StringHashTable():
         assert get_allocated_khash_memory() == 0
 
 
+def test_no_reallocation_StringHashTable():
+    N = 110
+    keys = np.arange(N).astype(np.compat.unicode).astype(np.object_)
+    table = ht.StringHashTable(N)
+    n_buckets_start = table.get_state()["n_buckets"]
+    table.map_locations(keys)
+    n_buckets_end = table.get_state()["n_buckets"]
+    # orgininal number of buckets was enough:
+    assert n_buckets_start == n_buckets_end
+
+
 @pytest.mark.parametrize(
     "table_type, dtype",
     [