Fix extract with duplicated selected indices (#6)

Author: jim22k

mlir_graphblas/implementations.py

@@ -16,7 +16,7 @@
 from .operators import UnaryOp, BinaryOp, SelectOp, IndexUnaryOp, Monoid, Semiring
 from .compiler import compile, engine_cache
 from . descriptor import Descriptor, NULL as NULL_DESC
-from .utils import get_sparse_output_pointer, get_scalar_output_pointer, renumber_indices
+from .utils import get_sparse_output_pointer, get_scalar_output_pointer, pick_and_renumber_indices
 from .types import RankedTensorType, BOOL, INT64, FP64
 from .exceptions import GrbIndexOutOfBounds, GrbDimensionMismatch
 
@@ -1010,14 +1010,11 @@ def extract(tensor: SparseTensorBase, row_indices, col_indices, row_size, col_si
         if row_indices is None:  # None indicates GrB_ALL
             return dup(tensor)
 
-        rowidx, vals = tensor.extract_tuples()
-        row_indices = np.array(row_indices, dtype=np.uint64)
-        selected = np.isin(rowidx, row_indices)
-        # Filter and renumber rowidx
-        rowidx, vals = rowidx[selected], vals[selected]
-        rowidx = renumber_indices(rowidx, row_indices)
+        idx, vals = tensor.extract_tuples()
+        pick_list = np.array(row_indices, dtype=np.uint64)
+        idx, vals = pick_and_renumber_indices(pick_list, idx, vals)
         v = Vector.new(tensor.dtype, row_size)
-        v.build(rowidx, vals)
+        v.build(idx, vals)
         return v
 
     # Matrix
@@ -1026,28 +1023,26 @@ def extract(tensor: SparseTensorBase, row_indices, col_indices, row_size, col_si
 
     rowidx, colidx, vals = tensor.extract_tuples()
     if row_indices is not None:
-        rindices_arr = np.array(row_indices, dtype=np.uint64)
-        rowsel = np.isin(rowidx, rindices_arr)
-        # Filter and renumber rowidx
-        rowidx, colidx, vals = rowidx[rowsel], colidx[rowsel], vals[rowsel]
-        if type(row_indices) is not int:
-            rowidx = renumber_indices(rowidx, rindices_arr)
+        if type(row_indices) is int:
+            pick_list = np.array([row_indices], dtype=np.uint64)
+        else:
+            pick_list = np.array(row_indices, dtype=np.uint64)
+        rowidx, colidx, vals = pick_and_renumber_indices(pick_list, rowidx, colidx, vals)
     if col_indices is not None:
-        cindices_arr = np.array(col_indices, dtype=np.uint64)
-        colsel = np.isin(colidx, cindices_arr)
-        # Filter and renumber colidx
-        rowidx, colidx, vals = rowidx[colsel], colidx[colsel], vals[colsel]
-        if type(col_indices) is not int:
-            colidx = renumber_indices(colidx, cindices_arr)
+        if type(col_indices) is int:
+            pick_list = np.array([col_indices], dtype=np.uint64)
+        else:
+            pick_list = np.array(col_indices, dtype=np.uint64)
+        colidx, rowidx, vals = pick_and_renumber_indices(pick_list, colidx, rowidx, vals)
     if type(row_indices) is int:
         # Extract row as Vector
-        assert np.all(rowidx == row_indices)
+        assert np.all(rowidx == 0)
         v = Vector.new(tensor.dtype, col_size)
         v.build(colidx, vals)
         return v
     if type(col_indices) is int:
         # Extract col as Vector
-        assert np.all(colidx == col_indices)
+        assert np.all(colidx == 0)
         v = Vector.new(tensor.dtype, row_size)
         v.build(rowidx, vals)
         return v

mlir_graphblas/tests/test_operations.py

@@ -251,6 +251,11 @@ def test_extract_vec(vs):
     operations.extract(z2, x, None)
     vector_compare(z2, xidx, xvals)
 
+    # Extract duplicates
+    z3 = Vector.new(x.dtype, 5)
+    operations.extract(z3, x, [0, 3, 3, 3, 2])
+    vector_compare(z3, [1, 2, 3, 4], [30., 30., 30., 20.])
+
 
 def test_extract_mat(mm):
     x, _ = mm
@@ -262,22 +267,22 @@ def test_extract_mat(mm):
     matrix_compare(z, xrows, xcols, xvals)
 
     # Extract some rows, some cols
-    z2 = Matrix.new(x.dtype, 2, 4)
-    operations.extract(z2, x, [0, 4], [1, 2, 3, 5])
-    matrix_compare(z2, [0, 0, 1], [2, 3, 1], [1.1, 2.2, 6.6])
+    z2 = Matrix.new(x.dtype, 2, 5)
+    operations.extract(z2, x, [0, 4], [1, 2, 3, 5, 3])
+    matrix_compare(z2, [0, 0, 0, 1], [2, 3, 4, 1], [1.1, 2.2, 1.1, 6.6])
 
     # Extract some rows, all cols
-    z3 = Matrix.new(x.dtype, 2, x.shape[1])
-    operations.extract(z3, x, [0, 4], None)
-    matrix_compare(z3, [0, 0, 1], [3, 5, 2], [1.1, 2.2, 6.6])
+    z3 = Matrix.new(x.dtype, 4, x.shape[1])
+    operations.extract(z3, x, [0, 4, 3, 0], None)
+    matrix_compare(z3, [0, 0, 1, 3, 3], [3, 5, 2, 3, 5], [1.1, 2.2, 6.6, 1.1, 2.2])
 
     # Extract all rows, some cols
-    z4 = Matrix.new(x.dtype, x.shape[0], 4)
-    operations.extract(z4, x, None, [1, 5, 3, 2])
+    z4 = Matrix.new(x.dtype, x.shape[0], 5)
+    operations.extract(z4, x, None, [1, 5, 3, 2, 1])
     matrix_compare(z4,
-                   [0, 0, 1, 2, 4],
-                   [1, 2, 2, 0, 3],
-                   [2.2, 1.1, 3.3, 5.5, 6.6])
+                   [0, 0, 1, 2, 2, 4],
+                   [1, 2, 2, 0, 4, 3],
+                   [2.2, 1.1, 3.3, 5.5, 5.5, 6.6])
 
 
 def test_extract_vec_from_mat(mm):
@@ -293,9 +298,9 @@ def test_extract_vec_from_mat(mm):
     vector_compare(z1, [0, 1], [1.1, 3.3])
 
     # Extract partial row
-    z2 = Vector.new(x.dtype, 5)
-    operations.extract(z2, x, 0, [0, 1, 3, 4, 5])
-    vector_compare(z2, [2, 4], [1.1, 2.2])
+    z2 = Vector.new(x.dtype, 8)
+    operations.extract(z2, x, 0, [0, 1, 3, 4, 5, 3, 5, 3])
+    vector_compare(z2, [2, 4, 5, 6, 7], [1.1, 2.2, 1.1, 2.2, 1.1])
 
     # Extract full row
     z3 = Vector.new(x.dtype, x.shape[1])

mlir_graphblas/tests/test_utils.py

@@ -3,20 +3,12 @@
 from mlir_graphblas import utils
 
 
-def test_renumber_indices():
-    a = np.array([1, 1, 1, 3, 5], dtype=np.uint64)
-    b = np.array([1, 2, 5, 3], dtype=np.uint64)
-    c = utils.renumber_indices(a, b)
-    assert c.dtype == np.uint64
-    np.testing.assert_equal(c, [0, 0, 0, 3, 2])
-
-    d = np.array([1, 2, 5, 47, 48, 49, 3], dtype=np.uint64)
-    e = utils.renumber_indices(a, d)
-    np.testing.assert_equal(e, [0, 0, 0, 6, 2])
-
-
-def test_renumber_indices_errors():
-    with pytest.raises(ValueError, match="4"):
-        utils.renumber_indices(np.array([1, 1, 1, 3, 5]), np.array([1, 4, 2, 5, 3, 4]))
-    with pytest.raises(KeyError, match="11"):
-        utils.renumber_indices(np.array([1, 2, 5, 11]), np.array([1, 2, 5, 3, 4]))
+def test_pick_and_renumber_indices():
+    rows = np.array([0, 0, 1, 1, 1, 1], dtype=np.uint64)
+    cols = np.array([1, 2, 0, 1, 2, 3], dtype=np.uint64)
+    vals = np.array([1, 2, 6, 7, 8, 9])
+    selected = np.array([1, 1, 0, 0, 1])
+    rows, cols, vals = utils.pick_and_renumber_indices(selected, rows, cols, vals)
+    np.testing.assert_equal(rows, [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 4, 4])
+    np.testing.assert_equal(cols, [0, 1, 2, 3, 0, 1, 2, 3, 1, 2, 1, 2, 0, 1, 2, 3])
+    np.testing.assert_equal(vals, [6, 7, 8, 9, 6, 7, 8, 9, 1, 2, 1, 2, 6, 7, 8, 9])

mlir_graphblas/utils.py

@@ -47,47 +47,43 @@ def ensure_unique(indices, name=None):
         raise ValueError(f"Found duplicate indices{name_str}: {unique[counts > 1]}")
 
 
-def renumber_indices(indices, selected):
+def pick_and_renumber_indices(selected, indices, *related):
     """
-    Given a set of non-unique `indices`, returns an array of the same size
-    as `indices` with values renumbered according to the positions in `selected`.
-
-    All values in indices must also be found in selected.
-
-    If these were Python lists instead of numpy arrays, this would be
-    equivalent to calling `[selected.index(x) for x in indices]`.
-    However, this will be much faster as it uses numpy to perform
-    the lookups.
-
-    :param indices: ndarray of non-unique positive integers
-    :param selected: ndarray of unique positive integers
-    :return: ndarray of same length as indices
-
-    Example
-    -------
-    >>> a = np.array([1, 1, 1, 3, 5])
-    >>> b = np.array([1, 2, 5, 3])
-    >>> renumber_indices(a, b)
-    array([0, 0, 0, 3, 2])
+    This function is used by the `extract` operation.
+
+    Picks elements from indices as indicated in `selected`. The new indices will
+    be renumbered according to the position in `selected`. Related ndarrays will
+    maintain matching order during the picking, but values will not be changed.
+
+    :param selected: values to pick from indices in the order indicated
+    :param indices: ndarray representing indices
+    :param related: list of ndarray of the same length as indices that must maintain
+                    proper relationship as indices are picked and rearranged
+    :return: tuple of modified ndarrays -- (indices, *related)
     """
-    # Check that values in selected are unique
-    # TODO: fix this; extract must allow non-unique selected indices
-    unique = np.unique(selected)
-    if unique.size < selected.size:
-        unique, counts = np.unique(selected, return_counts=True)
-        raise ValueError(f"Found duplicate values in `selected`: {unique[counts > 1]}")
-
-    # Check for required inclusion criteria
-    not_found = np.setdiff1d(indices, selected)
-    if not_found.size > 0:
-        raise KeyError(f"Found values in `indices` not contained in `selected`: {not_found}")
-
-    # To be efficient, the searching must be done on a sorted array
-    # Build the sort_order to map back to the original order
-    sort_order = np.argsort(selected)
-    renumbered_indices = np.arange(len(selected), dtype=indices.dtype)[sort_order]
-    pos = np.searchsorted(selected[sort_order], indices)
-    return renumbered_indices[pos]
+    for r in related:
+        assert len(r) == len(indices)
+    # Ensure indices are sorted (this is needed for np.searchsorted)
+    sort_order = np.argsort(indices)
+    indices = indices[sort_order]
+    related = [r[sort_order] for r in related]
+    # Find which `selected` exist in indices
+    is_found = np.isin(selected, indices)
+    selfound = selected[is_found]
+    # Perform sorted search of first and last occurrence
+    # Duplicates are handled by building a ramp of indexes between first and last
+    pick_start = np.searchsorted(indices, selfound, side='left')
+    pick_end = np.searchsorted(indices, selfound, side='right')
+    pick_delta = pick_end - pick_start
+    expanded_pick = np.ones(pick_delta.sum(), dtype=np.int64)
+    expanded_pick[0] = pick_start[0]
+    expanded_pick[np.cumsum(pick_delta)[:-1]] += pick_start[1:] - pick_end[:-1]
+    pick_index = np.cumsum(expanded_pick)
+    # Use pick_index to index into related arrays
+    new_related = [r[pick_index] for r in related]
+    # Build new indices by repeating range
+    new_indices = np.repeat(np.arange(len(selected), dtype=np.uint64)[is_found], pick_delta)
+    return new_indices, *new_related
 
 
 # https://mlir.llvm.org/docs/Dialects/ArithOps/#arithcmpi-mlirarithcmpiop