Correctly resize output during extract (#2)

The previous code extracted the correct rows and columns from the input,
but did not renumber the indices or adjust the output size. This code
fixes those deficiencies and adds more tests to ensure full coverage of
`extract`.

Author: jim22k

mlir_graphblas/implementations.py

@@ -24,7 +24,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
+from .utils import get_sparse_output_pointer, get_scalar_output_pointer, renumber_indices
 from .types import RankedTensorType, BOOL, INT64, FP64
 
 
@@ -902,44 +902,58 @@ def main(x):
         return compile(module)
 
 
-def extract(tensor: SparseTensorBase, row_indices, col_indices=None):
+def extract(tensor: SparseTensorBase, row_indices, col_indices=None, row_size=None, col_size=None):
     # There may be a way to do this in MLIR, but for now we use numpy
     if tensor.ndims == 1:
         # Vector
         assert col_indices is None
-        if row_indices is None:  # None indicate GrB_ALL
+        assert col_size is None
+
+        if row_indices is None:  # None indicates GrB_ALL
             return tensor.dup()
+
         rowidx, vals = tensor.extract_tuples()
+        row_indices = np.array(row_indices)
         selected = np.isin(rowidx, row_indices)
-        v = Vector.new(tensor.dtype, *tensor.shape)
-        v.build(rowidx[selected], vals[selected])
+        # Filter and renumber rowidx
+        rowidx, vals = rowidx[selected], vals[selected]
+        rowidx = renumber_indices(rowidx, row_indices)
+        v = Vector.new(tensor.dtype, row_size)
+        v.build(rowidx, vals)
         return v
 
     # Matrix
     if row_indices is None and col_indices is None:
         return tensor.dup()
+
     rowidx, colidx, vals = tensor.extract_tuples()
     if row_indices is not None:
-        rowsel = np.isin(rowidx, row_indices)
-        # Apply rowsel filter
+        rindices_arr = np.array(row_indices)
+        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 col_indices is not None:
-        colsel = np.isin(colidx, col_indices)
-        # Apply colsel filter
+        cindices_arr = np.array(col_indices)
+        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(row_indices) is int:
         # Extract row as Vector
         assert np.all(rowidx == row_indices)
-        v = Vector.new(tensor.dtype, tensor.shape[1])
+        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)
-        v = Vector.new(tensor.dtype, tensor.shape[0])
+        v = Vector.new(tensor.dtype, row_size)
         v.build(rowidx, vals)
         return v
-    m = Matrix.new(tensor.dtype, *tensor.shape)
+    m = Matrix.new(tensor.dtype, row_size, col_size)
     m.build(rowidx, colidx, vals)
     return m
 

mlir_graphblas/operations.py

