Enhancement: Add a custom initial allocation size for typed dicts.

See numba#8110

numba/typed/dictobject.py

@@ -69,8 +69,20 @@ class Status(IntEnum):
     ERR_CMP_FAILED = -5
 
 
-def new_dict(key, value):
-    """Construct a new dict.
+def new_dict_minsize(key, value):
+    """Construct a new dict with minimum size.
+
+    Parameters
+    ----------
+    key, value : TypeRef
+        Key type and value type of the new dict.
+    """
+    # With JIT disabled, ignore all arguments and return a Python dict.
+    return dict()
+
+
+def new_dict(key, value, allocated):
+    """Construct a new dict with custom size.
 
     Parameters
     ----------
@@ -253,6 +265,48 @@ def codegen(context, builder, sig, args):
     return sig, codegen
 
 
+@intrinsic
+def _dict_new(typingctx, keyty, valty, allocated):
+    """Wrap numba_dict_new.
+
+    Allocate a new dictionary object of size 'allocated'.
+
+    Parameters
+    ----------
+    keyty, valty: Type
+        Type of the key and value, respectively.
+    allocated: int
+        Number of buckets to allocate.
+    """
+    resty = types.voidptr
+    sig = resty(keyty, valty, allocated)
+
+    def codegen(context, builder, sig, args):
+        fnty = ir.FunctionType(
+            ll_status,
+            [ll_dict_type.as_pointer(), ll_ssize_t, ll_ssize_t, ll_ssize_t],
+        )
+        fn = ir.Function(builder.module, fnty, 'numba_dict_new')
+        # Determine sizeof key and value types
+        ll_key = context.get_data_type(keyty.instance_type)
+        ll_val = context.get_data_type(valty.instance_type)
+        sz_key = context.get_abi_sizeof(ll_key)
+        sz_val = context.get_abi_sizeof(ll_val)
+        refdp = cgutils.alloca_once(builder, ll_dict_type, zfill=True)
+        status = builder.call(
+            fn,
+            [refdp, args[2], ll_ssize_t(sz_key), ll_ssize_t(sz_val)],
+        )
+        _raise_if_error(
+            context, builder, status,
+            msg="Failed to allocate dictionary",
+        )
+        dp = builder.load(refdp)
+        return dp
+
+    return sig, codegen
+
+
 @intrinsic
 def _dict_set_method_table(typingctx, dp, keyty, valty):
     """Wrap numba_dict_set_method_table
@@ -643,10 +697,11 @@ def codegen(context, builder, signature, args):
     return sig, codegen
 
 
-@overload(new_dict)
-def impl_new_dict(key, value):
+@overload(new_dict_minsize)
+def impl_new_dict_minsize(key, value):
     """Creates a new dictionary with *key* and *value* as the type
-    of the dictionary key and value, respectively.
+    of the dictionary key and value, respectively. The number of buckets
+    initialized is defined in /cext/dictobject.c as D_MINSIZE.
     """
     if any([
         not isinstance(key, Type),
@@ -665,6 +720,35 @@ def imp(key, value):
     return imp
 
 
+@overload(new_dict)
+def impl_new_dict(key, value, allocated): #allocated=None?
+    """Creates a new dictionary with *key* and *value* as the type
+    of the dictionary key and value, respectively. Allocated is the
+    number of buckets initialized.
+    """
+    if any([
+        not isinstance(key, Type),
+        not isinstance(value, Type),
+    ]):
+        raise TypeError("expecting *key* and *value* to be a numba Type")
+
+    keyty, valty = key, value
+
+    def imp(key, value, allocated): #allocated=None?
+        # !!! Round allocated up (down?) to the nearest power of 2
+        # Or fail if not power of 2?
+        # Power of 2 checks earlier or now?
+        if allocated < 0:
+            # Does this check even work?
+            raise RuntimeError("expecting *allocated* to be >= 0")
+        dp = _dict_new(keyty, valty, allocated)
+        _dict_set_method_table(dp, keyty, valty)
+        d = _make_dict(keyty, valty, dp)
+        return d
+
+    return imp
+
+
 @overload(len)
 def impl_len(d):
     """len(dict)

numba/typed/typeddict.py

@@ -19,8 +19,12 @@
 
 
 @njit
-def _make_dict(keyty, valty):
-    return dictobject._as_meminfo(dictobject.new_dict(keyty, valty))
+def _make_dict(keyty, valty, allocated=None):
+    if allocated is None:
+        return dictobject._as_meminfo(dictobject.new_dict_minsize(keyty, valty))
+    else:
+        return dictobject._as_meminfo(dictobject.new_dict(keyty, valty, 
+                                                          allocated=allocated))
 
 
 @njit
@@ -84,21 +88,22 @@ class Dict(MutableMapping):
     Implements the MutableMapping interface.
     """
 
-    def __new__(cls, dcttype=None, meminfo=None):
+    def __new__(cls, dcttype=None, meminfo=None, allocated=None):
         if config.DISABLE_JIT:
             return dict.__new__(dict)
         else:
             return object.__new__(cls)
 
     @classmethod
-    def empty(cls, key_type, value_type):
+    def empty(cls, key_type, value_type, allocated=None):
         """Create a new empty Dict with *key_type* and *value_type*
         as the types for the keys and values of the dictionary respectively.
         """
         if config.DISABLE_JIT:
             return dict()
         else:
-            return cls(dcttype=DictType(key_type, value_type))
+            return cls(dcttype=DictType(key_type, value_type), 
+                       allocated=allocated)
 
     def __init__(self, **kwargs):
         """
@@ -117,14 +122,15 @@ def __init__(self, **kwargs):
         else:
             self._dict_type = None
 
-    def _parse_arg(self, dcttype, meminfo=None):
+    def _parse_arg(self, dcttype, meminfo=None, allocated=None):
         if not isinstance(dcttype, DictType):
             raise TypeError('*dcttype* must be a DictType')
 
         if meminfo is not None:
             opaque = meminfo
         else:
-            opaque = _make_dict(dcttype.key_type, dcttype.value_type)
+            opaque = _make_dict(dcttype.key_type, dcttype.value_type, 
+                                allocated=allocated)
         return dcttype, opaque
 
     @property
@@ -209,12 +215,16 @@ def copy(self):
 
 
 @overload_classmethod(types.DictType, 'empty')
-def typeddict_empty(cls, key_type, value_type):
+def typeddict_empty(cls, key_type, value_type, allocated=None):
     if cls.instance_type is not DictType:
         return
-
-    def impl(cls, key_type, value_type):
-        return dictobject.new_dict(key_type, value_type)
+
+    def impl(cls, key_type, value_type, allocated=None):
+        if allocated is None:
+            return dictobject.new_dict_minsize(key_type, value_type)
+        else:
+            return dictobject.new_dict(key_type, value_type, 
+                                       allocated=allocated)
 
     return impl