Add faster unique keys

.gitignore

@@ -7,6 +7,7 @@ __pycache__/
 *.so
 # Generated by Cython
 *.c
+*.html
 
 # Benchmarks
 .benchmarks

searcharray/utils/roaringish.py

@@ -161,9 +161,8 @@ def keys(self, encoded: np.ndarray) -> np.ndarray:
 
     def keys_unique(self, encoded: np.ndarray) -> np.ndarray:
         """Return keys from encoded."""
-        keys = self.keys(encoded)
-        intersected = sorted_unique(keys)
-        return intersected
+        rshift = _64 - self.key_bits
+        return unique(encoded, rshift)
 
     def payload_msb(self, encoded: np.ndarray) -> np.ndarray:
         """Return payload MSBs from encoded."""

searcharray/utils/snp_ops.pyx

@@ -213,21 +213,21 @@ cdef _scan_unique(DTYPE_t[:] arr,
     return result, result_idx
 
 
-cdef _scan_unique_masked(DTYPE_t[:] arr,
-                         DTYPE_t arr_len,
-                         DTYPE_t mask):
+cdef _scan_unique_shifted(DTYPE_t[:] arr,
+                          DTYPE_t arr_len,
+                          DTYPE_t rshift):
     cdef DTYPE_t i = 0
 
     cdef np.uint64_t[:] result = np.empty(arr_len, dtype=np.uint64)
     cdef DTYPE_t result_idx = 0
-    cdef DTYPE_t target = arr[i]
+    cdef DTYPE_t target = arr[i] >> rshift
 
     while i < arr_len:
-        target = arr[i]
+        target = arr[i] >> rshift
         result[result_idx] = target
         result_idx += 1
         i += 1
-        while i < arr_len and (arr[i] & mask) == (target & mask):
+        while i < arr_len and (arr[i] >> rshift) == target:
             i += 1
 
     return result, result_idx
@@ -236,6 +236,9 @@ cdef _scan_unique_masked(DTYPE_t[:] arr,
 
 
 def unique(np.ndarray[DTYPE_t, ndim=1] arr,
-           DTYPE_t mask=ALL_BITS):
-    result,  result_idx = _scan_unique(arr, arr.shape[0])
-    return result[:result_idx]
+           DTYPE_t rshift=0):
+    if rshift > 0:
+        result, result_idx = _scan_unique_shifted(arr, arr.shape[0], rshift)
+    else:
+        result, result_idx = _scan_unique(arr, arr.shape[0])
+    return np.array(result[:result_idx])

test/test_snp_ops.py

@@ -183,6 +183,13 @@ def test_unique(array):
     assert np.all(result == expected)
 
 
+@pytest.mark.parametrize("array,shift", [(u64([0xEE00, 0xFF00, 0xFF01]), 8)])
+def test_unique_shifted(array, shift):
+    expected = np.unique(array >> shift)
+    result = unique(array, shift)
+    assert np.all(result == expected)
+
+
 @pytest.mark.parametrize("seed", [0, 1, 2, 3, 4])
 def test_unique_matches_snp(seed):
     np.random.seed(seed)