diff --git a/bquery/ctable_ext.pyx b/bquery/ctable_ext.pyx index 38666f8..2273915 100644 --- a/bquery/ctable_ext.pyx +++ b/bquery/ctable_ext.pyx @@ -1,6 +1,6 @@ import numpy as np from numpy cimport ndarray, dtype, npy_intp, npy_int32, \ - npy_uint64, npy_int64, npy_float64, npy_bool + npy_uint64, npy_int64, npy_float64, npy_bool, uint64_t import cython import bcolz as bz @@ -130,13 +130,11 @@ def factorize_str(carray carray_, carray labels=None): @cython.wraparound(False) @cython.boundscheck(False) cdef void _factorize_int64_helper(Py_ssize_t iter_range, - Py_ssize_t allocation_size, - ndarray[npy_int64] in_buffer, - ndarray[npy_uint64] out_buffer, + npy_int64[:] in_buffer, + uint64_t[:] out_buffer, kh_int64_t *table, Py_ssize_t * count, - dict reverse, - ): + ) nogil: cdef: Py_ssize_t i, idx int ret @@ -154,7 +152,6 @@ cdef void _factorize_int64_helper(Py_ssize_t iter_range, else: k = kh_put_int64(table, element, &ret) table.vals[k] = idx = count[0] - reverse[count[0]] = element count[0] += 1 out_buffer[i] = idx @@ -168,6 +165,8 @@ def factorize_int64(carray carray_, carray labels=None): ndarray[npy_int64] in_buffer ndarray[npy_uint64] out_buffer kh_int64_t *table + npy_int64[:] in_buffer_view + uint64_t[:] out_buffer_view count = 0 ret = 0 @@ -179,6 +178,7 @@ def factorize_int64(carray carray_, carray labels=None): labels = carray([], dtype='int64', expectedlen=n) # in-buffer isn't typed, because cython doesn't support string arrays (?) out_buffer = np.empty(chunklen, dtype='uint64') + out_buffer_view = out_buffer in_buffer = np.empty(chunklen, dtype='int64') table = kh_init_int64() @@ -186,31 +186,35 @@ def factorize_int64(carray carray_, carray labels=None): chunk_ = carray_.chunks[i] # decompress into in_buffer chunk_._getitem(0, chunklen, in_buffer.data) - _factorize_int64_helper(chunklen, - carray_.dtype.itemsize + 1, - in_buffer, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = in_buffer + with nogil: + _factorize_int64_helper(chunklen, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer.astype(np.int64)) leftover_elements = cython.cdiv(carray_.leftover, carray_.atomsize) if leftover_elements > 0: - _factorize_int64_helper(leftover_elements, - carray_.dtype.itemsize + 1, - carray_.leftover_array, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = carray_.leftover_array + with nogil: + _factorize_int64_helper(leftover_elements, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer[:leftover_elements].astype(np.int64)) + for i in range(table.n_buckets): + if not kh_exist_int64(table, i): # adjust function name to hash-table data-type + continue + reverse[table.vals[i]] = table.keys[i] kh_destroy_int64(table) return labels, reverse @@ -218,13 +222,11 @@ def factorize_int64(carray carray_, carray labels=None): @cython.wraparound(False) @cython.boundscheck(False) cdef void _factorize_int32_helper(Py_ssize_t iter_range, - Py_ssize_t allocation_size, - ndarray[npy_int32] in_buffer, - ndarray[npy_uint64] out_buffer, + npy_int32[:] in_buffer, + uint64_t[:] out_buffer, kh_int32_t *table, Py_ssize_t * count, - dict reverse, - ): + ) nogil: cdef: Py_ssize_t i, idx int ret @@ -242,7 +244,6 @@ cdef void _factorize_int32_helper(Py_ssize_t iter_range, else: k = kh_put_int32(table, element, &ret) table.vals[k] = idx = count[0] - reverse[count[0]] = element count[0] += 1 out_buffer[i] = idx @@ -256,6 +257,8 @@ def factorize_int32(carray carray_, carray labels=None): ndarray[npy_int32] in_buffer ndarray[npy_uint64] out_buffer kh_int32_t *table + npy_int32[:] in_buffer_view + uint64_t[:] out_buffer_view count = 0 ret = 0 @@ -267,6 +270,7 @@ def factorize_int32(carray carray_, carray labels=None): labels = carray([], dtype='int64', expectedlen=n) # in-buffer isn't typed, because cython doesn't support string arrays (?) out_buffer = np.empty(chunklen, dtype='uint64') + out_buffer_view = out_buffer in_buffer = np.empty(chunklen, dtype='int32') table = kh_init_int32() @@ -274,31 +278,35 @@ def factorize_int32(carray carray_, carray labels=None): chunk_ = carray_.chunks[i] # decompress into in_buffer chunk_._getitem(0, chunklen, in_buffer.data) - _factorize_int32_helper(chunklen, - carray_.dtype.itemsize + 1, - in_buffer, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = in_buffer + with nogil: + _factorize_int32_helper(chunklen, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer.astype(np.int64)) leftover_elements = cython.cdiv(carray_.leftover, carray_.atomsize) if leftover_elements > 0: - _factorize_int32_helper(leftover_elements, - carray_.dtype.itemsize + 1, - carray_.leftover_array, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = carray_.leftover_array + with nogil: + _factorize_int32_helper(leftover_elements, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer[:leftover_elements].astype(np.int64)) + for i in range(table.n_buckets): + if not kh_exist_int32(table, i): # adjust function name to hash-table data-type + continue + reverse[table.vals[i]] = table.keys[i] kh_destroy_int32(table) return labels, reverse @@ -306,13 +314,11 @@ def factorize_int32(carray carray_, carray labels=None): @cython.wraparound(False) @cython.boundscheck(False) cdef void _factorize_float64_helper(Py_ssize_t iter_range, - Py_ssize_t allocation_size, - ndarray[npy_float64] in_buffer, - ndarray[npy_uint64] out_buffer, + npy_float64[:] in_buffer, + uint64_t[:] out_buffer, kh_float64_t *table, Py_ssize_t * count, - dict reverse, - ): + ) nogil: cdef: Py_ssize_t i, idx int ret @@ -330,7 +336,6 @@ cdef void _factorize_float64_helper(Py_ssize_t iter_range, else: k = kh_put_float64(table, element, &ret) table.vals[k] = idx = count[0] - reverse[count[0]] = element count[0] += 1 out_buffer[i] = idx @@ -344,6 +349,8 @@ def factorize_float64(carray carray_, carray labels=None): ndarray[npy_float64] in_buffer ndarray[npy_uint64] out_buffer kh_float64_t *table + npy_float64[:] in_buffer_view + uint64_t[:] out_buffer_view count = 0 ret = 0 @@ -355,6 +362,7 @@ def factorize_float64(carray carray_, carray labels=None): labels = carray([], dtype='int64', expectedlen=n) # in-buffer isn't typed, because cython doesn't support string arrays (?) out_buffer = np.empty(chunklen, dtype='uint64') + out_buffer_view = out_buffer in_buffer = np.empty(chunklen, dtype='float64') table = kh_init_float64() @@ -362,31 +370,35 @@ def factorize_float64(carray carray_, carray labels=None): chunk_ = carray_.chunks[i] # decompress into in_buffer chunk_._getitem(0, chunklen, in_buffer.data) - _factorize_float64_helper(chunklen, - carray_.dtype.itemsize + 1, - in_buffer, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = in_buffer + with nogil: + _factorize_float64_helper(chunklen, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer.astype(np.int64)) leftover_elements = cython.cdiv(carray_.leftover, carray_.atomsize) if leftover_elements > 0: - _factorize_float64_helper(leftover_elements, - carray_.dtype.itemsize + 1, - carray_.leftover_array, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = carray_.leftover_array + with nogil: + _factorize_float64_helper(leftover_elements, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer[:leftover_elements].astype(np.int64)) + for i in range(table.n_buckets): + if not kh_exist_float64(table, i): # adjust function name to hash-table data-type + continue + reverse[table.vals[i]] = table.keys[i] kh_destroy_float64(table) return labels, reverse diff --git a/bquery/khash.pxd b/bquery/khash.pxd index a8fd51a..de7b6a1 100644 --- a/bquery/khash.pxd +++ b/bquery/khash.pxd @@ -65,9 +65,9 @@ cdef extern from "khash_python.h": inline kh_int64_t* kh_init_int64() inline void kh_destroy_int64(kh_int64_t*) inline void kh_clear_int64(kh_int64_t*) - inline khint_t kh_get_int64(kh_int64_t*, int64_t) + inline khint_t kh_get_int64(kh_int64_t*, int64_t) nogil inline void kh_resize_int64(kh_int64_t*, khint_t) - inline khint_t kh_put_int64(kh_int64_t*, int64_t, int*) + inline khint_t kh_put_int64(kh_int64_t*, int64_t, int*) nogil inline void kh_del_int64(kh_int64_t*, khint_t) bint kh_exist_int64(kh_int64_t*, khiter_t) @@ -81,9 +81,9 @@ cdef extern from "khash_python.h": inline kh_float64_t* kh_init_float64() inline void kh_destroy_float64(kh_float64_t*) inline void kh_clear_float64(kh_float64_t*) - inline khint_t kh_get_float64(kh_float64_t*, float64_t) + inline khint_t kh_get_float64(kh_float64_t*, float64_t) nogil inline void kh_resize_float64(kh_float64_t*, khint_t) - inline khint_t kh_put_float64(kh_float64_t*, float64_t, int*) + inline khint_t kh_put_float64(kh_float64_t*, float64_t, int*) nogil inline void kh_del_float64(kh_float64_t*, khint_t) bint kh_exist_float64(kh_float64_t*, khiter_t) @@ -97,9 +97,9 @@ cdef extern from "khash_python.h": inline kh_int32_t* kh_init_int32() inline void kh_destroy_int32(kh_int32_t*) inline void kh_clear_int32(kh_int32_t*) - inline khint_t kh_get_int32(kh_int32_t*, int32_t) + inline khint_t kh_get_int32(kh_int32_t*, int32_t) nogil inline void kh_resize_int32(kh_int32_t*, khint_t) - inline khint_t kh_put_int32(kh_int32_t*, int32_t, int*) + inline khint_t kh_put_int32(kh_int32_t*, int32_t, int*) nogil inline void kh_del_int32(kh_int32_t*, khint_t) bint kh_exist_int32(kh_int32_t*, khiter_t) diff --git a/bquery/templates/ctable_ext.template.pyx b/bquery/templates/ctable_ext.template.pyx index 51e9e68..5c37688 100644 --- a/bquery/templates/ctable_ext.template.pyx +++ b/bquery/templates/ctable_ext.template.pyx @@ -1,6 +1,6 @@ import numpy as np from numpy cimport ndarray, dtype, npy_intp, npy_int32, \ - npy_uint64, npy_int64, npy_float64, npy_bool + npy_uint64, npy_int64, npy_float64, npy_bool, uint64_t import cython import bcolz as bz @@ -131,13 +131,11 @@ def factorize_str(carray carray_, carray labels=None): @cython.wraparound(False) @cython.boundscheck(False) cdef void _factorize_{{ factor_type }}_helper(Py_ssize_t iter_range, - Py_ssize_t allocation_size, - ndarray[npy_{{ factor_type }}] in_buffer, - ndarray[npy_uint64] out_buffer, + npy_{{ factor_type }}[:] in_buffer, + uint64_t[:] out_buffer, kh_{{ factor_type }}_t *table, Py_ssize_t * count, - dict reverse, - ): + ) nogil: cdef: Py_ssize_t i, idx int ret @@ -155,7 +153,6 @@ cdef void _factorize_{{ factor_type }}_helper(Py_ssize_t iter_range, else: k = kh_put_{{ factor_type }}(table, element, &ret) table.vals[k] = idx = count[0] - reverse[count[0]] = element count[0] += 1 out_buffer[i] = idx @@ -169,6 +166,8 @@ def factorize_{{ factor_type }}(carray carray_, carray labels=None): ndarray[npy_{{ factor_type }}] in_buffer ndarray[npy_uint64] out_buffer kh_{{ factor_type }}_t *table + npy_{{ factor_type }}[:] in_buffer_view + uint64_t[:] out_buffer_view count = 0 ret = 0 @@ -180,6 +179,7 @@ def factorize_{{ factor_type }}(carray carray_, carray labels=None): labels = carray([], dtype='int64', expectedlen=n) # in-buffer isn't typed, because cython doesn't support string arrays (?) out_buffer = np.empty(chunklen, dtype='uint64') + out_buffer_view = out_buffer in_buffer = np.empty(chunklen, dtype='{{ factor_type }}') table = kh_init_{{ factor_type }}() @@ -187,31 +187,35 @@ def factorize_{{ factor_type }}(carray carray_, carray labels=None): chunk_ = carray_.chunks[i] # decompress into in_buffer chunk_._getitem(0, chunklen, in_buffer.data) - _factorize_{{ factor_type }}_helper(chunklen, - carray_.dtype.itemsize + 1, - in_buffer, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = in_buffer + with nogil: + _factorize_{{ factor_type }}_helper(chunklen, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer.astype(np.int64)) leftover_elements = cython.cdiv(carray_.leftover, carray_.atomsize) if leftover_elements > 0: - _factorize_{{ factor_type }}_helper(leftover_elements, - carray_.dtype.itemsize + 1, - carray_.leftover_array, - out_buffer, - table, - &count, - reverse, - ) + in_buffer_view = carray_.leftover_array + with nogil: + _factorize_{{ factor_type }}_helper(leftover_elements, + in_buffer_view, + out_buffer_view, + table, + &count + ) # compress out_buffer into labels labels.append(out_buffer[:leftover_elements].astype(np.int64)) + for i in range(table.n_buckets): + if not kh_exist_{{ factor_type }}(table, i): # adjust function name to hash-table data-type + continue + reverse[table.vals[i]] = table.keys[i] kh_destroy_{{ factor_type }}(table) return labels, reverse