From 3baa34883b406650d96743ca38a883879b4cde2a Mon Sep 17 00:00:00 2001
From: Max <max.lindqvist@ipp.mpg.de>
Date: Tue, 20 Oct 2026 22:50:47 +0200
Subject: [PATCH 1/3] Add take_columns, a fast gather of columns

---
 CHANGELOG.md                                  |   8 +
 .../content/docs/guides/execution-helpers.md  |  15 ++
 .../content/docs/guides/particle-codes.mdx    |   6 +-
 docs/src/content/docs/reference/overview.md   |   2 +-
 src/cunumpy/LLM_GUIDE.md                      |   1 +
 src/cunumpy/_algorithms.py                    | 131 +++++++++++
 src/cunumpy/_kernels/__init__.py              |   1 +
 src/cunumpy/_kernels/take_columns/__init__.py |  11 +
 .../take_columns/take_columns_numpy.py        |  10 +
 .../take_columns/take_columns_pyccel.py       |  27 +++
 src/cunumpy/algorithms.py                     |   7 +-
 tests/unit/test_take_columns.py               | 214 ++++++++++++++++++
 12 files changed, 429 insertions(+), 4 deletions(-)
 create mode 100644 src/cunumpy/_kernels/__init__.py
 create mode 100644 src/cunumpy/_kernels/take_columns/__init__.py
 create mode 100644 src/cunumpy/_kernels/take_columns/take_columns_numpy.py
 create mode 100644 src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
 create mode 100644 tests/unit/test_take_columns.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2c84d1e..8f994a0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- `xp.algorithms.take_columns(array, index, *, out=None)`: gathers the columns of a
+  component-major `(ncomp, n)` array at integer indices or a boolean mask, as a
+  C-contiguous copy (`array[:, index]` is Fortran-ordered on NumPy). On NumPy a
+  compiled OpenMP kernel (Pyccel, one row per thread) is about 2.5 times faster than
+  the best NumPy gather for a million columns and falls back to NumPy without a
+  compiler; on CuPy it is `cupy.take`.
+
 ## [0.7.0] - 2026-10-20
 
 ### Added
diff --git a/docs/src/content/docs/guides/execution-helpers.md b/docs/src/content/docs/guides/execution-helpers.md
index 3d785f5..04b1118 100644
--- a/docs/src/content/docs/guides/execution-helpers.md
+++ b/docs/src/content/docs/guides/execution-helpers.md
@@ -163,6 +163,21 @@ unspecified. The count is needed on the host, so on CuPy each call
 synchronizes once. For component-major arrays, `(ncomp, N)` next to `(N,)`,
 pass `axis=-1` and continue with `positions[:, :n]`.
 
+## Gather marker columns: `take_columns`
+
+```python
+positions = xp.algorithms.take_columns(positions[:, :length], order)
+alive_positions = xp.algorithms.take_columns(positions, alive)  # boolean mask
+xp.algorithms.take_columns(velocities, order, out=buffer[:, :length])
+```
+
+`array[:, index]` of a component-major `(ncomp, n)` array, as a C-contiguous
+copy (NumPy's own comes back Fortran-ordered), for integer indices or a boolean
+mask. On NumPy, from about 30000 columns on, a compiled OpenMP kernel (Pyccel,
+one contiguous row per thread; float64/32, int64/32 and bool) is used, otherwise
+a row-by-row NumPy gather; on CuPy it is `cupy.take`. Out-of-bounds indices raise
+`IndexError` on NumPy and wrap around on CuPy (checking would synchronize).
+
 ## Prepare cell ranges and segment reductions
 
 ```python
diff --git a/docs/src/content/docs/guides/particle-codes.mdx b/docs/src/content/docs/guides/particle-codes.mdx
index 58ca835..b38a55a 100644
--- a/docs/src/content/docs/guides/particle-codes.mdx
+++ b/docs/src/content/docs/guides/particle-codes.mdx
@@ -45,7 +45,11 @@ Apply `order` to every per-marker array (positions, velocities, weights, ids),
 e.g. by keeping them as columns of one `(n, k)` array. A stable sort keeps the
 result independent of how the markers were ordered before.
 `xp.algorithms.sort_by_key(keys, positions, velocities, weights)` does the argsort and
-the reordering of several arrays in one call.
+the reordering of several arrays in one call. For component-major `(ncomp, n)`
+markers, `xp.algorithms.take_columns(positions, order)` gathers the columns in a
+compiled, threaded kernel on NumPy (about 2.5 times faster than `positions[:, order]`
+for a million markers) and with `cupy.take` on CuPy; a boolean mask works too, and
+`out=` writes into a preallocated array.
 
 For a tree code, or for better locality in 2D and 3D, sort by Morton key
 (`xp.algorithms.morton_keys`, see [CUDA headers](/cunumpy/kernels/cuda-headers/) for
diff --git a/docs/src/content/docs/reference/overview.md b/docs/src/content/docs/reference/overview.md
index ffe525f..cddd320 100644
--- a/docs/src/content/docs/reference/overview.md
+++ b/docs/src/content/docs/reference/overview.md
@@ -51,7 +51,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not
 | [`cunumpy.arguments`](/cunumpy/api/cunumpy/arguments/) | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` |
 | [`cunumpy.cuda`](/cunumpy/api/cunumpy/cuda/) | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) |
 | [`cunumpy.rng`](/cunumpy/api/cunumpy/rng/) | both | `random_streams`, `get_rng`, `philox_*` |
-| [`cunumpy.algorithms`](/cunumpy/api/cunumpy/algorithms/) | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask` |
+| [`cunumpy.algorithms`](/cunumpy/api/cunumpy/algorithms/) | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask`, `take_columns` |
 | [`cunumpy.optimize`](/cunumpy/api/cunumpy/optimize/) | both | `newton` (batched Newton/secant), `NewtonResult`; `fsolve`, `root`, `minimize` (SciPy on the host, callbacks on the device) |
 | [`cunumpy.integrate`](/cunumpy/api/cunumpy/integrate/) | both | `quad`, `odeint` (SciPy on the host, callbacks on the device) |
 | [`cunumpy.mpi`](/cunumpy/api/cunumpy/mpi/) | both | `mpi_buffer`, `MPIStaging`, CUDA-aware MPI detection, `synchronize_for_mpi`, and the MPI stand-ins of maybempi (`get_mpi`, `local_rank`, ...) |
diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md
index 7fad843..858b76e 100644
--- a/src/cunumpy/LLM_GUIDE.md
+++ b/src/cunumpy/LLM_GUIDE.md
@@ -75,6 +75,7 @@ https://max-models.github.io/cunumpy/ and in `docs/` of the repository.
 | run a Metal (MSL) kernel on an Apple silicon GPU, NumPy float32 in and out (float64 raises; `float64="cast"` computes in float32); not part of `Kernel` dispatch | `xp.kernels.MetalKernel(body, inputs=[...], outputs=[...])(*args, out=arrays, n_threads=n)`; check `xp.kernels.metal_available()` |
 | call host-only code (SciPy, file readers) with arguments of either backend | `xp.host_call(fun, *args)`; `@xp.evaluate_on_host` on a method; `@xp.setup_on_host` on `__init__` |
 | keep the live rows of particle arrays at the front | `n = xp.algorithms.compact_by_mask(alive, markers, weights)`; component-major `(ncomp, N)`: `axis=-1` |
