Closes Bears-R-Us#3009: indexof1d to handle null values (Bears-R-Us…

…#3169)

* Closes Bears-R-Us#3009: `indexof1d` to handle null values

This PR (closes Bears-R-Us#3009) refactors `indexof1d` to use `find` since they have similar functionality and `find` is fairly optimized and correctly handles null values (once we `dropna=False` to the `Groupby`). The two major difference is when there are how missing values are handled and how many indices get returned when there are duplicates in the search space. `find` would only return the index of the first occurrence and use `-1` to denote missing values, but `indexof1d` returns the indices of all occurrences and removes missing values. To enable this, I added the flags `all_occurrences` and `remove_missing` to `find`

The approach I took involved adding a segmented `mink/maxk`, which I went back and forth on whether it should be user facing. I implemented this by permuting the values and calling the existing `mink/maxk`. I'm not sure if this is the most efficient approach, but my goal was focus on correctness first and we can optimize later if needed.

Wrote tests for `indexof1d` both for the reproducer and in general.

* update and add examples in response to PR feedback

---------

Co-authored-by: Tess Hayes <[email protected]>

PROTO_tests/tests/setops_test.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 import pytest
 
 import arkouda as ak
@@ -50,8 +51,8 @@ def make_np_arrays_cross_type(dtype1, dtype2):
             a = np.array([-1, -3, 0, 1, 2, 3]).astype(dtype1)
             c = np.array([-1, 0, 0, 7, 8, 3]).astype(dtype1)
         elif dtype1 == ak.bigint:
-            a = np.array([-1, -3, 0, 1, 2, 3]).astype(ak.uint64) + 2 ** 200
-            c = np.array([-1, 0, 0, 7, 8, 3]).astype(ak.uint64) + 2 ** 200
+            a = np.array([-1, -3, 0, 1, 2, 3]).astype(ak.uint64) + 2**200
+            c = np.array([-1, 0, 0, 7, 8, 3]).astype(ak.uint64) + 2**200
         elif dtype1 == ak.bool:
             a = np.array([True, False, False, True, True])
             c = np.array([True, True, False, False, True])
@@ -62,8 +63,8 @@ def make_np_arrays_cross_type(dtype1, dtype2):
             b = np.array([-1, -11, 0, 4, 5, 3]).astype(dtype2)
             d = np.array([-1, -4, 0, 7, 8, 3]).astype(dtype2)
         elif dtype2 == ak.bigint:
-            b = np.array([-1, -11, 0, 4, 5, 3]).astype(ak.uint64) + 2 ** 200
-            d = np.array([-1, -4, 0, 7, 8, 3]).astype(ak.uint64) + 2 ** 200
+            b = np.array([-1, -11, 0, 4, 5, 3]).astype(ak.uint64) + 2**200
+            d = np.array([-1, -4, 0, 7, 8, 3]).astype(ak.uint64) + 2**200
         elif dtype2 == ak.bool:
             b = np.array([True, True, False, False, True])
             d = np.array([True, True, False, False, True])
@@ -674,3 +675,47 @@ def test_multiarray_validation(self):
         x = [ak.arange(3, dtype=ak.uint64), ak.arange(3)]
         with pytest.raises(TypeError):
             ak.pdarraysetops.multiarray_setop_validation(x, y)
+
+    def test_index_of(self):
+        # index of nan (reproducer from #3009)
+        s = ak.Series(ak.array([1, 2, 3]), index=ak.array([1, 2, np.nan]))
+        assert ak.indexof1d(ak.array([np.nan]), s.index.values).to_list() == [2]
+
+        select_from_list = [
+            ak.randint(-(2**32), 2**32, 10),
+            ak.linspace(-(2**32), 2**32, 10),
+            ak.random_strings_uniform(1, 16, 10),
+        ]
+        for select_from in select_from_list:
+            arr1 = select_from[ak.randint(0, select_from.size, 20)]
+
+            # test unique search space, this should be identical to find
+            # be sure to test when all items are present and when there are items missing
+            for arr2 in select_from, select_from[:5], select_from[5:]:
+                found_in_second = ak.in1d(arr1, arr2)
+                idx_of_first_in_second = ak.indexof1d(arr1, arr2)
+
+                # ensure we match find
+                assert (idx_of_first_in_second == ak.find(arr1, arr2, remove_missing=True)).all()
+
+                # if an element of arr1 is found in arr2, return the index of that item in arr2
+                assert (arr2[idx_of_first_in_second] == arr1[found_in_second]).all()
+
+            # test duplicate items in search space, the easiest way I can think
+            # of to do this is to compare against pandas series getitem
+            arr2 = select_from[ak.randint(0, select_from.size, 20)]
+            pd_s = pd.Series(index=arr1.to_ndarray(), data=arr2.to_ndarray())
+            ak_s = ak.Series(index=arr1, data=arr2)
+
+            arr1_keys = ak.GroupBy(arr1).unique_keys
+            arr2_keys = ak.GroupBy(arr2).unique_keys
+            in_both = ak.intersect1d(arr1_keys, arr2_keys)
+
+            for i in in_both.to_list():
+                pd_i = pd_s[i]
+                ak_i = ak_s[i]
+                if isinstance(pd_i, pd.Series):
+                    assert isinstance(ak_i, ak.Series)
+                    assert pd_i.values.tolist() == ak_i.values.to_list()
+                else:
+                    assert pd_i == ak_i

arkouda/alignment.py

@@ -5,13 +5,14 @@
 import numpy as np  # type: ignore
 
 from arkouda.categorical import Categorical
+from arkouda.client import generic_msg
 from arkouda.dtypes import bigint
 from arkouda.dtypes import float64 as akfloat64
 from arkouda.dtypes import int64 as akint64
 from arkouda.dtypes import uint64 as akuint64
 from arkouda.groupbyclass import GroupBy, broadcast, unique
-from arkouda.numeric import where
-from arkouda.pdarrayclass import pdarray
+from arkouda.numeric import cumsum, where
+from arkouda.pdarrayclass import create_pdarray, pdarray
 from arkouda.pdarraycreation import arange, full, ones, zeros
 from arkouda.pdarraysetops import concatenate, in1d
 from arkouda.sorting import argsort, coargsort
@@ -109,23 +110,102 @@ class NonUniqueError(ValueError):
     pass
 
 
-def find(query, space):
+def find(query, space, all_occurrences=False, remove_missing=False):
     """
-    Return indices of query items in a search list of items (-1 if not found).
+    Return indices of query items in a search list of items.
 
     Parameters
     ----------
     query : (sequence of) array-like
         The items to search for. If multiple arrays, each "row" is an item.
     space : (sequence of) array-like
         The set of items in which to search. Must have same shape/dtype as query.
+    all_occurrences: bool
+        When duplicate terms are present in search space, if all_occurrences is True,
+        return all occurrences found as a SegArray, otherwise return only the first
+        occurrences as a pdarray. Defaults to only finding the first occurrence.
+        Finding all occurrences is not yet supported on sequences of arrays
+    remove_missing: bool
+        If False, return -1 for any items in query not found in space. If True,
+        remove these and only return indices of items that are found.
 
     Returns
     -------
-    indices : pdarray, int64
-        For each item in query, its index in space or -1 if not found.
-    """
+    indices : pdarray or SegArray
+        For each item in query, its index in space. If remove_missing is True,
+        exclued missing values otherwise return -1. If all_occurrences is False,
+        the return will be a pdarray of the first index where each value in the
+        query appears in the space. if all_occurrences is True, the return will be
+        a SegArray containing every index where each value in the query appears in
+        the space.
 
+    Examples
+    --------
+    >>> select_from = ak.arange(10)
+    >>> arr1 = select_from[ak.randint(0, select_from.size, 20, seed=10)]
+    >>> arr2 = select_from[ak.randint(0, select_from.size, 20, seed=11)]
+    # remove some values to ensure we have some values
+    # which don't appear in the search space
+    >>> arr2 = arr2[arr2 != 9]
+    >>> arr2 = arr2[arr2 != 3]
+
+    # find with defaults (all_occurrences and remove_missing both False)
+    >>> ak.find(arr1, arr2)
+    array([-1 -1 -1 0 1 -1 -1 -1 2 -1 5 -1 8 -1 5 -1 -1 11 5 0])
+
+     # set remove_missing to True, only difference from default
+     # is missing values are excluded
+     >>> ak.find(arr1, arr2, remove_missing=True)
+    array([0 1 2 5 8 5 11 5 0])
+
+    # set all_occurrences to True, the first index of each list
+    # is the first occurence and should match the default
+    >>> ak.find(arr1, arr2, all_occurrences=True).to_list()
+    [[-1],
+     [-1],
+     [-1],
+     [0, 4],
+     [1, 3, 10],
+     [-1],
+     [-1],
+     [-1],
+     [2, 6, 12, 13],
+     [-1],
+     [5, 7],
+     [-1],
+     [8, 9, 14],
+     [-1],
+     [5, 7],
+     [-1],
+     [-1],
+     [11, 15],
+     [5, 7],
+     [0, 4]]
+
+    # set both remove_missing and all_occurrences to True, missing values
+    # will be empty segments
+    >>> ak.find(arr1, arr2, remove_missing=True, all_occurrences=True).to_list()
+    [[],
+     [],
+     [],
+     [0, 4],
+     [1, 3, 10],
+     [],
+     [],
+     [],
+     [2, 6, 12, 13],
+     [],
+     [5, 7],
+     [],
+     [8, 9, 14],
+     [],
+     [5, 7],
+     [],
+     [],
+     [11, 15],
+     [5, 7],
+     [0, 4]]
+    """
     # Concatenate the space and query in fast (block interleaved) mode
     if isinstance(query, (pdarray, Strings, Categorical)):
         if type(query) is not type(space):
@@ -151,15 +231,48 @@ def find(query, space):
     # All space indices are less than all query indices
     i = concatenate((arange(spacesize), arange(spacesize, spacesize + querysize)), ordered=False)
     # Group on terms
-    g = GroupBy(c)
+    g = GroupBy(c, dropna=False)
     # For each term, count how many times it appears in the search space
     space_multiplicity = g.sum(i < spacesize)[1]
-    # Warn of any duplicate terms in space
-    if (space_multiplicity > 1).any():
-        warn(
-            "Duplicate terms present in search space. Only first instance of each query term\
-            will be reported."
-        )
+    has_duplicates = (space_multiplicity > 1).any()
+    # handle duplicate terms in space
+    if has_duplicates:
+        if all_occurrences:
+            if isinstance(query, Sequence):
+                raise TypeError("finding all_occurrences is not yet supported on sequences of arrays")
+
+            from arkouda.segarray import SegArray
+
+            # use segmented mink to select space_multiplicity number of elements
+            # and create a segarray which contains all the indices
+            # in our query space, instead of just the min for each segment
+
+            # only calculate where to place the negatives if remove_missing is false
+            negative_at = "" if remove_missing else space_multiplicity == 0
+            repMsg = generic_msg(
+                cmd="segmentedExtremaK",
+                args={
+                    "vals": i[g.permutation],
+                    "segs": g.segments,
+                    "segLens": g.size()[1],
+                    "kArray": space_multiplicity,
+                    "isMin": True,
+                    "removeMissing": remove_missing,
+                    "negativeAt": negative_at,
+                },
+            )
+            min_k_vals = create_pdarray(repMsg)
+            seg_idx = g.broadcast(arange(g.segments.size))[i >= spacesize]
+            if not remove_missing:
+                space_multiplicity += negative_at
+            min_k_segs = cumsum(space_multiplicity) - space_multiplicity
+            sa = SegArray(min_k_segs, min_k_vals)
+            return sa[seg_idx]
+        else:
+            warn(
+                "Duplicate terms present in search space. Only first instance of each query term"
+                " will be reported. To return all occurrences, set all_occurrences=True."
+            )
     # For query terms in the space, the min combined index will be the first index of that
     # term in the space
     uspaceidx = g.min(i)[1]
@@ -169,7 +282,8 @@ def find(query, space):
     # Broadcast unique term indices to combined list of space and query terms
     spaceidx = g.broadcast(uspaceidx)
     # Return only the indices of the query terms (remove the search space)
-    return spaceidx[i >= spacesize]
+    pda = spaceidx[i >= spacesize]
+    return pda[pda != -1] if remove_missing else pda
 
 
 def lookup(keys, values, arguments, fillvalue=-1):