Gather component-major arrays row by row in sort_by_key

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+- `xp.algorithms.sort_by_key` gathers a 2-D NumPy array along its last axis (component-major
+  markers, `axis=-1`) row by row, about 2x faster than `take` (1e6 markers: 3 rows 12.9 to 6.7 ms,
+  8 rows 37 to 19 ms).
+
 ## [0.7.0] - 2026-10-20
 
 ### Added

src/cunumpy/_algorithms.py

@@ -380,6 +380,8 @@ def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
     radix sort; on NumPy, integer keys of more than 4096 entries get an LSD
     radix sort too (one stable ``uint16`` argsort per 16-bit digit the key range
     needs), about ten times faster than NumPy's stable sort of 64-bit integers.
+    On NumPy, a 2-D array sorted along its last axis (component-major markers) is
+    gathered row by row, about twice as fast as ``take`` along that axis.
 
     Parameters
     ----------
@@ -437,12 +439,27 @@ def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
         keys[order],
         order,
         *(
-            xpm.take(array, order, axis=array_axis)
+            _take_along_axis(xpm, array, order, array_axis)
             for array, array_axis in zip(arrays, axes, strict=True)
         ),
     )
 
 
+def _take_along_axis(xpm: Any, array: Any, order: Any, axis: int) -> Any:
+    """Return `array` taken at `order` along `axis`, as a new C-contiguous array.
+
+    A component-major ``(ncomp, n)`` NumPy array is gathered one contiguous row at a
+    time, which is about twice as fast as ``take`` along its last axis; everything
+    else, and CuPy, goes through ``take``.
+    """
+    if xpm is np and array.ndim == 2 and axis == 1:
+        gathered = np.empty((array.shape[0], order.shape[0]), dtype=array.dtype)
+        for row in range(array.shape[0]):
+            gathered[row] = array[row][order]
+        return gathered
+    return xpm.take(array, order, axis=axis)
+
+
 def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int:
     """Move the entries where `mask` is True to the front of every array, in place.
 

tests/unit/test_morton.py

@@ -270,3 +270,45 @@ def test_cupy_arrays_stay_on_the_device():
     assert bool((sorted_keys[1:] >= sorted_keys[:-1]).all())
     cx, _ = xp.algorithms.morton_decode(sorted_keys, 2)
     assert isinstance(cx, cp.ndarray)
+
+
+@pytest.mark.parametrize(
+    "dtype", [np.float64, np.float32, np.int64, np.int32, np.bool_]
+)
+@pytest.mark.parametrize("ncomp", [1, 3, 8])
+def test_sort_by_key_gathers_component_major_arrays_like_take(dtype, ncomp):
+    rng = np.random.default_rng(ncomp)
+    n = 5000
+    keys = rng.integers(0, 50, n)
+    array = (rng.random((ncomp, n)) * 100).astype(dtype)
+
+    _, order, gathered = xp.algorithms.sort_by_key(keys, array, axis=-1)
+
+    expected = np.take(array, order, axis=1)
+    np.testing.assert_array_equal(gathered, expected)
+    assert gathered.dtype == array.dtype
+    assert gathered.flags.c_contiguous
+
+
+def test_sort_by_key_gathers_strided_views_other_axes_and_empty_arrays():
+    rng = np.random.default_rng(0)
+    keys = rng.integers(0, 10, 100)
+    storage = rng.random((3, 150))  # markers[:, :100] of a store with spare capacity
+    transposed = rng.random((100, 3))  # axis=0 is the plain take path
+    volume = rng.random((2, 100, 4))  # three axes: the plain take path as well
+
+    _, order, from_view = xp.algorithms.sort_by_key(keys, storage[:, :100], axis=-1)
+    np.testing.assert_array_equal(from_view, storage[:, :100][:, order])
+
+    _, order, from_rows = xp.algorithms.sort_by_key(keys, transposed, axis=0)
+    np.testing.assert_array_equal(from_rows, transposed[order])
+
+    _, order, from_volume = xp.algorithms.sort_by_key(keys, volume, axis=1)
+    np.testing.assert_array_equal(from_volume, volume[:, order])
+
+    empty_keys = np.zeros(0, dtype=np.int64)
+    _, empty_order, empty = xp.algorithms.sort_by_key(
+        empty_keys, np.zeros((3, 0)), axis=-1
+    )
+    assert empty_order.shape == (0,)
+    assert empty.shape == (3, 0)