introducing Int16HashTable and UInt16HashTable

Author: realead

pandas/_libs/hashtable.pxd

@@ -3,16 +3,20 @@ from numpy cimport intp_t, ndarray
 from pandas._libs.khash cimport (
     float32_t,
     float64_t,
+    int16_t,
     int32_t,
     int64_t,
     kh_float32_t,
     kh_float64_t,
+    kh_int16_t,
     kh_int32_t,
     kh_int64_t,
     kh_pymap_t,
     kh_str_t,
+    kh_uint16_t,
     kh_uint32_t,
     kh_uint64_t,
+    uint16_t,
     uint32_t,
     uint64_t,
 )
@@ -46,6 +50,18 @@ cdef class Int32HashTable(HashTable):
     cpdef get_item(self, int32_t val)
     cpdef set_item(self, int32_t key, Py_ssize_t val)
 
+cdef class UInt16HashTable(HashTable):
+    cdef kh_uint16_t *table
+
+    cpdef get_item(self, uint16_t val)
+    cpdef set_item(self, uint16_t key, Py_ssize_t val)
+
+cdef class Int16HashTable(HashTable):
+    cdef kh_int16_t *table
+
+    cpdef get_item(self, int16_t val)
+    cpdef set_item(self, int16_t key, Py_ssize_t val)
+
 cdef class Float64HashTable(HashTable):
     cdef kh_float64_t *table
 

pandas/_libs/hashtable_class_helper.pxi.in

@@ -10,11 +10,13 @@ WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
 # name
 cimported_types = ['float32',
                    'float64',
+                   'int16',
                    'int32',
                    'int64',
                    'pymap',
                    'str',
                    'strbox',
+                   'uint16',
                    'uint32',
                    'uint64']
 }}
@@ -48,9 +50,11 @@ dtypes = [('Float64', 'float64', 'float64_t'),
           ('Float32', 'float32', 'float32_t'),
           ('Int64', 'int64', 'int64_t'),
           ('Int32', 'int32', 'int32_t'),
+          ('Int16', 'int16', 'int16_t'),
           ('String', 'string', 'char *'),
           ('UInt64', 'uint64', 'uint64_t'),
-          ('UInt32', 'uint32', 'uint32_t')]
+          ('UInt32', 'uint32', 'uint32_t'),
+          ('UInt16', 'uint16', 'uint16_t')]
 }}
 
 {{for name, dtype, c_type in dtypes}}
@@ -78,8 +82,10 @@ cdef inline void append_data_{{dtype}}({{name}}VectorData *data,
 ctypedef fused vector_data:
     Int64VectorData
     Int32VectorData
+    Int16VectorData
     UInt64VectorData
     UInt32VectorData
+    UInt16VectorData
     Float64VectorData
     Float32VectorData
     StringVectorData
@@ -99,7 +105,9 @@ dtypes = [('Float64', 'float64', 'float64_t'),
           ('Int64', 'int64', 'int64_t'),
           ('Float32', 'float32', 'float32_t'),
           ('UInt32', 'uint32', 'uint32_t'),
-          ('Int32', 'int32', 'int32_t')]
+          ('Int32', 'int32', 'int32_t'),
+          ('UInt16', 'uint16', 'uint16_t'),
+          ('Int16', 'int16', 'int16_t')]
 
 }}
 
@@ -293,7 +301,9 @@ dtypes = [('Float64', 'float64', True, 'np.nan'),
           ('Int64', 'int64', False, 'NPY_NAT'),
           ('Float32', 'float32', True, 'np.nan'),
           ('UInt32', 'uint32', False, 0),
-          ('Int32', 'int32', False, 0)]
+          ('Int32', 'int32', False, 0),
+          ('UInt16', 'uint16', False, 0),
+          ('Int16', 'int16', False, 0)]
 
 }}
 

pandas/_libs/hashtable_func_helper.pxi.in

@@ -11,9 +11,11 @@ dtypes = [('float64', 'float64', 'float64_t'),
           ('float32', 'float32', 'float32_t'),
           ('uint64', 'uint64', 'uint64_t'),
           ('uint32', 'uint32', 'uint32_t'),
+          ('uint16', 'uint16', 'uint16_t'),
           ('object', 'pymap', 'object'),
           ('int64', 'int64', 'int64_t'),
-          ('int32', 'int32', 'int32_t')]
+          ('int32', 'int32', 'int32_t'),
+          ('int16', 'int16', 'int16_t')]
 
 }}
 
