Nogil merge

searcharray/phrase/bigram_freqs.py

@@ -111,6 +111,8 @@ def _inner_bigram_freqs(lhs: np.ndarray,
     -----------
     lhs: roaringish encoded posn array of left term
     rhs: roaringish encoded posn array of right term
+    lhs_int: np.ndarray, intersected lhs
+    rhs_int: np.ndarray, intersected rhs
     phrase_freqs: np.ndarray, preallocated phrase freqs for output
     cont_rhs: bool, whether to continue matching on the rhs or lhs
 
@@ -187,26 +189,26 @@ def _adjacent_bigram_freqs(lhs: np.ndarray, rhs: np.ndarray,
     return phrase_freqs, (lhs_next, rhs_next)
 
 
-def _set_adjbit_at_header(enc1: np.ndarray, just_lsb_enc: np.ndarray,
+def _set_adjbit_at_header(next_inner: np.ndarray, next_adj: np.ndarray,
                           cont: Continuation = Continuation.RHS) -> np.ndarray:
     """Merge two encoded arrays on their headers."""
-    if len(enc1) == 0:
-        return just_lsb_enc
-    if len(just_lsb_enc) == 0:
-        return enc1
+    if len(next_inner) == 0:
+        return next_adj
+    if len(next_adj) == 0:
+        return next_inner
 
-    same_header_enc1, same_header_just_lsb = intersect(enc1, just_lsb_enc,
-                                                       mask=encoder.header_mask)
+    same_header_inner, same_header_adj = intersect(next_inner, next_adj,
+                                                   mask=encoder.header_mask)
     # Set _1 on intersection
-    ignore_mask = np.ones(len(just_lsb_enc), dtype=bool)
-    ignore_mask[same_header_just_lsb] = False
-    if len(same_header_enc1) > 0 and cont == Continuation.RHS:
-        enc1[same_header_enc1] |= _1
-        just_lsb_enc = just_lsb_enc[ignore_mask]
-    if len(same_header_just_lsb) > 0 and cont == Continuation.LHS:
-        enc1[same_header_enc1] |= _upper_bit
-        just_lsb_enc = just_lsb_enc[ignore_mask]
-    return merge(enc1, just_lsb_enc)
+    ignore_mask = np.ones(len(next_adj), dtype=bool)
+    ignore_mask[same_header_adj] = False
+    if len(same_header_inner) > 0 and cont == Continuation.RHS:
+        next_inner[same_header_inner] |= _1
+        next_adj = next_adj[ignore_mask]
+    if len(same_header_adj) > 0 and cont == Continuation.LHS:
+        next_inner[same_header_adj] |= _upper_bit
+        next_adj = next_adj[ignore_mask]
+    return merge(next_inner, next_adj)
 
 
 def bigram_freqs(lhs: np.ndarray,

searcharray/roaringish/merge.pyx

@@ -52,19 +52,16 @@ cdef _merge_naive(DTYPE_t[:] lhs,
     return np.asarray(results), result_idx
 
 
-cdef _merge(DTYPE_t[:] lhs,
-            DTYPE_t[:] rhs):
-    cdef np.intp_t len_lhs = lhs.shape[0]
-    cdef np.intp_t len_rhs = rhs.shape[0]
-
+cdef DTYPE_t _merge(DTYPE_t* lhs,
+                    DTYPE_t* rhs,
+                    DTYPE_t len_lhs,
+                    DTYPE_t len_rhs,
+                    DTYPE_t* results) nogil:
     cdef DTYPE_t* lhs_ptr = &lhs[0]
     cdef DTYPE_t* end_lhs_ptr = &lhs[len_lhs]
     cdef DTYPE_t* rhs_ptr = &rhs[0]
     cdef DTYPE_t* end_rhs_ptr = &rhs[len_rhs]
-
-    # Outputs as numpy arrays
-    cdef np.uint64_t[:] results = np.empty(len_lhs + len_rhs, dtype=np.uint64)
-    cdef np.uint64_t* result_ptr = &results[0]
+    cdef DTYPE_t* result_ptr = &results[0]
 
     # Copy elements from both arrays
     while lhs_ptr < end_lhs_ptr and rhs_ptr < end_rhs_ptr:
@@ -93,22 +90,20 @@ cdef _merge(DTYPE_t[:] lhs,
         lhs_ptr += 1
         result_ptr += 1
 
-    return np.asarray(results), result_ptr - &results[0]
-
+    return result_ptr - &results[0]
 
-cdef _merge_w_drop(DTYPE_t[:] lhs,
-                   DTYPE_t[:] rhs):
-    cdef np.intp_t len_lhs = lhs.shape[0]
-    cdef np.intp_t len_rhs = rhs.shape[0]
 
+cdef DTYPE_t _merge_w_drop(DTYPE_t* lhs,
+                           DTYPE_t* rhs,
+                           DTYPE_t len_lhs,
+                           DTYPE_t len_rhs,
+                           DTYPE_t* results) nogil:
     cdef DTYPE_t* lhs_ptr = &lhs[0]
     cdef DTYPE_t* end_lhs_ptr = &lhs[len_lhs]
     cdef DTYPE_t* rhs_ptr = &rhs[0]
     cdef DTYPE_t* end_rhs_ptr = &rhs[len_rhs]
 
-    # Outputs as numpy arrays
-    cdef np.uint64_t[:] results = np.empty(len_lhs + len_rhs, dtype=np.uint64)
-    cdef np.uint64_t* result_ptr = &results[0]
+    cdef DTYPE_t* result_ptr = &results[0]
 
     # Copy elements from both arrays
     while lhs_ptr < end_lhs_ptr and rhs_ptr < end_rhs_ptr:
@@ -135,16 +130,31 @@ cdef _merge_w_drop(DTYPE_t[:] lhs,
         lhs_ptr += 1
         result_ptr += 1
 
-    return np.asarray(results), result_ptr - &results[0]
-
+    return result_ptr - &results[0]
 
 
 
 def merge(np.ndarray[DTYPE_t, ndim=1] lhs,
           np.ndarray[DTYPE_t, ndim=1] rhs,
           bint drop_duplicates=False):
-    if drop_duplicates:
-        result, result_idx = _merge_w_drop(lhs, rhs)
-    else:
-        result, result_idx = _merge(lhs, rhs)
-    return np.array(result[:result_idx])
+    # Outputs as numpy arrays
+    cdef DTYPE_t result_idx
+
+    cdef DTYPE_t* lhs_ptr = &lhs[0]
+    cdef DTYPE_t* rhs_ptr = &rhs[0]
+    cdef DTYPE_t lhs_len = lhs.shape[0]
+    cdef DTYPE_t rhs_len = rhs.shape[0]
+    cdef DTYPE_t[:] results = np.empty(lhs.shape[0] + rhs.shape[0], dtype=np.uint64)
+    cdef DTYPE_t* result_ptr = &results[0]
+
+    with nogil:
+        if drop_duplicates:
+            result_idx = _merge_w_drop(lhs_ptr, rhs_ptr,
+                                       lhs_len, rhs_len,
+                                       result_ptr)
+        else:
+            result_idx = _merge(lhs_ptr, rhs_ptr,
+                                lhs_len, rhs_len,
+                                result_ptr)
+
+    return np.array(results[:result_idx])