+| reorder or filter the columns of `(ncomp, N)` markers (fast, C-contiguous, replaces `a[:, order]`) | `xp.algorithms.take_columns(a, order_or_mask, out=None)` |
 | run CUDA kernel launches on the CPU in a test (fake CuPy) | `with kernel_testing.emulated_launches(): ...` (also makes `kernel.compile()`/`compile_all()` build the emulation instead of CUDA); `with kernel_testing.fake_cupy_session(): ...` adds the CuPy backend; `kernel_testing.host_buffer(a)` reads a fake array; struct arguments are read through their fields |
 | run a test on the fake CuPy from a NumPy test process | `kernel_testing.run_in_fake_cupy_subprocess(code)` (serial child, rank 0 only under MPI); `@kernel_testing.requires_device_backend` = a GPU or the fake CuPy |
 | find the arguments a host kernel writes (copy only those back) | `PyccelKernel(fn, outputs=("out",))` (names work positionally); `xp.kernels.outputs_from_annotations(fn)`; `Kernel.from_folder(..., outputs="annotations")` |
diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py
index 11dd85f..cf3d400 100644
--- a/src/cunumpy/_algorithms.py
+++ b/src/cunumpy/_algorithms.py
@@ -510,3 +510,134 @@ def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int:
         # the right side is a copy: no overlap problem
         array[front] = xpm.take(array, rows, axis=array_axis)
     return n_kept
+
+
+#: dtypes the compiled ``take_columns`` kernel is built for; others use NumPy.
+_TAKE_COLUMNS_DTYPES = frozenset(
+    np.dtype(name) for name in ("float64", "float32", "int64", "int32", "bool")
+)
+
+
+#: Fewer columns are gathered by NumPy: the kernel's call and thread start-up
+#: cost about 60 microseconds, which is what NumPy needs for ~30000 columns.
+_TAKE_COLUMNS_MIN_COLUMNS = 32768
+
+
+def _checked_columns(index: Any, n: int, xpm: Any) -> Any:
+    """Return the integer column indices of `index`, validated against `n` columns."""
+    index = xpm.asarray(index)
+    if index.ndim != 1:
+        raise ValueError(f"index must be 1D, got shape {index.shape}")
+    if index.dtype == np.bool_:
+        if index.shape[0] != n:
+            raise ValueError(
+                f"a boolean index needs one entry per column ({n}), "
+                f"got {index.shape[0]}",
+            )
+        return xpm.nonzero(index)[0]
+    if index.dtype.kind not in "iu":
+        raise TypeError("index must be an integer or boolean array")
+    return index
+
+
+def take_columns(array: Any, index: Any, *, out: Any = None) -> Any:
+    """Gather columns of a component-major 2-D array, as ``array[:, index]``.
+
+    The step of a particle code that reorders or filters its markers stored as
+    ``(ncomp, n)`` arrays: after :func:`sort_by_key`, or to drop the dead
+    markers. Unlike ``array[:, index]``, which NumPy returns Fortran-ordered,
+    the result is C-contiguous. On NumPy a compiled kernel (Pyccel, OpenMP)
+    gathers one contiguous row per thread, about 2.5 times faster than the
+    best NumPy gather for a random permutation of a million columns (below
+    about 30000 columns NumPy is faster and is used); without a compiler it
+    falls back to NumPy, one row at a time. On CuPy it is
+    ``cupy.take``.
+
+    Parameters
+    ----------
+    array : array, shape (ncomp, n)
+        The source; a view such as ``storage[:, :length]`` is fine.
+    index : array of int or bool, shape (m,) or (n,)
+        Indices of the columns to take (negative ones count from the end), or a
+        boolean mask with one entry per column.
+    out : array, optional
+        Output to overwrite: shape ``(ncomp, m)``, the dtype of `array`,
+        writable, C-contiguous and not overlapping `array`.
+
+    Returns
+    -------
+    array, shape (ncomp, m)
+        ``array[:, index]``, C-contiguous; `out` if given.
+
+    Raises
+    ------
+    ValueError
+        If `array` is not 2-D, `index` is not 1D (or a mask of the wrong
+        length), or `out` is unsuitable.
+    TypeError
+        If `index` is neither integer nor boolean.
+    IndexError
+        If an index is out of bounds. Checked on NumPy only: on CuPy the check
+        would synchronize, and out-of-range indices wrap around.
+
+    See Also
+    --------
+    sort_by_key : Computes the permutation to gather with.
+    compact_by_mask : Filters in place instead.
+
+    Examples
+    --------
+    >>> a = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    >>> xp.algorithms.take_columns(a, xp.asarray([2, 0]))
+    array([[3., 1.],
+           [6., 4.]])
+    >>> xp.algorithms.take_columns(a, xp.asarray([True, False, True]))
+    array([[1., 3.],
+           [4., 6.]])
+    """
+    assert_same_backend(array, index, *(() if out is None else (out,)))
+    xpm = get_array_module(array)
+    array = xpm.asarray(array)
+    if array.ndim != 2:
+        raise ValueError(f"array must be 2D (ncomp, n), got shape {array.shape}")
+    ncomp, n = array.shape
+    columns = _checked_columns(index, n, xpm)
+    shape = (ncomp, columns.shape[0])
+    if out is not None:
+        if out.shape != shape or out.dtype != array.dtype or not out.flags.c_contiguous:
+            raise ValueError(
+                f"out must have shape {shape}, dtype {array.dtype} and be C-contiguous",
+            )
+        if not getattr(out.flags, "writeable", True):
+            raise ValueError("out must be writable")
+        if xpm.may_share_memory(out, array):
+            raise ValueError("out must not overlap array")
+    if get_array_backend(array) == "cupy":
+        # CuPy's take wraps out-of-range indices, negative ones as NumPy does
+        return xpm.take(array, columns, axis=1, out=out, mode="wrap")
+    result = np.empty(shape, dtype=array.dtype) if out is None else out
+    if result.size == 0:
+        return result
+    low, high = int(columns.min()), int(columns.max())
+    if high >= n or low < -n:
+        raise IndexError(
+            f"column index out of bounds for {n} columns "
+            f"(indices from {low} to {high})",
+        )
+    if low < 0:
+        columns = np.where(columns < 0, columns + n, columns)
+    if (
+        array.dtype in _TAKE_COLUMNS_DTYPES
+        and array.strides[1] == array.itemsize
+        and columns.shape[0] >= _TAKE_COLUMNS_MIN_COLUMNS
+    ):
+        from cunumpy._kernels.take_columns import kernel
+
+        kernel(array, np.ascontiguousarray(columns, dtype=np.int64), result)
+    else:
+        from cunumpy._kernels.take_columns.take_columns_numpy import (
+            take_columns as take_numpy,
+        )
+
+        take_numpy(array, columns, result)
+    return result
diff --git a/src/cunumpy/_kernels/__init__.py b/src/cunumpy/_kernels/__init__.py
new file mode 100644
index 0000000..2f1c2fb
--- /dev/null
+++ b/src/cunumpy/_kernels/__init__.py
@@ -0,0 +1 @@
+"""Host kernels of cunumpy's own algorithms, one folder per kernel."""
diff --git a/src/cunumpy/_kernels/take_columns/__init__.py b/src/cunumpy/_kernels/take_columns/__init__.py
new file mode 100644
index 0000000..bb73bb7
--- /dev/null
+++ b/src/cunumpy/_kernels/take_columns/__init__.py
@@ -0,0 +1,11 @@
+"""The ``take_columns`` kernel: Pyccel and NumPy versions with the same arguments."""
+
+from cunumpy._dispatch import Kernel
+from cunumpy._pyccel_cache import PyccelCache
+
+kernel = Kernel.from_folder(
+    __name__,
+    host_suffix="_pyccel",
+    dispatch="arrays",
+    compile_host=PyccelCache("cunumpy"),
+)
diff --git a/src/cunumpy/_kernels/take_columns/take_columns_numpy.py b/src/cunumpy/_kernels/take_columns/take_columns_numpy.py
new file mode 100644
index 0000000..1d60590
--- /dev/null
+++ b/src/cunumpy/_kernels/take_columns/take_columns_numpy.py
@@ -0,0 +1,10 @@
+"""Gather the columns ``order`` of a component-major array, in NumPy."""
+
+import numpy as np
+
+
+def take_columns(source: np.ndarray, order: np.ndarray, target: np.ndarray) -> None:
+    """Write ``source[:, order[i]]`` to ``target[:, i]`` for every entry of ``order``."""
+    # one contiguous row at a time: faster than a 2-D gather
+    for component in range(source.shape[0]):
+        target[component] = source[component][order]
diff --git a/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py b/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
new file mode 100644
index 0000000..46d1aff
--- /dev/null
+++ b/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
@@ -0,0 +1,27 @@
+"""Gather the columns ``order`` of a component-major array, in parallel (Pyccel)."""
+
+from typing import TypeVar
+
+T = TypeVar(
+    "T",
+    "float[:,:]",
+    "float32[:,:]",  # noqa: F821
+    "int[:,:]",
+    "int32[:,:]",  # noqa: F821
+    "bool[:,:]",
+)
+
+
+def take_columns(source: "T", order: "int[:]", target: "T") -> None:
+    """Write ``source[:, order[i]]`` to ``target[:, i]`` for every entry of ``order``.
+
+    The indices must lie in ``[0, source.shape[1])`` (the caller checks them) and
+    ``target`` must not overlap ``source``. One loop per component, so every
+    thread streams through one contiguous row.
+    """
+    for component in range(source.shape[0]):
+        # fmt: off
+        #$ omp parallel for schedule(static)
+        # fmt: on
+        for marker in range(order.shape[0]):
+            target[component, marker] = source[component, order[marker]]
diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py
index c13cebc..e371bea 100644
--- a/src/cunumpy/algorithms.py
+++ b/src/cunumpy/algorithms.py
@@ -2,13 +2,14 @@
 
 Morton (Z-order) keys (the same as ``cunumpy/morton.cuh`` computes in a
 kernel, at most ``MAX_MORTON_LEVELS`` bits per axis), a stable sort of several
-arrays by one key, sums per key, and the compaction of the live rows of
-particle arrays::
+arrays by one key, sums per key, the compaction of the live rows of
+particle arrays and the gather of columns of component-major markers::
 
     keys = xp.algorithms.morton_keys(positions, lower, upper, levels)
     keys, order, positions = xp.algorithms.sort_by_key(keys, positions)
     charge = xp.algorithms.segment_sum(q, cell, n_cells)
     n = xp.algorithms.compact_by_mask(alive, positions, charges)
+    markers = xp.algorithms.take_columns(markers, order)
 
 See :doc:`/guides/particle-codes`.
 """
@@ -20,6 +21,7 @@
     segment_boundaries,
     segment_sum,
     sort_by_key,
+    take_columns,
 )
 from cunumpy._morton import (
     MAX_MORTON_LEVELS,
@@ -41,4 +43,5 @@
     "segment_boundaries",
     "segment_sum",
     "sort_by_key",
+    "take_columns",
 ]
diff --git a/tests/unit/test_take_columns.py b/tests/unit/test_take_columns.py
new file mode 100644
index 0000000..445ebb7
--- /dev/null
+++ b/tests/unit/test_take_columns.py
@@ -0,0 +1,214 @@
+"""take_columns gathers columns of (ncomp, n) arrays like array[:, index], C-contiguous."""
+
+import numpy as np
+import pytest
+
+import cunumpy as xp
+from cunumpy import _algorithms
+from cunumpy.kernels import use_host_kernel_implementation
+
+DTYPES = [np.float64, np.float32, np.int64, np.int32, np.bool_, np.complex128]
+IMPLEMENTATIONS = ["default", "numpy"]
+
+
+def make(shape, dtype, seed=0):
+    rng = np.random.default_rng(seed)
+    return (rng.random(shape) * 100).astype(dtype)
+
+
+def run(implementation, *args, **kwargs):
+    if implementation == "default":
+        return xp.algorithms.take_columns(*args, **kwargs)
+    with use_host_kernel_implementation(implementation):
+        return xp.algorithms.take_columns(*args, **kwargs)
+
+
+@pytest.fixture(params=IMPLEMENTATIONS)
+def implementation(request, monkeypatch):
+    # run the compiled kernel on the small test arrays too
+    monkeypatch.setattr(_algorithms, "_TAKE_COLUMNS_MIN_COLUMNS", 0)
+    return request.param
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("shape", [(1, 1), (3, 7), (8, 1000), (3, 5000)])
+def test_integer_index_matches_fancy_indexing(implementation, dtype, shape):
+    array = make(shape, dtype)
+    index = np.random.default_rng(1).integers(0, shape[1], 2 * shape[1] // 3 + 1)
+    result = run(implementation, array, index)
+    assert result.dtype == array.dtype
+    assert result.flags.c_contiguous
+    np.testing.assert_array_equal(result, array[:, index])
+
+
+def test_permutation_of_a_large_array(implementation):
+    array = make((3, 100_000), np.float64)
+    order = np.random.default_rng(2).permutation(100_000)
+    np.testing.assert_array_equal(run(implementation, array, order), array[:, order])
+
+
+@pytest.mark.parametrize("index_dtype", [np.int32, np.int64, np.uint32, np.uint64])
+def test_index_dtypes(implementation, index_dtype):
+    array = make((2, 50), np.float64)
+    index = np.array([4, 0, 49, 4], dtype=index_dtype)
+    np.testing.assert_array_equal(run(implementation, array, index), array[:, index])
+
+
+def test_negative_indices_count_from_the_end(implementation):
+    array = make((2, 10), np.float64)
+    index = np.array([-1, 0, -10, 3])
+    np.testing.assert_array_equal(run(implementation, array, index), array[:, index])
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.bool_])
+def test_strided_column_prefix_view(implementation, dtype):
+    storage = make((3, 40), dtype)
+    view = storage[:, :25]  # the rows are further apart than the view is wide
+    index = np.random.default_rng(3).permutation(25)
+    np.testing.assert_array_equal(run(implementation, view, index), view[:, index])
+
+
+def test_non_row_contiguous_input(implementation):
+    base = make((9, 6), np.float64)
+    for array in (base.T, base[:, ::2], np.asfortranarray(base)):
+        index = np.arange(array.shape[1])[::-1]
+        np.testing.assert_array_equal(
+            run(implementation, array, index), array[:, index]
+        )
+
+
+def test_read_only_input(implementation):
+    array = make((2, 20), np.float64)
+    array.flags.writeable = False
+    index = np.arange(20)[::-1]
+    np.testing.assert_array_equal(run(implementation, array, index), array[:, index])
+
+
+def test_empty_index_and_empty_array(implementation):
+    array = make((3, 5), np.float64)
+    result = run(implementation, array, np.array([], dtype=np.int64))
+    assert result.shape == (3, 0)
+    assert result.dtype == array.dtype
+    assert run(implementation, np.zeros((0, 4)), np.array([1, 2])).shape == (0, 2)
+    assert run(implementation, np.zeros((3, 0)), np.array([], dtype=int)).shape == (
+        3,
+        0,
+    )
+    assert run(implementation, array, np.zeros(5, bool)).shape == (3, 0)
+
+
+def test_out_is_filled_and_returned(implementation):
+    array = make((3, 30), np.float64)
+    index = np.random.default_rng(4).permutation(30)[:20]
+    out = np.full((3, 20), -1.0)
+    result = run(implementation, array, index, out=out)
+    assert result is out
+    np.testing.assert_array_equal(out, array[:, index])
+
+
+def test_out_may_be_a_row_block_of_a_larger_array(implementation):
+    array = make((2, 10), np.float64)
+    storage = np.zeros((4, 10))
+    out = storage[:2]  # C-contiguous
+    run(implementation, array, np.arange(10)[::-1], out=out)
+    np.testing.assert_array_equal(storage[:2], array[:, ::-1])
+    assert not storage[2:].any()
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.bool_])
+@pytest.mark.parametrize("shape", [(3, 1), (3, 9), (8, 5000)])
+def test_boolean_mask_matches_compress(implementation, dtype, shape):
+    array = make(shape, dtype)
+    mask = np.random.default_rng(5).random(shape[1]) < 0.4
+    result = run(implementation, array, mask)
+    assert result.flags.c_contiguous
+    np.testing.assert_array_equal(result, np.compress(mask, array, axis=1))
+
+
+def test_mask_of_a_strided_view_with_out(implementation):
+    storage = make((3, 40), np.float64)
+    view = storage[:, :30]
+    mask = np.arange(30) % 3 == 0
+    out = np.empty((3, 10))
+    run(implementation, view, mask, out=out)
+    np.testing.assert_array_equal(out, view[:, mask])
+
+
+def test_the_kernel_is_used_from_the_minimum_number_of_columns(monkeypatch):
+    import cunumpy._kernels.take_columns as folder
+
+    calls = []
+    monkeypatch.setattr(folder, "kernel", lambda *args: calls.append(args))
+    size = _algorithms._TAKE_COLUMNS_MIN_COLUMNS
+    array = make((2, size), np.float64)
+    xp.algorithms.take_columns(array, np.arange(size - 1))
+    assert not calls
+    xp.algorithms.take_columns(array, np.arange(size))
+    assert len(calls) == 1
+
+
+def test_error_arguments():
+    array = np.zeros((2, 4))
+    with pytest.raises(ValueError, match="2D"):
+        xp.algorithms.take_columns(np.zeros(4), np.array([0]))
+    with pytest.raises(ValueError, match="1D"):
+        xp.algorithms.take_columns(array, np.zeros((2, 2), dtype=int))
+    with pytest.raises(TypeError, match="integer or boolean"):
+        xp.algorithms.take_columns(array, np.array([0.0]))
+    with pytest.raises(ValueError, match="one entry per column"):
+        xp.algorithms.take_columns(array, np.ones(3, bool))
+    with pytest.raises(IndexError, match="out of bounds"):
+        xp.algorithms.take_columns(array, np.array([0, 4]))
+    with pytest.raises(IndexError, match="out of bounds"):
+        xp.algorithms.take_columns(array, np.array([-5]))
+
+
+def test_out_errors():
+    array = np.zeros((2, 4))
+    index = np.array([0, 1])
+    with pytest.raises(ValueError, match="shape"):
+        xp.algorithms.take_columns(array, index, out=np.zeros((2, 3)))
+    with pytest.raises(ValueError, match="dtype"):
+        xp.algorithms.take_columns(array, index, out=np.zeros((2, 2), dtype=np.int64))
+    with pytest.raises(ValueError, match="C-contiguous"):
+        xp.algorithms.take_columns(array, index, out=np.zeros((2, 4))[:, ::2])
+    read_only = np.zeros((2, 2))
+    read_only.flags.writeable = False
+    with pytest.raises(ValueError, match="writable"):
+        xp.algorithms.take_columns(array, index, out=read_only)
+    with pytest.raises(ValueError, match="overlap"):
+        xp.algorithms.take_columns(array, np.arange(4), out=array)
+
+
+def test_compiled_kernel_agrees_with_numpy_implementation(monkeypatch):
+    monkeypatch.setattr(_algorithms, "_TAKE_COLUMNS_MIN_COLUMNS", 0)
+    from cunumpy._kernels.take_columns import kernel
+
+    pytest.importorskip("pyccel")
+    assert kernel.selected() == "pyccel"
+    array = make((3, 3000), np.float64)
+    index = np.random.default_rng(6).permutation(3000)
+    with use_host_kernel_implementation("pyccel"):
+        compiled = xp.algorithms.take_columns(array, index)
+    with use_host_kernel_implementation("numpy"):
+        reference = xp.algorithms.take_columns(array, index)
+    np.testing.assert_array_equal(compiled, reference)
+
+
+def test_cupy_backend_matches_numpy():
+    if not xp.cupy_available():
+        pytest.skip("CuPy not installed or not functional")
+    array = make((3, 200), np.float64)
+    index = np.random.default_rng(7).permutation(200)[:150]
+    mask = index.size and np.random.default_rng(8).random(200) < 0.5
+    with xp.use_backend("cupy"):
+        device = xp.to_cupy(array)
+        taken = xp.algorithms.take_columns(device, xp.to_cupy(index))
+        masked = xp.algorithms.take_columns(device, xp.to_cupy(mask))
+        out = xp.zeros((3, 150))
+        xp.algorithms.take_columns(device, xp.to_cupy(index), out=out)
+        negative = xp.algorithms.take_columns(device, xp.to_cupy(np.array([-1, 0])))
+    np.testing.assert_array_equal(xp.to_numpy(taken), array[:, index])
+    np.testing.assert_array_equal(xp.to_numpy(masked), array[:, mask])
+    np.testing.assert_array_equal(xp.to_numpy(out), array[:, index])
+    np.testing.assert_array_equal(xp.to_numpy(negative), array[:, [-1, 0]])

From 98d496037328475358e6438ab8ef9cd86cc0f8ff Mon Sep 17 00:00:00 2001
From: Max <max.lindqvist@ipp.mpg.de>
Date: Tue, 20 Oct 2026 22:56:45 +0200
Subject: [PATCH 2/3] Fix the numpydoc sections and skip the compiled-kernel
 test without a compiler

---
 .../_kernels/take_columns/take_columns_numpy.py      | 12 +++++++++++-
 .../_kernels/take_columns/take_columns_pyccel.py     |  9 +++++++++
 tests/unit/test_take_columns.py                      |  5 ++++-
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/cunumpy/_kernels/take_columns/take_columns_numpy.py b/src/cunumpy/_kernels/take_columns/take_columns_numpy.py
index 1d60590..834daf7 100644
--- a/src/cunumpy/_kernels/take_columns/take_columns_numpy.py
+++ b/src/cunumpy/_kernels/take_columns/take_columns_numpy.py
@@ -4,7 +4,17 @@
 
 
 def take_columns(source: np.ndarray, order: np.ndarray, target: np.ndarray) -> None:
-    """Write ``source[:, order[i]]`` to ``target[:, i]`` for every entry of ``order``."""
+    """Write ``source[:, order[i]]`` to ``target[:, i]`` for every entry of ``order``.
+
+    Parameters
+    ----------
+    source : numpy.ndarray
+        The component-major array, ``(ncomp, n)``.
+    order : numpy.ndarray
+        The column indices to gather, ``(m,)``, within ``[0, n)``.
+    target : numpy.ndarray
+        Receives the gathered columns, ``(ncomp, m)``; must not overlap ``source``.
+    """
     # one contiguous row at a time: faster than a 2-D gather
     for component in range(source.shape[0]):
         target[component] = source[component][order]
diff --git a/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py b/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
index 46d1aff..b4df0e2 100644
--- a/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
+++ b/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
@@ -18,6 +18,15 @@ def take_columns(source: "T", order: "int[:]", target: "T") -> None:
     The indices must lie in ``[0, source.shape[1])`` (the caller checks them) and
     ``target`` must not overlap ``source``. One loop per component, so every
     thread streams through one contiguous row.
+
+    Parameters
+    ----------
+    source : array of shape (ncomp, n)
+        The component-major array.
+    order : array of int, shape (m,)
+        The column indices to gather.
+    target : array of shape (ncomp, m)
+        Receives the gathered columns.
     """
     for component in range(source.shape[0]):
         # fmt: off
diff --git a/tests/unit/test_take_columns.py b/tests/unit/test_take_columns.py
index 445ebb7..d97414f 100644
--- a/tests/unit/test_take_columns.py
+++ b/tests/unit/test_take_columns.py
@@ -185,7 +185,10 @@ def test_compiled_kernel_agrees_with_numpy_implementation(monkeypatch):
     from cunumpy._kernels.take_columns import kernel
 
     pytest.importorskip("pyccel")
-    assert kernel.selected() == "pyccel"
+    if kernel.selected() != "pyccel":
+        pytest.skip(
+            "the compiled kernel is not available (needs Pyccel and a compiler)"
+        )
     array = make((3, 3000), np.float64)
     index = np.random.default_rng(6).permutation(3000)
     with use_host_kernel_implementation("pyccel"):

From 8caeb23b2fd1bd818ddcf3c8a33b547a65a6c387 Mon Sep 17 00:00:00 2001
From: Max <max.lindqvist@ipp.mpg.de>
Date: Tue, 20 Oct 2026 23:05:50 +0200
Subject: [PATCH 3/3] Gather component-major arrays row by row in sort_by_key
 instead of adding take_columns

---
 CHANGELOG.md                                  |  11 +-
 .../content/docs/guides/execution-helpers.md  |  15 --
 .../content/docs/guides/particle-codes.mdx    |   6 +-
 docs/src/content/docs/reference/overview.md   |   2 +-
 src/cunumpy/LLM_GUIDE.md                      |   1 -
 src/cunumpy/_algorithms.py                    | 150 ++----------
 src/cunumpy/_kernels/__init__.py              |   1 -
 src/cunumpy/_kernels/take_columns/__init__.py |  11 -
 .../take_columns/take_columns_numpy.py        |  20 --
 .../take_columns/take_columns_pyccel.py       |  36 ---
 src/cunumpy/algorithms.py                     |   7 +-
 tests/unit/test_morton.py                     |  42 ++++
 tests/unit/test_take_columns.py               | 217 ------------------
 13 files changed, 68 insertions(+), 451 deletions(-)
 delete mode 100644 src/cunumpy/_kernels/__init__.py
 delete mode 100644 src/cunumpy/_kernels/take_columns/__init__.py
 delete mode 100644 src/cunumpy/_kernels/take_columns/take_columns_numpy.py
 delete mode 100644 src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
 delete mode 100644 tests/unit/test_take_columns.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8f994a0..6b46bef 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,13 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### Added