@@ -281,8 +283,10 @@ dtypes = [('float64', 'float64_t', 'float64', 'float64'),
           ('float32', 'float32_t', 'float32', 'float32'),
           ('int64', 'int64_t', 'int64', 'int64'),
           ('int32', 'int32_t', 'int32', 'int32'),
+          ('int16', 'int16_t', 'int16', 'int16'),
           ('uint64', 'uint64_t', 'uint64', 'uint64'),
           ('uint32', 'uint32_t', 'uint32', 'uint32'),
+          ('uint16', 'uint16_t', 'uint16', 'uint16'),
           ('object', 'object', 'pymap', 'object_')]
 }}
 

pandas/_libs/khash.pxd

@@ -1,5 +1,14 @@
 from cpython.object cimport PyObject
-from numpy cimport float32_t, float64_t, int32_t, int64_t, uint32_t, uint64_t
+from numpy cimport (
+    float32_t,
+    float64_t,
+    int16_t,
+    int32_t,
+    int64_t,
+    uint16_t,
+    uint32_t,
+    uint64_t,
+)
 
 
 cdef extern from "khash_python.h":

pandas/_libs/khash_for_primitive_helper.pxi.in

@@ -13,6 +13,8 @@ primitive_types = [('int64', 'int64_t'),
                    ('int32', 'int32_t'),
                    ('uint32', 'uint32_t'),
                    ('float32', 'float32_t'),
+                   ('int16', 'int16_t'),
+                   ('uint16', 'uint16_t'),
                   ]
 }}
 

pandas/_libs/src/klib/khash.h

@@ -129,6 +129,12 @@ typedef unsigned long long khuint64_t;
 typedef signed long long khint64_t;
 #endif
 
+#if UINT_MAX == 0xffffu
+typedef unsigned int khint16_t;
+#elif USHRT_MAX == 0xffffu
+typedef unsigned short khint16_t;
+#endif
+
 typedef double khfloat64_t;
 typedef double khfloat32_t;
 
@@ -616,6 +622,17 @@ PANDAS_INLINE khint_t __ac_Wang_hash(khint_t key)
 #define KHASH_MAP_INIT_INT64(name, khval_t)								\
 	KHASH_INIT(name, khint64_t, khval_t, 1, kh_int64_hash_func, kh_int64_hash_equal)
 
+/*! @function
+  @abstract     Instantiate a hash map containing 16bit-integer keys
+  @param  name  Name of the hash table [symbol]
+  @param  khval_t  Type of values [type]
+ */
+#define KHASH_MAP_INIT_INT16(name, khval_t)								\
+	KHASH_INIT(name, khint16_t, khval_t, 1, kh_int_hash_func, kh_int_hash_equal)
+
+#define KHASH_MAP_INIT_UINT16(name, khval_t)								\
+	KHASH_INIT(name, khint16_t, khval_t, 1, kh_int_hash_func, kh_int_hash_equal)
+
 
 typedef const char *kh_cstr_t;
 /*! @function
@@ -641,12 +658,16 @@ typedef const char *kh_cstr_t;
 #define kh_exist_float32(h, k) (kh_exist(h, k))
 #define kh_exist_int32(h, k) (kh_exist(h, k))
 #define kh_exist_uint32(h, k) (kh_exist(h, k))
+#define kh_exist_int16(h, k) (kh_exist(h, k))
+#define kh_exist_uint16(h, k) (kh_exist(h, k))
 
 KHASH_MAP_INIT_STR(str, size_t)
 KHASH_MAP_INIT_INT(int32, size_t)
 KHASH_MAP_INIT_UINT(uint32, size_t)
 KHASH_MAP_INIT_INT64(int64, size_t)
 KHASH_MAP_INIT_UINT64(uint64, size_t)
+KHASH_MAP_INIT_INT16(int16, size_t)
+KHASH_MAP_INIT_UINT16(uint16, size_t)
 
 
 #endif /* __AC_KHASH_H */

pandas/tests/libs/test_hashtable.py

@@ -15,6 +15,8 @@
         (ht.Int32HashTable, np.int32),
         (ht.UInt32HashTable, np.uint32),
         (ht.Float32HashTable, np.float32),
+        (ht.Int16HashTable, np.int16),
+        (ht.UInt16HashTable, np.uint16),
     ],
 )
 class TestHashTable:
@@ -153,6 +155,8 @@ def get_ht_function(fun_name, type_suffix):
         (np.int32, "int32"),
         (np.uint32, "uint32"),
         (np.float32, "float32"),
+        (np.int16, "int16"),
+        (np.uint16, "uint16"),
     ],
 )
 class TestHelpFunctions: