New blosc2.argsort() function

Author: FrancescAlted

src/blosc2/__init__.py

@@ -533,6 +533,7 @@ def _raise(exc):
     eye,
     asarray,
     astype,
+    argsort,
     indices,
     sort,
     reshape,

src/blosc2/ndarray.py

@@ -5218,6 +5218,18 @@ def indices(self, order: str | list[str] | None = None, **kwargs: Any) -> NDArra
         """
         return indices(self, order, **kwargs)
 
+    def argsort(self, order: str | list[str] | None = None, **kwargs: Any) -> NDArray:
+        """
+        Return the permutation that sorts the array.
+
+        This follows :func:`numpy.argsort` semantics more closely than
+        :meth:`indices`: plain 1-D arrays are supported, and ``order=None``
+        means "use the array's natural order" rather than "leave unsorted".
+
+        See full documentation in :func:`argsort`.
+        """
+        return argsort(self, order, **kwargs)
+
     def itersorted(
         self,
         order: str | list[str] | None = None,
@@ -6748,6 +6760,60 @@ def indices(array: blosc2.Array, order: str | list[str] | None = None, **kwargs:
     return larr.indices(order).compute(**kwargs)
 
 
+def argsort(array: blosc2.Array, order: str | list[str] | None = None, **kwargs: Any) -> NDArray:
+    """
+    Return the indices that would sort the array.
+
+    This mirrors :func:`numpy.argsort` for 1-D arrays. Plain arrays sort by
+    their values. Structured arrays sort by ``order`` when provided, or by
+    their dtype field order when ``order=None``. Expression orders such as
+    ``"abs(x)"`` are also supported when a matching ``full`` expression index
+    exists.
+
+    Parameters
+    ----------
+    array: :ref:`blosc2.Array`
+        The 1-D array to be ordered.
+    order: str, list of str, optional
+        Primary and optional secondary order keys for structured arrays. When
+        omitted, NumPy's default record order is used for structured dtypes and
+        the array values themselves are used for plain dtypes.
+    kwargs: Any, optional
+        Keyword arguments that are supported by the :func:`empty` constructor.
+
+    Returns
+    -------
+    out: :ref:`NDArray`
+        The ordered logical positions as ``int64``.
+
+    Notes
+    -----
+    When the primary order key has a matching ``full`` field or expression
+    index, the permutation is returned directly from that index in ascending
+    stable order. Secondary keys refine ties after the primary indexed order.
+    Without a matching ``full`` index, :func:`argsort` falls back to
+    materializing the input values and delegating ordering to
+    :func:`numpy.argsort`.
+
+    The result is always a new array materialization. For persistent inputs,
+    the returned permutation is in memory by default; pass storage kwargs such
+    as ``urlpath`` (and typically ``mode="w"``) if the permutation should also
+    be persisted on disk.
+    """
+    if isinstance(array, blosc2.NDArray):
+        from . import indexing
+
+        ordered = indexing.ordered_indices(array, order=order)
+        if ordered is not None:
+            return blosc2.asarray(ordered, **kwargs)
+        if indexing.is_expression_order(array, order):
+            raise ValueError("expression order requires a matching full expression index")
+
+    values = array[:]
+    positions = np.argsort(values, order=order, kind="stable")
+    return blosc2.asarray(positions.astype(np.int64, copy=False), **kwargs)
+
+
 def sort(array: blosc2.Array, order: str | list[str] | None = None, **kwargs: Any) -> NDArray:
     """
     Return a sorted array following the specified order.

tests/ndarray/test_indexing.py

@@ -674,6 +674,29 @@ def test_full_index_reuses_primary_order_for_indices_and_sort():
     np.testing.assert_array_equal(arr.sort(order=["a", "b"])[:], np.sort(data, order=["a", "b"]))
 
 
+def test_full_index_reuses_primary_order_for_argsort():
+    dtype = np.dtype([("a", np.int64), ("b", np.int64)])
+    data = np.array(
+        [(2, 9), (1, 8), (2, 7), (1, 6), (2, 5), (1, 4), (2, 3), (1, 2), (2, 1), (1, 0)],
+        dtype=dtype,
+    )
+    arr = blosc2.asarray(data, chunks=(4,), blocks=(2,))
+    arr.create_csindex("a")
+
+    np.testing.assert_array_equal(arr.argsort(order=["a", "b"])[:], np.argsort(data, order=["a", "b"]))
+
+
+def test_persistent_scalar_argsort_uses_full_index(tmp_path):
+    path = tmp_path / "scalar_argsort.b2nd"
+    data = np.array([9, 1, 7, 3, 1, 5], dtype=np.int64)
+    arr = blosc2.asarray(data, urlpath=path, mode="w", chunks=(3,), blocks=(2,))
+    arr.create_index(kind="full")
+
+    result = blosc2.argsort(arr)
+
+    np.testing.assert_array_equal(result[:], np.argsort(data, kind="stable"))
+
+
 def test_filtered_ordered_queries_support_cross_field_exact_indexes():
     dtype = np.dtype([("a", np.int64), ("b", np.int64), ("payload", np.int32)])
     data = np.array(

tests/ndarray/test_ndarray.py

@@ -482,6 +482,23 @@ def test_indices(order):
     assert np.array_equal(b[:], nb)
 
 
+@pytest.mark.parametrize("order", ["f0", "f1", "f2", None])
+def test_argsort_structured(order):
+    it = ((x + 1, x - 2, -x) for x in range(10))
+    a = blosc2.fromiter(it, dtype="i4, i4, i8", shape=(10,))
+    b = blosc2.argsort(a, order=order)
+    narr = a[:]
+    nb = np.argsort(narr, order=order, kind="stable")
+    assert np.array_equal(b[:], nb)
+
+
+def test_argsort_scalar():
+    data = np.array([7, 2, 9, 2, 1, 8], dtype=np.int64)
+    a = blosc2.asarray(data)
+    b = a.argsort()
+    np.testing.assert_array_equal(b[:], np.argsort(data, kind="stable"))
+
+
 def test_save():
     a = blosc2.arange(0, 10, 1, dtype="i4", shape=(10,))
     blosc2.save(a, "test.b2nd")