Merge pull request #27993 from charris/backport-27992

BUG: Fix segfault in stringdtype lexsort
numpy · Dec 13, 2024 · 42fd84c · 42fd84c
2 parents 852ad5b + 0282fc8
commit 42fd84c
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 3 deletions.
diff --git a/numpy/_core/src/multiarray/item_selection.c b/numpy/_core/src/multiarray/item_selection.c
@@ -2014,8 +2014,7 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
                 }
                 rcode = argsort(its[j]->dataptr,
                         (npy_intp *)rit->dataptr, N, mps[j]);
-                if (rcode < 0 || (PyDataType_REFCHK(PyArray_DESCR(mps[j]))
-                            && PyErr_Occurred())) {
+                if (rcode < 0 || (object && PyErr_Occurred())) {
                     goto fail;
                 }
                 PyArray_ITER_NEXT(its[j]);

diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
@@ -5374,6 +5374,13 @@ def test_object(self):  # gh-6312
             u, v = np.array(u, dtype='object'), np.array(v, dtype='object')
             assert_array_equal(idx, np.lexsort((u, v)))
 
+    def test_strings(self):  # gh-27984
+        for dtype in "TU":
+            surnames = np.array(['Hertz',    'Galilei', 'Hertz'], dtype=dtype)
+            first_names = np.array(['Heinrich', 'Galileo', 'Gustav'], dtype=dtype)
+            assert_array_equal(np.lexsort((first_names, surnames)), [1, 2, 0])
+
+
     def test_invalid_axis(self): # gh-7528
         x = np.linspace(0., 1., 42*3).reshape(42, 3)
         assert_raises(AxisError, np.lexsort, x, axis=2)

diff --git a/numpy/_core/tests/test_stringdtype.py b/numpy/_core/tests/test_stringdtype.py
@@ -415,8 +415,19 @@ def test_sort(dtype, strings):
 
     def test_sort(strings, arr_sorted):
         arr = np.array(strings, dtype=dtype)
-        np.random.default_rng().shuffle(arr)
         na_object = getattr(arr.dtype, "na_object", "")
+        if na_object is None and None in strings:
+            with pytest.raises(
+                ValueError,
+                match="Cannot compare null that is not a nan-like value",
+            ):
+                np.argsort(arr)
+            argsorted = None
+        elif na_object is pd_NA or na_object != '':
+            argsorted = None
+        else:
+            argsorted = np.argsort(arr)
+        np.random.default_rng().shuffle(arr)
         if na_object is None and None in strings:
             with pytest.raises(
                 ValueError,
@@ -426,6 +437,9 @@ def test_sort(strings, arr_sorted):
         else:
             arr.sort()
             assert np.array_equal(arr, arr_sorted, equal_nan=True)
+        if argsorted is not None:
+            assert np.array_equal(argsorted, np.argsort(strings))
+
 
     # make a copy so we don't mutate the lists in the fixture
     strings = strings.copy()