@@ -502,27 +502,48 @@ def extract(out: SparseTensor,
         tensor = TransposedMatrix.wrap(tensor)
 
     # Check indices
-    if tensor.ndims < 1:
+    if tensor.ndims == 0:  # Scalar input
         raise TypeError("Use `extract_element` rather than `extract` for Scalars")
-    if tensor.ndims < 2 and col_indices is not None:
-        raise ValueError("col_indices not allowed for Vector, use row_indices")
-
-    # Compare shapes
-    if type(row_indices) is int and type(col_indices) is int:
-        raise TypeError("Cannot provide int for both row_indices and col_indices")
+    elif tensor.ndims == 1:  # Vector input
+        if col_indices is not None:
+            raise ValueError("col_indices not allowed for Vector, use row_indices")
+        if type(row_indices) is int:
+            raise TypeError("Use extract_element to get a single element from the Vector")
+    else:  # Matrix input
+        if type(row_indices) is int and type(col_indices) is int:
+            raise TypeError("Use extract_element to get a single element from the Matrix")
+
+    # Compute output sizes
     if type(row_indices) is int:
-        expected_out_shape = (tensor.shape[1],)
-    elif type(col_indices) is int:
-        expected_out_shape = (tensor.shape[0],)
+        row_size = None
+    elif row_indices is None:
+        row_size = tensor.shape[0]
     else:
-        expected_out_shape = tensor.shape
+        row_size = len(row_indices)
+
+    if type(col_indices) is int or tensor.ndims < 2:
+        col_size = None
+    elif col_indices is None:
+        col_size = tensor.shape[1]
+    else:
+        col_size = len(col_indices)
+
+    # Compare shapes
+    if tensor.ndims == 1:  # Vector input
+        expected_out_shape = (row_size,)
+    else:  # Matrix input
+        if type(row_indices) is int:
+            expected_out_shape = (col_size,)
+        elif type(col_indices) is int:
+            expected_out_shape = (row_size,)
+        else:
+            expected_out_shape = (row_size, col_size)
     if out.shape != expected_out_shape:
         raise GrbDimensionMismatch(f"output shape mismatch: {out.shape} != {expected_out_shape}")
 
+    result = impl.extract(tensor, row_indices, col_indices, row_size, col_size)
     if mask is not None:
-        tensor = impl.apply_mask(tensor, mask, desc)
-
-    result = impl.extract(tensor, row_indices, col_indices)
+        result = impl.apply_mask(result, mask, desc)
     update(out, result, mask, accum, desc)
 
 

mlir_graphblas/tests/test_operations.py

@@ -292,56 +292,90 @@ def test_reduce_scalar_vec(vs):
 def test_extract_vec(vs):
     x, _ = vs
     xidx, xvals = x.extract_tuples()
-    z = Vector.new(x.dtype, *x.shape)
+    z = Vector.new(x.dtype, 3)
     operations.extract(z, x, [0, 1, 3])
     idx, vals = z.extract_tuples()
-    np_assert_equal(idx, [1, 3])
+    np_assert_equal(idx, [1, 2])
     np_assert_allclose(vals, [10., 30.])
 
-    # None == GrB_ALL
-    operations.extract(z, x, None)
-    idx, vals = z.extract_tuples()
+    # Extract all
+    z2 = Vector.new(x.dtype, *x.shape)
+    operations.extract(z2, x, None)
+    idx, vals = z2.extract_tuples()
     np_assert_equal(idx, xidx)
     np_assert_allclose(vals, xvals)
 
 
 def test_extract_mat(mm):
     x, _ = mm
     xrows, xcols, xvals = x.extract_tuples()
-    z = Matrix.new(x.dtype, *x.shape)
-    operations.extract(z, x, [0, 4], [1, 3, 5])
-    rowidx, colidx, vals = z.extract_tuples()
-    np_assert_equal(rowidx, [0, 0])
-    np_assert_equal(colidx, [3, 5])
-    np_assert_allclose(vals, [1.1, 2.2])
 
-    # None == GrB_ALL
+    # Extract all rows, all cols
+    z = Matrix.new(x.dtype, *x.shape)
     operations.extract(z, x, None, None)
     rowidx, colidx, vals = z.extract_tuples()
     np_assert_equal(rowidx, xrows)
     np_assert_equal(colidx, xcols)
     np_assert_allclose(vals, xvals)
 
+    # Extract some rows, some cols
+    z2 = Matrix.new(x.dtype, 2, 4)
+    operations.extract(z2, x, [0, 4], [1, 2, 3, 5])
+    rowidx, colidx, vals = z2.extract_tuples()
+    np_assert_equal(rowidx, [0, 0, 1])
+    np_assert_equal(colidx, [2, 3, 1])
+    np_assert_allclose(vals, [1.1, 2.2, 6.6])
+
+    # Extract some rows, all cols
+    z3 = Matrix.new(x.dtype, 2, x.shape[1])
+    operations.extract(z3, x, [0, 4], None)
+    rowidx, colidx, vals = z3.extract_tuples()
+    np_assert_equal(rowidx, [0, 0, 1])
+    np_assert_equal(colidx, [3, 5, 2])
+    np_assert_allclose(vals, [1.1, 2.2, 6.6])
+
+    # Extract all rows, some cols
+    z4 = Matrix.new(x.dtype, x.shape[0], 4)
+    operations.extract(z4, x, None, [1, 5, 3, 2])
+    rowidx, colidx, vals = z4.extract_tuples()
+    np_assert_equal(rowidx, [0, 0, 1, 2, 4])
+    np_assert_equal(colidx, [1, 2, 2, 0, 3])
+    np_assert_allclose(vals, [2.2, 1.1, 3.3, 5.5, 6.6])
+
 
 def test_extract_vec_from_mat(mm):
     x, _ = mm
-    # Extract column
-    z = Vector.new(x.dtype, x.shape[0])
-    operations.extract(z, x, [0, 1, 4], 3)
+    # Extract partial column
+    z = Vector.new(x.dtype, 3)
+    operations.extract(z, x, [0, 1, 4], 2)
     idx, vals = z.extract_tuples()
+    np_assert_equal(idx, [2])
+    np_assert_allclose(vals, [6.6])
+
+    # Extract full column
+    z1 = Vector.new(x.dtype, x.shape[0])
+    operations.extract(z1, x, None, 3)
+    idx, vals = z1.extract_tuples()
     np_assert_equal(idx, [0, 1])
     np_assert_allclose(vals, [1.1, 3.3])
 
-    # Extract row
-    z = Vector.new(x.dtype, x.shape[1])
-    operations.extract(z, x, 2, [0, 1, 4])
-    idx, vals = z.extract_tuples()
-    np_assert_equal(idx, [0, 1])
-    np_assert_allclose(vals, [4.4, 5.5])
+    # Extract partial row
+    z2 = Vector.new(x.dtype, 5)
+    operations.extract(z2, x, 0, [0, 1, 3, 4, 5])
+    idx, vals = z2.extract_tuples()
+    np_assert_equal(idx, [2, 4])
+    np_assert_allclose(vals, [1.1, 2.2])
 
-    # Extract column via transposed input
-    z = Vector.new(x.dtype, x.shape[0])
-    operations.extract(z, x, 3, [0, 1, 4], desc=desc.T0)
-    idx, vals = z.extract_tuples()
-    np_assert_equal(idx, [0, 1])
-    np_assert_allclose(vals, [1.1, 3.3])
+    # Extract full row
+    z3 = Vector.new(x.dtype, x.shape[1])
+    operations.extract(z3, x, 0, None)
+    idx, vals = z3.extract_tuples()
+    np_assert_equal(idx, [3, 5])
+    np_assert_allclose(vals, [1.1, 2.2])
+
+    # Extract partial column via transposed input
+    z3 = Vector.new(x.dtype, 3)
+    operations.extract(z3, x, 2, [0, 1, 4], desc=desc.T0)
+    idx, vals = z3.extract_tuples()
+    np_assert_equal(idx, [2])
+    np_assert_allclose(vals, [6.6])

mlir_graphblas/tests/test_utils.py

@@ -0,0 +1,22 @@
+import pytest
+import numpy as np
+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]))

mlir_graphblas/utils.py

@@ -1,4 +1,5 @@
 import ctypes
+import numpy as np
 from enum import Enum
 from mlir import ir
 from .exceptions import (
@@ -38,6 +39,48 @@ def ensure_scalar_of_type(obj, dtype):
     return s
 
 
+def renumber_indices(indices, selected):
+    """
+    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])
+    """
+    # Check that values in selected are unique
+    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]
+
+
 # https://mlir.llvm.org/docs/Dialects/ArithOps/#arithcmpi-mlirarithcmpiop
 class CmpIPredicate(Enum):
     eq = 0  # equal