-- `xp.algorithms.take_columns(array, index, *, out=None)`: gathers the columns of a
-  component-major `(ncomp, n)` array at integer indices or a boolean mask, as a
-  C-contiguous copy (`array[:, index]` is Fortran-ordered on NumPy). On NumPy a
-  compiled OpenMP kernel (Pyccel, one row per thread) is about 2.5 times faster than
-  the best NumPy gather for a million columns and falls back to NumPy without a
-  compiler; on CuPy it is `cupy.take`.
+### 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
 
diff --git a/docs/src/content/docs/guides/execution-helpers.md b/docs/src/content/docs/guides/execution-helpers.md
index 04b1118..3d785f5 100644
--- a/docs/src/content/docs/guides/execution-helpers.md
+++ b/docs/src/content/docs/guides/execution-helpers.md
@@ -163,21 +163,6 @@ unspecified. The count is needed on the host, so on CuPy each call
 synchronizes once. For component-major arrays, `(ncomp, N)` next to `(N,)`,
 pass `axis=-1` and continue with `positions[:, :n]`.
 
-## Gather marker columns: `take_columns`
-
-```python
-positions = xp.algorithms.take_columns(positions[:, :length], order)
-alive_positions = xp.algorithms.take_columns(positions, alive)  # boolean mask
-xp.algorithms.take_columns(velocities, order, out=buffer[:, :length])
-```
-
-`array[:, index]` of a component-major `(ncomp, n)` array, as a C-contiguous
-copy (NumPy's own comes back Fortran-ordered), for integer indices or a boolean
-mask. On NumPy, from about 30000 columns on, a compiled OpenMP kernel (Pyccel,
-one contiguous row per thread; float64/32, int64/32 and bool) is used, otherwise
-a row-by-row NumPy gather; on CuPy it is `cupy.take`. Out-of-bounds indices raise
-`IndexError` on NumPy and wrap around on CuPy (checking would synchronize).
-
 ## Prepare cell ranges and segment reductions
 
 ```python
diff --git a/docs/src/content/docs/guides/particle-codes.mdx b/docs/src/content/docs/guides/particle-codes.mdx
index b38a55a..58ca835 100644
--- a/docs/src/content/docs/guides/particle-codes.mdx
+++ b/docs/src/content/docs/guides/particle-codes.mdx
@@ -45,11 +45,7 @@ Apply `order` to every per-marker array (positions, velocities, weights, ids),
 e.g. by keeping them as columns of one `(n, k)` array. A stable sort keeps the
 result independent of how the markers were ordered before.
 `xp.algorithms.sort_by_key(keys, positions, velocities, weights)` does the argsort and
-the reordering of several arrays in one call. For component-major `(ncomp, n)`
-markers, `xp.algorithms.take_columns(positions, order)` gathers the columns in a
-compiled, threaded kernel on NumPy (about 2.5 times faster than `positions[:, order]`
-for a million markers) and with `cupy.take` on CuPy; a boolean mask works too, and
-`out=` writes into a preallocated array.
+the reordering of several arrays in one call.
 
 For a tree code, or for better locality in 2D and 3D, sort by Morton key
 (`xp.algorithms.morton_keys`, see [CUDA headers](/cunumpy/kernels/cuda-headers/) for
diff --git a/docs/src/content/docs/reference/overview.md b/docs/src/content/docs/reference/overview.md
index cddd320..ffe525f 100644
--- a/docs/src/content/docs/reference/overview.md
+++ b/docs/src/content/docs/reference/overview.md
@@ -51,7 +51,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not
 | [`cunumpy.arguments`](/cunumpy/api/cunumpy/arguments/) | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` |
 | [`cunumpy.cuda`](/cunumpy/api/cunumpy/cuda/) | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) |
 | [`cunumpy.rng`](/cunumpy/api/cunumpy/rng/) | both | `random_streams`, `get_rng`, `philox_*` |
-| [`cunumpy.algorithms`](/cunumpy/api/cunumpy/algorithms/) | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask`, `take_columns` |
+| [`cunumpy.algorithms`](/cunumpy/api/cunumpy/algorithms/) | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask` |
 | [`cunumpy.optimize`](/cunumpy/api/cunumpy/optimize/) | both | `newton` (batched Newton/secant), `NewtonResult`; `fsolve`, `root`, `minimize` (SciPy on the host, callbacks on the device) |
 | [`cunumpy.integrate`](/cunumpy/api/cunumpy/integrate/) | both | `quad`, `odeint` (SciPy on the host, callbacks on the device) |
 | [`cunumpy.mpi`](/cunumpy/api/cunumpy/mpi/) | both | `mpi_buffer`, `MPIStaging`, CUDA-aware MPI detection, `synchronize_for_mpi`, and the MPI stand-ins of maybempi (`get_mpi`, `local_rank`, ...) |
diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md
index 858b76e..7fad843 100644
--- a/src/cunumpy/LLM_GUIDE.md
+++ b/src/cunumpy/LLM_GUIDE.md
@@ -75,7 +75,6 @@ https://max-models.github.io/cunumpy/ and in `docs/` of the repository.
 | run a Metal (MSL) kernel on an Apple silicon GPU, NumPy float32 in and out (float64 raises; `float64="cast"` computes in float32); not part of `Kernel` dispatch | `xp.kernels.MetalKernel(body, inputs=[...], outputs=[...])(*args, out=arrays, n_threads=n)`; check `xp.kernels.metal_available()` |
 | call host-only code (SciPy, file readers) with arguments of either backend | `xp.host_call(fun, *args)`; `@xp.evaluate_on_host` on a method; `@xp.setup_on_host` on `__init__` |
 | keep the live rows of particle arrays at the front | `n = xp.algorithms.compact_by_mask(alive, markers, weights)`; component-major `(ncomp, N)`: `axis=-1` |
-| reorder or filter the columns of `(ncomp, N)` markers (fast, C-contiguous, replaces `a[:, order]`) | `xp.algorithms.take_columns(a, order_or_mask, out=None)` |
 | run CUDA kernel launches on the CPU in a test (fake CuPy) | `with kernel_testing.emulated_launches(): ...` (also makes `kernel.compile()`/`compile_all()` build the emulation instead of CUDA); `with kernel_testing.fake_cupy_session(): ...` adds the CuPy backend; `kernel_testing.host_buffer(a)` reads a fake array; struct arguments are read through their fields |
 | run a test on the fake CuPy from a NumPy test process | `kernel_testing.run_in_fake_cupy_subprocess(code)` (serial child, rank 0 only under MPI); `@kernel_testing.requires_device_backend` = a GPU or the fake CuPy |
 | find the arguments a host kernel writes (copy only those back) | `PyccelKernel(fn, outputs=("out",))` (names work positionally); `xp.kernels.outputs_from_annotations(fn)`; `Kernel.from_folder(..., outputs="annotations")` |
diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py
index cf3d400..d7e1621 100644
--- a/src/cunumpy/_algorithms.py
+++ b/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.
 
@@ -510,134 +527,3 @@ def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int:
         # the right side is a copy: no overlap problem
         array[front] = xpm.take(array, rows, axis=array_axis)
     return n_kept
-
-
-#: dtypes the compiled ``take_columns`` kernel is built for; others use NumPy.
-_TAKE_COLUMNS_DTYPES = frozenset(
-    np.dtype(name) for name in ("float64", "float32", "int64", "int32", "bool")
-)
-
-
-#: Fewer columns are gathered by NumPy: the kernel's call and thread start-up
-#: cost about 60 microseconds, which is what NumPy needs for ~30000 columns.
-_TAKE_COLUMNS_MIN_COLUMNS = 32768
-
-
-def _checked_columns(index: Any, n: int, xpm: Any) -> Any:
-    """Return the integer column indices of `index`, validated against `n` columns."""
-    index = xpm.asarray(index)
-    if index.ndim != 1:
-        raise ValueError(f"index must be 1D, got shape {index.shape}")
-    if index.dtype == np.bool_:
-        if index.shape[0] != n:
-            raise ValueError(
-                f"a boolean index needs one entry per column ({n}), "
-                f"got {index.shape[0]}",
-            )
-        return xpm.nonzero(index)[0]
-    if index.dtype.kind not in "iu":
-        raise TypeError("index must be an integer or boolean array")
-    return index
-
-
-def take_columns(array: Any, index: Any, *, out: Any = None) -> Any:
-    """Gather columns of a component-major 2-D array, as ``array[:, index]``.
-
-    The step of a particle code that reorders or filters its markers stored as
-    ``(ncomp, n)`` arrays: after :func:`sort_by_key`, or to drop the dead
-    markers. Unlike ``array[:, index]``, which NumPy returns Fortran-ordered,
-    the result is C-contiguous. On NumPy a compiled kernel (Pyccel, OpenMP)
-    gathers one contiguous row per thread, about 2.5 times faster than the
-    best NumPy gather for a random permutation of a million columns (below
-    about 30000 columns NumPy is faster and is used); without a compiler it
-    falls back to NumPy, one row at a time. On CuPy it is
-    ``cupy.take``.
-
-    Parameters
-    ----------
-    array : array, shape (ncomp, n)
-        The source; a view such as ``storage[:, :length]`` is fine.
-    index : array of int or bool, shape (m,) or (n,)
-        Indices of the columns to take (negative ones count from the end), or a
-        boolean mask with one entry per column.
-    out : array, optional
-        Output to overwrite: shape ``(ncomp, m)``, the dtype of `array`,
-        writable, C-contiguous and not overlapping `array`.
-
-    Returns
-    -------
-    array, shape (ncomp, m)
-        ``array[:, index]``, C-contiguous; `out` if given.
-
-    Raises
-    ------
-    ValueError
-        If `array` is not 2-D, `index` is not 1D (or a mask of the wrong
-        length), or `out` is unsuitable.
-    TypeError
-        If `index` is neither integer nor boolean.
-    IndexError
-        If an index is out of bounds. Checked on NumPy only: on CuPy the check
-        would synchronize, and out-of-range indices wrap around.
-
-    See Also
-    --------
-    sort_by_key : Computes the permutation to gather with.
-    compact_by_mask : Filters in place instead.
-
-    Examples
-    --------
-    >>> a = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    >>> xp.algorithms.take_columns(a, xp.asarray([2, 0]))
-    array([[3., 1.],
-           [6., 4.]])
-    >>> xp.algorithms.take_columns(a, xp.asarray([True, False, True]))
-    array([[1., 3.],
-           [4., 6.]])
-    """
-    assert_same_backend(array, index, *(() if out is None else (out,)))
-    xpm = get_array_module(array)
-    array = xpm.asarray(array)
-    if array.ndim != 2:
-        raise ValueError(f"array must be 2D (ncomp, n), got shape {array.shape}")
-    ncomp, n = array.shape
-    columns = _checked_columns(index, n, xpm)
-    shape = (ncomp, columns.shape[0])
-    if out is not None:
-        if out.shape != shape or out.dtype != array.dtype or not out.flags.c_contiguous:
-            raise ValueError(
-                f"out must have shape {shape}, dtype {array.dtype} and be C-contiguous",
-            )
-        if not getattr(out.flags, "writeable", True):
-            raise ValueError("out must be writable")
-        if xpm.may_share_memory(out, array):
-            raise ValueError("out must not overlap array")
-    if get_array_backend(array) == "cupy":
-        # CuPy's take wraps out-of-range indices, negative ones as NumPy does
-        return xpm.take(array, columns, axis=1, out=out, mode="wrap")
-    result = np.empty(shape, dtype=array.dtype) if out is None else out
-    if result.size == 0:
-        return result
-    low, high = int(columns.min()), int(columns.max())
-    if high >= n or low < -n:
-        raise IndexError(
-            f"column index out of bounds for {n} columns "
-            f"(indices from {low} to {high})",
-        )
-    if low < 0:
-        columns = np.where(columns < 0, columns + n, columns)
-    if (
-        array.dtype in _TAKE_COLUMNS_DTYPES
-        and array.strides[1] == array.itemsize
-        and columns.shape[0] >= _TAKE_COLUMNS_MIN_COLUMNS
-    ):
-        from cunumpy._kernels.take_columns import kernel
-
-        kernel(array, np.ascontiguousarray(columns, dtype=np.int64), result)
-    else:
-        from cunumpy._kernels.take_columns.take_columns_numpy import (
-            take_columns as take_numpy,
-        )
-
-        take_numpy(array, columns, result)
-    return result
diff --git a/src/cunumpy/_kernels/__init__.py b/src/cunumpy/_kernels/__init__.py
deleted file mode 100644
index 2f1c2fb..0000000
--- a/src/cunumpy/_kernels/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-"""Host kernels of cunumpy's own algorithms, one folder per kernel."""
diff --git a/src/cunumpy/_kernels/take_columns/__init__.py b/src/cunumpy/_kernels/take_columns/__init__.py
deleted file mode 100644
index bb73bb7..0000000
--- a/src/cunumpy/_kernels/take_columns/__init__.py
+++ /dev/null
@@ -1,11 +0,0 @@
-"""The ``take_columns`` kernel: Pyccel and NumPy versions with the same arguments."""
-
-from cunumpy._dispatch import Kernel
-from cunumpy._pyccel_cache import PyccelCache
-
-kernel = Kernel.from_folder(
-    __name__,
-    host_suffix="_pyccel",
-    dispatch="arrays",
-    compile_host=PyccelCache("cunumpy"),
-)
diff --git a/src/cunumpy/_kernels/take_columns/take_columns_numpy.py b/src/cunumpy/_kernels/take_columns/take_columns_numpy.py
deleted file mode 100644
index 834daf7..0000000
--- a/src/cunumpy/_kernels/take_columns/take_columns_numpy.py
+++ /dev/null
@@ -1,20 +0,0 @@
-"""Gather the columns ``order`` of a component-major array, in NumPy."""
-
-import numpy as np
-
-
-def take_columns(source: np.ndarray, order: np.ndarray, target: np.ndarray) -> None:
-    """Write ``source[:, order[i]]`` to ``target[:, i]`` for every entry of ``order``.
-
-    Parameters
-    ----------
-    source : numpy.ndarray
-        The component-major array, ``(ncomp, n)``.
-    order : numpy.ndarray
-        The column indices to gather, ``(m,)``, within ``[0, n)``.
-    target : numpy.ndarray
-        Receives the gathered columns, ``(ncomp, m)``; must not overlap ``source``.
-    """
-    # one contiguous row at a time: faster than a 2-D gather
-    for component in range(source.shape[0]):
-        target[component] = source[component][order]
diff --git a/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py b/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
deleted file mode 100644
index b4df0e2..0000000
--- a/src/cunumpy/_kernels/take_columns/take_columns_pyccel.py
+++ /dev/null
@@ -1,36 +0,0 @@
-"""Gather the columns ``order`` of a component-major array, in parallel (Pyccel)."""
-
-from typing import TypeVar
-
-T = TypeVar(
-    "T",
-    "float[:,:]",
-    "float32[:,:]",  # noqa: F821
-    "int[:,:]",
-    "int32[:,:]",  # noqa: F821
-    "bool[:,:]",
-)
-
-
-def take_columns(source: "T", order: "int[:]", target: "T") -> None:
-    """Write ``source[:, order[i]]`` to ``target[:, i]`` for every entry of ``order``.
-
-    The indices must lie in ``[0, source.shape[1])`` (the caller checks them) and
-    ``target`` must not overlap ``source``. One loop per component, so every
-    thread streams through one contiguous row.
-
-    Parameters
-    ----------
-    source : array of shape (ncomp, n)
-        The component-major array.
-    order : array of int, shape (m,)
-        The column indices to gather.
-    target : array of shape (ncomp, m)
-        Receives the gathered columns.
-    """
-    for component in range(source.shape[0]):
-        # fmt: off
-        #$ omp parallel for schedule(static)
-        # fmt: on
-        for marker in range(order.shape[0]):
-            target[component, marker] = source[component, order[marker]]
diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py
index e371bea..c13cebc 100644
--- a/src/cunumpy/algorithms.py
+++ b/src/cunumpy/algorithms.py
@@ -2,14 +2,13 @@
 
 Morton (Z-order) keys (the same as ``cunumpy/morton.cuh`` computes in a
 kernel, at most ``MAX_MORTON_LEVELS`` bits per axis), a stable sort of several
