From 6137c77ddc30528e93a68c498b312d9abb45ad11 Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Sat, 4 Jul 2026 10:42:33 -0700 Subject: [PATCH 1/4] group_ordering implementation --- performance/__main__.py | 22 ++++ src/__init__.py | 1 + src/__init__.pyi | 3 + src/_arraykit.c | 4 + src/methods.c | 127 ++++++++++++++++++++++ src/methods.h | 3 + test/test_group_ordering.py | 204 ++++++++++++++++++++++++++++++++++++ 7 files changed, 364 insertions(+) create mode 100644 test/test_group_ordering.py diff --git a/performance/__main__.py b/performance/__main__.py index 537037f9..652f09ef 100644 --- a/performance/__main__.py +++ b/performance/__main__.py @@ -49,6 +49,7 @@ from arraykit import split_after_count as split_after_count_ak from arraykit import count_iteration as count_iteration_ak from arraykit import slice_to_ascending_slice as slice_to_ascending_slice_ak +from arraykit import group_ordering as group_ordering_ak from arraykit import ArrayGO as ArrayGOAK @@ -845,6 +846,27 @@ class SliceToAscendingREF(SliceToAscending): entry = staticmethod(slice_to_ascending_slice_ref) +# ------------------------------------------------------------------------------- +class GroupOrdering(Perf): + NUMBER = 20 + + def __init__(self): + rng = np.random.default_rng(0) + # 1M rows over ~5k dense groups, shuffled (mimics factorize codes) + self.codes = rng.integers(0, 5_000, size=1_000_000).astype(np.intp) + + def main(self): + _ = self.entry(self.codes) + + +class GroupOrderingAK(GroupOrdering): + entry = staticmethod(group_ordering_ak) + + +class GroupOrderingREF(GroupOrdering): + entry = staticmethod(lambda codes: np.argsort(codes, kind='stable')) + + # ------------------------------------------------------------------------------- diff --git a/src/__init__.py b/src/__init__.py index 65289266..1fc933f5 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -26,6 +26,7 @@ from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_screen from ._arraykit import write_array_to_file as write_array_to_file from ._arraykit import factorize as factorize +from ._arraykit import group_ordering as group_ordering from ._arraykit import count_iteration as count_iteration from ._arraykit import first_true_1d as first_true_1d from ._arraykit import first_true_2d as first_true_2d diff --git a/src/__init__.pyi b/src/__init__.pyi index 8065bcf2..85304f9c 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -227,6 +227,9 @@ def write_array_to_file( def factorize( array: np.ndarray, *, sort: bool = ... ) -> tp.Tuple[np.ndarray, np.ndarray]: ... +def group_ordering( + codes: np.ndarray, *, size: tp.Optional[int] = ... +) -> tp.Tuple[np.ndarray, np.ndarray]: ... def first_true_1d(__array: np.ndarray, *, forward: bool) -> int: ... def first_true_2d(__array: np.ndarray, *, forward: bool, axis: int) -> np.ndarray: ... def nonzero_1d(__array: np.ndarray, /) -> np.ndarray: ... diff --git a/src/_arraykit.c b/src/_arraykit.c index d04315e6..db80aca4 100644 --- a/src/_arraykit.c +++ b/src/_arraykit.c @@ -74,6 +74,10 @@ static PyMethodDef arraykit_methods[] = { (PyCFunction)factorize, METH_VARARGS | METH_KEYWORDS, NULL}, + {"group_ordering", + (PyCFunction)group_ordering, + METH_VARARGS | METH_KEYWORDS, + NULL}, {NULL}, }; diff --git a/src/methods.c b/src/methods.c index b06742b5..ede556cf 100644 --- a/src/methods.c +++ b/src/methods.c @@ -985,6 +985,133 @@ first_true_2d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) return (PyObject *)array_pos; } +static char *group_ordering_kwarg_names[] = { + "codes", + "size", + NULL +}; + +// Stable counting sort of dense factorize codes. Given `codes` in [0, size), +// return (permutation, offsets) such that permutation[offsets[g]:offsets[g+1]] +// are the input positions of group g, in ascending (stable) order. This is an +// O(n) alternative to np.argsort(codes, kind='stable') for already-dense codes. +PyObject * +group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) +{ + PyArrayObject *codes = NULL; + PyObject *size_obj = NULL; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, + "O!|$O:group_ordering", + group_ordering_kwarg_names, + &PyArray_Type, &codes, + &size_obj + )) { + return NULL; + } + if (PyArray_NDIM(codes) != 1) { + PyErr_SetString(PyExc_ValueError, "Array must be 1-dimensional"); + return NULL; + } + if (PyArray_TYPE(codes) != NPY_INTP) { + PyErr_SetString(PyExc_ValueError, "Array must be of type intp"); + return NULL; + } + if (!PyArray_IS_C_CONTIGUOUS(codes)) { + PyErr_SetString(PyExc_ValueError, "Array must be contiguous"); + return NULL; + } + + npy_intp n = PyArray_SIZE(codes); + npy_intp *codes_buffer = (npy_intp*)PyArray_DATA(codes); + + // Determine the number of groups: caller-provided, else max(codes) + 1. + npy_intp size = 0; + if (size_obj != NULL && size_obj != Py_None) { + size = (npy_intp)PyNumber_AsSsize_t(size_obj, PyExc_OverflowError); + if (size == -1 && PyErr_Occurred()) { + return NULL; + } + if (size < 0) { + PyErr_SetString(PyExc_ValueError, "size must be non-negative"); + return NULL; + } + } + else { + for (npy_intp i = 0; i < n; i++) { + npy_intp c = codes_buffer[i]; + if (c < 0) { + PyErr_SetString(PyExc_ValueError, "codes must be non-negative"); + return NULL; + } + if (c + 1 > size) { + size = c + 1; + } + } + } + + PyObject *perm_arr = NULL; + PyObject *offsets_arr = NULL; + npy_intp *cursor = NULL; + + perm_arr = PyArray_EMPTY(1, &n, NPY_INTP, 0); + if (!perm_arr) { + goto fail; + } + npy_intp size_plus = size + 1; + offsets_arr = PyArray_ZEROS(1, &size_plus, NPY_INTP, 0); + if (!offsets_arr) { + goto fail; + } + npy_intp *perm = (npy_intp*)PyArray_DATA((PyArrayObject*)perm_arr); + npy_intp *offsets = (npy_intp*)PyArray_DATA((PyArrayObject*)offsets_arr); + + // Count pass: tally each group into offsets[c + 1], validating the range. + for (npy_intp i = 0; i < n; i++) { + npy_intp c = codes_buffer[i]; + if (c < 0 || c >= size) { + PyErr_Format(PyExc_ValueError, + "code %zd out of range [0, %zd)", + (Py_ssize_t)c, (Py_ssize_t)size); + goto fail; + } + offsets[c + 1]++; + } + // Prefix sum: offsets[g] becomes the start index of group g; offsets[size] == n. + for (npy_intp g = 0; g < size; g++) { + offsets[g + 1] += offsets[g]; + } + // Scatter pass: place each input position at its group's running cursor. + // Ascending i preserves original order within each group (stability). + if (size > 0) { + cursor = PyMem_New(npy_intp, size); + if (!cursor) { + PyErr_NoMemory(); + goto fail; + } + memcpy(cursor, offsets, size * sizeof(npy_intp)); + for (npy_intp i = 0; i < n; i++) { + perm[cursor[codes_buffer[i]]++] = i; + } + } + + PyArray_CLEARFLAGS((PyArrayObject*)perm_arr, NPY_ARRAY_WRITEABLE); + PyArray_CLEARFLAGS((PyArrayObject*)offsets_arr, NPY_ARRAY_WRITEABLE); + + PyMem_Free(cursor); + + PyObject *result = PyTuple_Pack(2, perm_arr, offsets_arr); + Py_DECREF(perm_arr); + Py_DECREF(offsets_arr); + return result; + +fail: + PyMem_Free(cursor); + Py_XDECREF(perm_arr); + Py_XDECREF(offsets_arr); + return NULL; +} + PyObject * dtype_from_element(PyObject *Py_UNUSED(m), PyObject *arg) { diff --git a/src/methods.h b/src/methods.h index 92e0adc3..574c7bdd 100644 --- a/src/methods.h +++ b/src/methods.h @@ -69,6 +69,9 @@ first_true_1d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); PyObject * first_true_2d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); +PyObject * +group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); + PyObject * dtype_from_element(PyObject *Py_UNUSED(m), PyObject *arg); diff --git a/test/test_group_ordering.py b/test/test_group_ordering.py new file mode 100644 index 00000000..86a2a7b4 --- /dev/null +++ b/test/test_group_ordering.py @@ -0,0 +1,204 @@ +import unittest + +import numpy as np + +from arraykit import group_ordering +from arraykit import factorize + + +def offsets_from_codes(codes, size): + # CSR-style offsets: [0, *cumsum(bincount(codes, minlength=size))] + counts = np.bincount(codes, minlength=size) + return np.concatenate([[0], np.cumsum(counts)]).astype(np.intp) + + +class TestUnit(unittest.TestCase): + # ------------------------------------------------------------------ + # basic behavior + + def test_group_ordering_basic_a(self) -> None: + codes = np.array([0, 0, 0, 1, 1, 2], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm.tolist(), [0, 1, 2, 3, 4, 5]) + self.assertEqual(offsets.tolist(), [0, 3, 5, 6]) + + def test_group_ordering_interleaved(self) -> None: + codes = np.array([2, 0, 1, 0, 2, 1], dtype=np.intp) + perm, offsets = group_ordering(codes) + # group 0 -> positions 1, 3; group 1 -> 2, 5; group 2 -> 0, 4 + self.assertEqual(offsets.tolist(), [0, 2, 4, 6]) + self.assertEqual(perm[offsets[0]:offsets[1]].tolist(), [1, 3]) + self.assertEqual(perm[offsets[1]:offsets[2]].tolist(), [2, 5]) + self.assertEqual(perm[offsets[2]:offsets[3]].tolist(), [0, 4]) + + def test_group_ordering_stability(self) -> None: + # original positions within each group must stay ascending + codes = np.array([0, 1, 0, 1, 0, 1], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm[offsets[0]:offsets[1]].tolist(), [0, 2, 4]) + self.assertEqual(perm[offsets[1]:offsets[2]].tolist(), [1, 3, 5]) + + # ------------------------------------------------------------------ + # parity against numpy oracle + + def test_group_ordering_parity_argsort(self) -> None: + rng = np.random.default_rng(0) + for size in (1, 5, 50, 500): + codes = rng.integers(0, size, size=10_000).astype(np.intp) + perm, offsets = group_ordering(codes) + expected = np.argsort(codes, kind='stable').astype(np.intp) + self.assertEqual(perm.tolist(), expected.tolist()) + self.assertEqual( + offsets.tolist(), offsets_from_codes(codes, size).tolist() + ) + + def test_group_ordering_parity_inferred_size(self) -> None: + rng = np.random.default_rng(1) + codes = rng.integers(0, 100, size=5_000).astype(np.intp) + perm, offsets = group_ordering(codes) + size = int(codes.max()) + 1 + self.assertEqual(len(offsets), size + 1) + expected = np.argsort(codes, kind='stable').astype(np.intp) + self.assertEqual(perm.tolist(), expected.tolist()) + + # ------------------------------------------------------------------ + # dtype / shape + + def test_group_ordering_dtypes(self) -> None: + codes = np.array([0, 1, 0, 2], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm.dtype, np.dtype(np.intp)) + self.assertEqual(offsets.dtype, np.dtype(np.intp)) + + def test_group_ordering_offsets_length(self) -> None: + codes = np.array([0, 1, 2, 3], dtype=np.intp) + _, offsets = group_ordering(codes, size=10) + self.assertEqual(len(offsets), 11) + self.assertEqual(offsets[-1], 4) + + # ------------------------------------------------------------------ + # size keyword + + def test_group_ordering_size_explicit(self) -> None: + codes = np.array([0, 0, 1, 1], dtype=np.intp) + perm, offsets = group_ordering(codes, size=2) + self.assertEqual(perm.tolist(), [0, 1, 2, 3]) + self.assertEqual(offsets.tolist(), [0, 2, 4]) + + def test_group_ordering_size_trailing_empty(self) -> None: + codes = np.array([0, 0, 1], dtype=np.intp) + _, offsets = group_ordering(codes, size=4) + # groups 2 and 3 are empty: offsets[g] == offsets[g+1] + self.assertEqual(offsets.tolist(), [0, 2, 3, 3, 3]) + + def test_group_ordering_size_none(self) -> None: + codes = np.array([0, 1, 1], dtype=np.intp) + perm, offsets = group_ordering(codes, size=None) + self.assertEqual(offsets.tolist(), [0, 1, 3]) + + def test_group_ordering_size_is_keyword_only(self) -> None: + codes = np.array([0, 1], dtype=np.intp) + with self.assertRaises(TypeError): + group_ordering(codes, 2) + + # ------------------------------------------------------------------ + # edge cases + + def test_group_ordering_empty(self) -> None: + codes = np.array([], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm.tolist(), []) + self.assertEqual(offsets.tolist(), [0]) + + def test_group_ordering_empty_with_size(self) -> None: + codes = np.array([], dtype=np.intp) + perm, offsets = group_ordering(codes, size=3) + self.assertEqual(perm.tolist(), []) + self.assertEqual(offsets.tolist(), [0, 0, 0, 0]) + + def test_group_ordering_single(self) -> None: + codes = np.array([0], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm.tolist(), [0]) + self.assertEqual(offsets.tolist(), [0, 1]) + + def test_group_ordering_single_group(self) -> None: + codes = np.array([0, 0, 0], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm.tolist(), [0, 1, 2]) + self.assertEqual(offsets.tolist(), [0, 3]) + + def test_group_ordering_all_distinct(self) -> None: + codes = np.array([3, 2, 1, 0], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm.tolist(), [3, 2, 1, 0]) + self.assertEqual(offsets.tolist(), [0, 1, 2, 3, 4]) + + # ------------------------------------------------------------------ + # validation + + def test_group_ordering_not_array(self) -> None: + with self.assertRaises(TypeError): + group_ordering([0, 1, 2]) + + def test_group_ordering_2d(self) -> None: + codes = np.array([[0, 1], [1, 0]], dtype=np.intp) + with self.assertRaises(ValueError): + group_ordering(codes) + + def test_group_ordering_wrong_dtype(self) -> None: + codes = np.array([0, 1, 2], dtype=np.int32) + with self.assertRaises(ValueError): + group_ordering(codes) + + def test_group_ordering_non_contiguous(self) -> None: + codes = np.arange(10, dtype=np.intp)[::2] + self.assertFalse(codes.flags['C_CONTIGUOUS']) + with self.assertRaises(ValueError): + group_ordering(codes) + + def test_group_ordering_negative_code_inferred(self) -> None: + codes = np.array([0, -1, 1], dtype=np.intp) + with self.assertRaises(ValueError): + group_ordering(codes) + + def test_group_ordering_out_of_range(self) -> None: + codes = np.array([0, 1, 5], dtype=np.intp) + with self.assertRaises(ValueError): + group_ordering(codes, size=3) + + def test_group_ordering_negative_size(self) -> None: + codes = np.array([0, 1], dtype=np.intp) + with self.assertRaises(ValueError): + group_ordering(codes, size=-1) + + def test_group_ordering_size_out_of_range_zero(self) -> None: + codes = np.array([0, 1], dtype=np.intp) + with self.assertRaises(ValueError): + group_ordering(codes, size=0) + + # ------------------------------------------------------------------ + # immutability + + def test_group_ordering_outputs_immutable(self) -> None: + codes = np.array([0, 1, 0], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertFalse(perm.flags.writeable) + self.assertFalse(offsets.flags.writeable) + + # ------------------------------------------------------------------ + # round-trip with factorize + + def test_group_ordering_with_factorize(self) -> None: + a = np.array(['b', 'a', 'b', 'c', 'a', 'a']) + uniques, codes = factorize(a) + perm, offsets = group_ordering(codes, size=len(uniques)) + ordered = a[perm] + # each group's slice of the reordered array is constant + for g in range(len(uniques)): + segment = ordered[offsets[g]:offsets[g + 1]] + self.assertTrue((segment == segment[0]).all()) + + +if __name__ == '__main__': + unittest.main() From 92edb9e00a93c28d0d9be36971577147206193b8 Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Sat, 4 Jul 2026 11:07:10 -0700 Subject: [PATCH 2/4] fixed test assumptions --- test/test_group_ordering.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/test/test_group_ordering.py b/test/test_group_ordering.py index 86a2a7b4..61d045e3 100644 --- a/test/test_group_ordering.py +++ b/test/test_group_ordering.py @@ -147,7 +147,15 @@ def test_group_ordering_2d(self) -> None: group_ordering(codes) def test_group_ordering_wrong_dtype(self) -> None: - codes = np.array([0, 1, 2], dtype=np.int32) + # pick an integer width that differs from intp on this platform + # (intp is 32-bit on some Windows builds, 64-bit elsewhere) + wrong = np.int32 if np.dtype(np.intp).itemsize != 4 else np.int64 + codes = np.array([0, 1, 2], dtype=wrong) + with self.assertRaises(ValueError): + group_ordering(codes) + + def test_group_ordering_wrong_dtype_float(self) -> None: + codes = np.array([0.0, 1.0, 2.0], dtype=np.float64) with self.assertRaises(ValueError): group_ordering(codes) From 59951e5c67981d0942d274dd5a11a9d917e39af5 Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Sat, 4 Jul 2026 12:55:08 -0700 Subject: [PATCH 3/4] handling overflows --- performance/__main__.py | 10 +++++++++- src/methods.c | 13 +++++++++++++ test/test_group_ordering.py | 12 ++++++++++++ 3 files changed, 34 insertions(+), 1 deletion(-) diff --git a/performance/__main__.py b/performance/__main__.py index 652f09ef..6a7f7c36 100644 --- a/performance/__main__.py +++ b/performance/__main__.py @@ -859,12 +859,20 @@ def main(self): _ = self.entry(self.codes) +def group_ordering_ref(codes): + # match group_ordering's contract: return both the stable permutation and + # the CSR-style group offsets (length size + 1) + perm = np.argsort(codes, kind='stable') + offsets = np.concatenate([[0], np.cumsum(np.bincount(codes))]).astype(np.intp) + return perm, offsets + + class GroupOrderingAK(GroupOrdering): entry = staticmethod(group_ordering_ak) class GroupOrderingREF(GroupOrdering): - entry = staticmethod(lambda codes: np.argsort(codes, kind='stable')) + entry = staticmethod(group_ordering_ref) # ------------------------------------------------------------------------------- diff --git a/src/methods.c b/src/methods.c index ede556cf..f84369cc 100644 --- a/src/methods.c +++ b/src/methods.c @@ -1044,12 +1044,25 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) PyErr_SetString(PyExc_ValueError, "codes must be non-negative"); return NULL; } + // guard c + 1 against signed overflow (undefined behavior) + if (c == NPY_MAX_INTP) { + PyErr_SetString(PyExc_OverflowError, + "cannot infer size: code value too large"); + return NULL; + } if (c + 1 > size) { size = c + 1; } } } + // offsets has length size + 1; guard that against signed overflow (covers + // both a caller-provided size and an inferred one) + if (size == NPY_MAX_INTP) { + PyErr_SetString(PyExc_OverflowError, "size too large"); + return NULL; + } + PyObject *perm_arr = NULL; PyObject *offsets_arr = NULL; npy_intp *cursor = NULL; diff --git a/test/test_group_ordering.py b/test/test_group_ordering.py index 61d045e3..88d7493b 100644 --- a/test/test_group_ordering.py +++ b/test/test_group_ordering.py @@ -185,6 +185,18 @@ def test_group_ordering_size_out_of_range_zero(self) -> None: with self.assertRaises(ValueError): group_ordering(codes, size=0) + def test_group_ordering_infer_overflow(self) -> None: + # a code at the intp max would overflow when inferring size = c + 1 + codes = np.array([np.iinfo(np.intp).max], dtype=np.intp) + with self.assertRaises(OverflowError): + group_ordering(codes) + + def test_group_ordering_size_overflow(self) -> None: + # an explicit size at the intp max would overflow computing size + 1 + codes = np.array([0, 1], dtype=np.intp) + with self.assertRaises(OverflowError): + group_ordering(codes, size=np.iinfo(np.intp).max) + # ------------------------------------------------------------------ # immutability From c44cb7b584fd7f0688728692b9e8d23105e9907d Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Sat, 4 Jul 2026 16:35:42 -0700 Subject: [PATCH 4/4] added test --- src/methods.c | 9 ++++++--- test/test_group_ordering.py | 7 +++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/methods.c b/src/methods.c index f84369cc..746846d7 100644 --- a/src/methods.c +++ b/src/methods.c @@ -1027,7 +1027,8 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) // Determine the number of groups: caller-provided, else max(codes) + 1. npy_intp size = 0; - if (size_obj != NULL && size_obj != Py_None) { + int size_given = (size_obj != NULL && size_obj != Py_None); + if (size_given) { size = (npy_intp)PyNumber_AsSsize_t(size_obj, PyExc_OverflowError); if (size == -1 && PyErr_Occurred()) { return NULL; @@ -1079,10 +1080,12 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) npy_intp *perm = (npy_intp*)PyArray_DATA((PyArrayObject*)perm_arr); npy_intp *offsets = (npy_intp*)PyArray_DATA((PyArrayObject*)offsets_arr); - // Count pass: tally each group into offsets[c + 1], validating the range. + // Count pass: tally each group into offsets[c + 1]. When size was inferred + // the codes are already known to be in [0, size); only a caller-provided + // size needs the range validated here. for (npy_intp i = 0; i < n; i++) { npy_intp c = codes_buffer[i]; - if (c < 0 || c >= size) { + if (size_given && (c < 0 || c >= size)) { PyErr_Format(PyExc_ValueError, "code %zd out of range [0, %zd)", (Py_ssize_t)c, (Py_ssize_t)size); diff --git a/test/test_group_ordering.py b/test/test_group_ordering.py index 88d7493b..8a4a927b 100644 --- a/test/test_group_ordering.py +++ b/test/test_group_ordering.py @@ -22,6 +22,12 @@ def test_group_ordering_basic_a(self) -> None: self.assertEqual(perm.tolist(), [0, 1, 2, 3, 4, 5]) self.assertEqual(offsets.tolist(), [0, 3, 5, 6]) + def test_group_ordering_basic_b(self) -> None: + codes = np.array([2, 0, 0, 2, 1, 1, 0, 0, 3, 0], dtype=np.intp) + perm, offsets = group_ordering(codes) + self.assertEqual(perm.tolist(), [1, 2, 6, 7, 9, 4, 5, 0, 3, 8]) + self.assertEqual(offsets.tolist(), [0, 5, 7, 9, 10]) + def test_group_ordering_interleaved(self) -> None: codes = np.array([2, 0, 1, 0, 2, 1], dtype=np.intp) perm, offsets = group_ordering(codes) @@ -38,6 +44,7 @@ def test_group_ordering_stability(self) -> None: self.assertEqual(perm[offsets[0]:offsets[1]].tolist(), [0, 2, 4]) self.assertEqual(perm[offsets[1]:offsets[2]].tolist(), [1, 3, 5]) + # ------------------------------------------------------------------ # parity against numpy oracle