Shuffles order around a bit. Replaces memcpy with cast.

Author: chaburkland

performance/reference/util.py

@@ -244,9 +244,13 @@ def get_new_indexers_and_screen_ak(
     ) -> tp.Tuple[np.ndarray, np.ndarray]:
     from arraykit import get_new_indexers_and_screen as ak_routine
 
+    if len(positions) > len(indexers):
+        return np.unique(indexers, return_inverse=True)
+
+    # Will return same *objects* back if it was able to finish early.
     new_indexers, index_screen = ak_routine(indexers, positions)
     if new_indexers is indexers and index_screen is positions:
-        return indexers, positions
+        return positions, indexers
 
     # Use a more helpful alias!
     element_locations = index_screen
@@ -257,4 +261,4 @@ def get_new_indexers_and_screen_ak(
     order_found = np.argsort(found_element_locations)
 
     found_positions = positions[found_mask]
-    return new_indexers, found_positions[order_found]
+    return found_positions[order_found], new_indexers

src/_arraykit.c

@@ -574,37 +574,32 @@ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kw
 
     Equivalent Python code:
 
-        num_unique = len(positions)                                          # 1
-        element_locations = np.full(num_unique, num_unique, dtype=np.int64)  # 2
-        new_indexers = np.empty(len(indexers), dtype=np.int64)               # 3
+        num_unique = len(positions)
+        element_locations = np.full(num_unique, num_unique, dtype=np.int64)
+        new_indexers = np.empty(len(indexers), dtype=np.int64)
 
-        num_found = 0                                                        # 4
+        num_found = 0
 
-        for i, element in enumerate(indexers):                               # 5
-            if element_locations[element] == num_unique:                     # 6
-                element_locations[element] = num_found                       # 7
-                num_found += 1                                               # 8
+        for i, element in enumerate(indexers):
+            if element_locations[element] == num_unique:
+                element_locations[element] = num_found
+                num_found += 1
 
-            if num_found == num_unique:                                      # 9
-                return indexers, positions                                   # 10
+            if num_found == num_unique:
+                return positions, indexers
 
-            new_indexers[i] = element_locations[element]                     # 11
+            new_indexers[i] = element_locations[element]
 
-        return new_indexers, element_locations
-        # This return values will be used in CPython like this
-
-        # The rest of this will be implemented in Python since there is no
-        # benefit to using Cpython for this.
+        return element_locations, new_indexers
+        # ...
+        # NOTE: These return values will be used in a Python wrapper like this:
         found_mask = element_locations != num_unique
 
         found_element_locations = element_locations[found_mask]
         order_found = np.argsort(found_element_locations)
 
         found_positions = positions[found_mask]
-        return found_positions[order_found]
-
-        index_screen = get_index_screen(element_locations, num_unique)
-        return new_indexers, index_screen
+        return found_positions[order_found], new_indexers
     */
     PyArrayObject *indexers;
     PyArrayObject *positions;
@@ -615,7 +610,7 @@ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kw
                 &PyArray_Type, &indexers,
                 &PyArray_Type, &positions
         ))