-arrays by one key, sums per key, the compaction of the live rows of
-particle arrays and the gather of columns of component-major markers::
+arrays by one key, sums per key, and the compaction of the live rows of
+particle arrays::
 
     keys = xp.algorithms.morton_keys(positions, lower, upper, levels)
     keys, order, positions = xp.algorithms.sort_by_key(keys, positions)
     charge = xp.algorithms.segment_sum(q, cell, n_cells)
     n = xp.algorithms.compact_by_mask(alive, positions, charges)
-    markers = xp.algorithms.take_columns(markers, order)
 
 See :doc:`/guides/particle-codes`.
 """
@@ -21,7 +20,6 @@
     segment_boundaries,
     segment_sum,
     sort_by_key,
-    take_columns,
 )
 from cunumpy._morton import (
     MAX_MORTON_LEVELS,
@@ -43,5 +41,4 @@
     "segment_boundaries",
     "segment_sum",
     "sort_by_key",
-    "take_columns",
 ]
diff --git a/tests/unit/test_morton.py b/tests/unit/test_morton.py
index 3c1bb8a..3ef3d85 100644
--- a/tests/unit/test_morton.py
+++ b/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)
diff --git a/tests/unit/test_take_columns.py b/tests/unit/test_take_columns.py
deleted file mode 100644
index d97414f..0000000
--- a/tests/unit/test_take_columns.py
+++ /dev/null
@@ -1,217 +0,0 @@
-"""take_columns gathers columns of (ncomp, n) arrays like array[:, index], C-contiguous."""
-
-import numpy as np
-import pytest
-
-import cunumpy as xp
-from cunumpy import _algorithms
-from cunumpy.kernels import use_host_kernel_implementation
-
-DTYPES = [np.float64, np.float32, np.int64, np.int32, np.bool_, np.complex128]
-IMPLEMENTATIONS = ["default", "numpy"]
-
-
-def make(shape, dtype, seed=0):
-    rng = np.random.default_rng(seed)
-    return (rng.random(shape) * 100).astype(dtype)
-
-
-def run(implementation, *args, **kwargs):
-    if implementation == "default":
-        return xp.algorithms.take_columns(*args, **kwargs)
-    with use_host_kernel_implementation(implementation):
-        return xp.algorithms.take_columns(*args, **kwargs)
-
-
-@pytest.fixture(params=IMPLEMENTATIONS)
-def implementation(request, monkeypatch):
-    # run the compiled kernel on the small test arrays too
-    monkeypatch.setattr(_algorithms, "_TAKE_COLUMNS_MIN_COLUMNS", 0)
-    return request.param
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("shape", [(1, 1), (3, 7), (8, 1000), (3, 5000)])
-def test_integer_index_matches_fancy_indexing(implementation, dtype, shape):
-    array = make(shape, dtype)
-    index = np.random.default_rng(1).integers(0, shape[1], 2 * shape[1] // 3 + 1)
-    result = run(implementation, array, index)
-    assert result.dtype == array.dtype
-    assert result.flags.c_contiguous
-    np.testing.assert_array_equal(result, array[:, index])
-
-
-def test_permutation_of_a_large_array(implementation):
-    array = make((3, 100_000), np.float64)
-    order = np.random.default_rng(2).permutation(100_000)
-    np.testing.assert_array_equal(run(implementation, array, order), array[:, order])
-
-
-@pytest.mark.parametrize("index_dtype", [np.int32, np.int64, np.uint32, np.uint64])
-def test_index_dtypes(implementation, index_dtype):
-    array = make((2, 50), np.float64)
-    index = np.array([4, 0, 49, 4], dtype=index_dtype)
-    np.testing.assert_array_equal(run(implementation, array, index), array[:, index])
-
-
-def test_negative_indices_count_from_the_end(implementation):
-    array = make((2, 10), np.float64)
-    index = np.array([-1, 0, -10, 3])
-    np.testing.assert_array_equal(run(implementation, array, index), array[:, index])
-
-
-@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.bool_])
-def test_strided_column_prefix_view(implementation, dtype):
-    storage = make((3, 40), dtype)
-    view = storage[:, :25]  # the rows are further apart than the view is wide
-    index = np.random.default_rng(3).permutation(25)
-    np.testing.assert_array_equal(run(implementation, view, index), view[:, index])
-
-
-def test_non_row_contiguous_input(implementation):
-    base = make((9, 6), np.float64)
-    for array in (base.T, base[:, ::2], np.asfortranarray(base)):
-        index = np.arange(array.shape[1])[::-1]
-        np.testing.assert_array_equal(
-            run(implementation, array, index), array[:, index]
-        )
-
-
-def test_read_only_input(implementation):
-    array = make((2, 20), np.float64)
-    array.flags.writeable = False
-    index = np.arange(20)[::-1]
-    np.testing.assert_array_equal(run(implementation, array, index), array[:, index])
-
-
-def test_empty_index_and_empty_array(implementation):
-    array = make((3, 5), np.float64)
-    result = run(implementation, array, np.array([], dtype=np.int64))
-    assert result.shape == (3, 0)
-    assert result.dtype == array.dtype
-    assert run(implementation, np.zeros((0, 4)), np.array([1, 2])).shape == (0, 2)
-    assert run(implementation, np.zeros((3, 0)), np.array([], dtype=int)).shape == (
-        3,
-        0,
-    )
-    assert run(implementation, array, np.zeros(5, bool)).shape == (3, 0)
-
-
-def test_out_is_filled_and_returned(implementation):
-    array = make((3, 30), np.float64)
-    index = np.random.default_rng(4).permutation(30)[:20]
-    out = np.full((3, 20), -1.0)
-    result = run(implementation, array, index, out=out)
-    assert result is out
-    np.testing.assert_array_equal(out, array[:, index])
-
-
-def test_out_may_be_a_row_block_of_a_larger_array(implementation):
-    array = make((2, 10), np.float64)
-    storage = np.zeros((4, 10))
-    out = storage[:2]  # C-contiguous
-    run(implementation, array, np.arange(10)[::-1], out=out)
-    np.testing.assert_array_equal(storage[:2], array[:, ::-1])
-    assert not storage[2:].any()
-
-
-@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.bool_])
-@pytest.mark.parametrize("shape", [(3, 1), (3, 9), (8, 5000)])
-def test_boolean_mask_matches_compress(implementation, dtype, shape):
-    array = make(shape, dtype)
-    mask = np.random.default_rng(5).random(shape[1]) < 0.4
-    result = run(implementation, array, mask)
-    assert result.flags.c_contiguous
-    np.testing.assert_array_equal(result, np.compress(mask, array, axis=1))
-
-
-def test_mask_of_a_strided_view_with_out(implementation):
-    storage = make((3, 40), np.float64)
-    view = storage[:, :30]
-    mask = np.arange(30) % 3 == 0
-    out = np.empty((3, 10))
-    run(implementation, view, mask, out=out)
-    np.testing.assert_array_equal(out, view[:, mask])
-
-
-def test_the_kernel_is_used_from_the_minimum_number_of_columns(monkeypatch):
-    import cunumpy._kernels.take_columns as folder
-
-    calls = []
-    monkeypatch.setattr(folder, "kernel", lambda *args: calls.append(args))
-    size = _algorithms._TAKE_COLUMNS_MIN_COLUMNS
-    array = make((2, size), np.float64)
-    xp.algorithms.take_columns(array, np.arange(size - 1))
-    assert not calls
-    xp.algorithms.take_columns(array, np.arange(size))
-    assert len(calls) == 1
-
-
-def test_error_arguments():
-    array = np.zeros((2, 4))
-    with pytest.raises(ValueError, match="2D"):
-        xp.algorithms.take_columns(np.zeros(4), np.array([0]))
-    with pytest.raises(ValueError, match="1D"):
-        xp.algorithms.take_columns(array, np.zeros((2, 2), dtype=int))
-    with pytest.raises(TypeError, match="integer or boolean"):
-        xp.algorithms.take_columns(array, np.array([0.0]))
-    with pytest.raises(ValueError, match="one entry per column"):
-        xp.algorithms.take_columns(array, np.ones(3, bool))
-    with pytest.raises(IndexError, match="out of bounds"):
-        xp.algorithms.take_columns(array, np.array([0, 4]))
-    with pytest.raises(IndexError, match="out of bounds"):
-        xp.algorithms.take_columns(array, np.array([-5]))
-
-
-def test_out_errors():
-    array = np.zeros((2, 4))
-    index = np.array([0, 1])
-    with pytest.raises(ValueError, match="shape"):
-        xp.algorithms.take_columns(array, index, out=np.zeros((2, 3)))
-    with pytest.raises(ValueError, match="dtype"):
-        xp.algorithms.take_columns(array, index, out=np.zeros((2, 2), dtype=np.int64))
-    with pytest.raises(ValueError, match="C-contiguous"):
-        xp.algorithms.take_columns(array, index, out=np.zeros((2, 4))[:, ::2])
-    read_only = np.zeros((2, 2))
-    read_only.flags.writeable = False
-    with pytest.raises(ValueError, match="writable"):
-        xp.algorithms.take_columns(array, index, out=read_only)
-    with pytest.raises(ValueError, match="overlap"):
-        xp.algorithms.take_columns(array, np.arange(4), out=array)
-
-
-def test_compiled_kernel_agrees_with_numpy_implementation(monkeypatch):
-    monkeypatch.setattr(_algorithms, "_TAKE_COLUMNS_MIN_COLUMNS", 0)
-    from cunumpy._kernels.take_columns import kernel
-
-    pytest.importorskip("pyccel")
-    if kernel.selected() != "pyccel":
-        pytest.skip(
-            "the compiled kernel is not available (needs Pyccel and a compiler)"
-        )
-    array = make((3, 3000), np.float64)
-    index = np.random.default_rng(6).permutation(3000)
-    with use_host_kernel_implementation("pyccel"):
-        compiled = xp.algorithms.take_columns(array, index)
-    with use_host_kernel_implementation("numpy"):
-        reference = xp.algorithms.take_columns(array, index)
-    np.testing.assert_array_equal(compiled, reference)
-
-
-def test_cupy_backend_matches_numpy():
-    if not xp.cupy_available():
-        pytest.skip("CuPy not installed or not functional")
-    array = make((3, 200), np.float64)
-    index = np.random.default_rng(7).permutation(200)[:150]
-    mask = index.size and np.random.default_rng(8).random(200) < 0.5
-    with xp.use_backend("cupy"):
-        device = xp.to_cupy(array)
-        taken = xp.algorithms.take_columns(device, xp.to_cupy(index))
-        masked = xp.algorithms.take_columns(device, xp.to_cupy(mask))
-        out = xp.zeros((3, 150))
-        xp.algorithms.take_columns(device, xp.to_cupy(index), out=out)
-        negative = xp.algorithms.take_columns(device, xp.to_cupy(np.array([-1, 0])))
-    np.testing.assert_array_equal(xp.to_numpy(taken), array[:, index])
-    np.testing.assert_array_equal(xp.to_numpy(masked), array[:, mask])
-    np.testing.assert_array_equal(xp.to_numpy(out), array[:, index])
-    np.testing.assert_array_equal(xp.to_numpy(negative), array[:, [-1, 0]])