-{
+    {
         return NULL;
     }
 
@@ -634,7 +629,7 @@ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kw
     if (num_unique > PyArray_SIZE(indexers)) {
         // This algorithm is only optimal if the number of unique elements is
         // less than the number of elements in the indexers.
-        // Otherwise, the most optimal code is ``np.unique(indexers, return_index=True)[1]``
+        // Otherwise, the most optimal code is ``np.unique(indexers, return_index=True)``
         // and we don't want to re-implement that in C.
         PyErr_SetString(
                 PyExc_ValueError,
@@ -691,11 +686,11 @@ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kw
     // C-contiguous, F-contiguous, both, or neither.
     // See https://numpy.org/doc/stable/reference/c-api/iterator.html#simple-iteration-example
     NpyIter *indexer_iter = NpyIter_New(
-            indexers,
-            NPY_ITER_READONLY| NPY_ITER_EXTERNAL_LOOP,
-            NPY_KEEPORDER,
-            NPY_NO_CASTING,
-            NULL
+            indexers,                                   // array
+            NPY_ITER_READONLY | NPY_ITER_EXTERNAL_LOOP, // iter flags
+            NPY_KEEPORDER,                              // order
+            NPY_NO_CASTING,                             // casting
+            NULL                                        // dtype
             );
     if (indexer_iter == NULL) {
         Py_DECREF(element_locations);
@@ -727,7 +722,7 @@ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kw
         npy_int64 element;
 
         while (inner_size--) {
-            memcpy (&element, data, sizeof (npy_int64));
+            element = *((npy_int64 *)data);
 
             if (element_location_values[element] == num_unique) {
                 element_location_values[element] = num_found;

test/test_util.py

@@ -16,7 +16,8 @@
 from arraykit import array_deepcopy
 from arraykit import isna_element
 from arraykit import dtype_from_element
-from performance.reference.util import get_new_indexers_and_screen_ak
+from performance.reference.util import get_new_indexers_and_screen_ak as get_new_indexers_and_screen_full
+from arraykit import get_new_indexers_and_screen
 
 from performance.reference.util import mloc as mloc_ref
 
@@ -384,50 +385,50 @@ def test_dtype_from_element_str_and_bytes_dtypes(self) -> None:
 
     def test_get_new_indexers_and_screen_a(self) -> None:
         indexersA = np.array([9, 9, 9, 9, 0, 0, 1, 4, 5, 0, 0, 0, 1], dtype=np.int64)
-        postA = get_new_indexers_and_screen_ak(indexersA, np.arange(10, dtype=np.int64))
+        postA = get_new_indexers_and_screen_full(indexersA, np.arange(10, dtype=np.int64))
         assert indexersA.flags.c_contiguous
         assert indexersA.flags.f_contiguous
         assert tuple(map(list, postA)) == (
-            [0, 0, 0, 0, 1, 1, 2, 3, 4, 1, 1, 1, 2],
             [9, 0, 1, 4, 5],
+            [0, 0, 0, 0, 1, 1, 2, 3, 4, 1, 1, 1, 2],
         )
 
         # Prove we can handle non-continuous arrays
         indexersB = np.full((len(indexersA), 3), -1, dtype=np.int64)
         indexersB[:,1] = indexersA.copy()
         assert not indexersB[:,1].flags.c_contiguous
         assert not indexersB[:,1].flags.f_contiguous
-        postB = get_new_indexers_and_screen_ak(indexersB[:,1], np.arange(10, dtype=np.int64))
+        postB = get_new_indexers_and_screen_full(indexersB[:,1], np.arange(10, dtype=np.int64))
         assert tuple(map(list, postA)) == tuple(map(list, postB))
 
         indexersC = np.array([9, 9, 9, 9, 0, 0, 1, 4, 5, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=np.int64)
-        postC = get_new_indexers_and_screen_ak(indexersC, positions=np.arange(15, dtype=np.int64))
+        postC = get_new_indexers_and_screen_full(indexersC, positions=np.arange(15, dtype=np.int64))
         assert tuple(map(list, postC)) == (
-            [0, 0, 0, 0, 1, 1, 2, 3, 4, 1, 1, 1, 2, 5, 6, 3, 4,7, 8, 9, 0, 10],
             [9, 0, 1, 4, 5, 2, 3, 6, 7, 8, 10],
+            [0, 0, 0, 0, 1, 1, 2, 3, 4, 1, 1, 1, 2, 5, 6, 3, 4,7, 8, 9, 0, 10],
         )
 
         indexersD = np.array([2, 1, 0, 2, 0, 1, 1, 2, 0], dtype=np.int64)
-        postD = get_new_indexers_and_screen_ak(indexers=indexersD, positions=np.arange(3, dtype=np.int64))
+        postD = get_new_indexers_and_screen_full(indexers=indexersD, positions=np.arange(3, dtype=np.int64))
         assert tuple(map(list, postD)) == (
-            [2, 1, 0, 2, 0, 1, 1, 2, 0],
             [0, 1, 2],
+            [2, 1, 0, 2, 0, 1, 1, 2, 0],
         )
 
     def test_get_new_indexers_and_screen_b(self) -> None:
         indexersA = np.array([5], dtype=np.int64)
 
         with self.assertRaises(ValueError):
-            get_new_indexers_and_screen_ak(indexersA, np.arange(6, dtype=np.int64))
+            get_new_indexers_and_screen(indexersA, np.arange(6, dtype=np.int64))
 
         with self.assertRaises(ValueError):
-            get_new_indexers_and_screen_ak(indexersA, np.arange(106, dtype=np.int64))
+            get_new_indexers_and_screen(indexersA, np.arange(106, dtype=np.int64))
 
         with self.assertRaises(ValueError):
-            get_new_indexers_and_screen_ak(indexersA.astype(np.int32), np.arange(5))
+            get_new_indexers_and_screen(indexersA.astype(np.int32), np.arange(5))
 
         indexersB = np.arange(25, dtype=np.int64)
-        postB = get_new_indexers_and_screen_ak(indexersB, indexersB)
+        postB = get_new_indexers_and_screen(indexersB, indexersB)
         assert tuple(map(list, postB)) == (list(indexersB), list(indexersB))