apply comments from review

numpy · seberg · Feb 27, 2025 · Mar 14, 2024 · Mar 14, 2024 · Mar 15, 2024
commit a8e69ff7661a3357f7617a16431b931986b6de2d
diff --git a/numpy/_core/meson.build b/numpy/_core/meson.build
@@ -690,7 +690,8 @@ cpp_args_common = c_args_common + [
 ]
 if cc.get_argument_syntax() != 'msvc'
   cpp_args_common += [
-    '-fno-exceptions',  # no exception support
+    # unique code handles C++ exceptions
+    # '-fno-exceptions',  # no exception support
     '-fno-rtti',  # no runtime type information
   ]
 endif

diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c
@@ -4588,8 +4588,8 @@ static struct PyMethodDef array_module_methods[] = {
         "Give a warning on reload and big warning in sub-interpreters."},
     {"from_dlpack", (PyCFunction)from_dlpack,
         METH_FASTCALL | METH_KEYWORDS, NULL},
-    {"unique_hash",  (PyCFunction)PyArray_Unique, METH_VARARGS,
-     "Collect unique values via a hash map."},
+    {"unique_hash",  (PyCFunction)array_unique,
+        METH_VARARGS, "Collect unique values via a hash map."},
     {NULL, NULL, 0, NULL}                /* sentinel */
 };
 

diff --git a/numpy/_core/src/multiarray/unique.cpp b/numpy/_core/src/multiarray/unique.cpp
@@ -1,5 +1,6 @@
 #define NPY_NO_DEPRECATED_API NPY_API_VERSION
 
+#include <iostream>
 #include <Python.h>
 
 #include <unordered_set>
@@ -43,9 +44,11 @@ PyObject* unique(PyArrayObject *self)
     npy_intp* strideptr,* innersizeptr;
     std::unordered_set<T> hashset;
 
-    iter = NpyIter_New(self, NPY_ITER_READONLY|
-                             NPY_ITER_EXTERNAL_LOOP|
-                             NPY_ITER_REFS_OK,
+    iter = NpyIter_New(self, NPY_ITER_READONLY |
+                             NPY_ITER_EXTERNAL_LOOP |
+                             NPY_ITER_REFS_OK |
+                             NPY_ITER_ZEROSIZE_OK |
+                             NPY_ITER_GROWINNER,
                         NPY_KEEPORDER, NPY_NO_CASTING,
                         NULL);
     // Making sure the iterator is deallocated when the function returns, with
@@ -94,7 +97,8 @@ PyObject* unique(PyArrayObject *self)
         dims, // shape
         NULL, // strides
         res, // data
-        NPY_ARRAY_OUT_ARRAY, // flags
+        // This flag is needed to be able to call .sort on it.
+        NPY_ARRAY_WRITEABLE, // flags
         NULL // obj
     );
 }
@@ -125,14 +129,17 @@ std::unordered_map<int, function_type> unique_funcs = {
 
 
 extern "C" NPY_NO_EXPORT PyObject *
-PyArray_Unique(PyObject *NPY_UNUSED(dummy), PyObject *args)
+array_unique(PyObject *NPY_UNUSED(dummy), PyObject *args)
 {
     /* This is called from Python space, and expects a single numpy array as input.
 
     It then returns a numpy array containing the unique values of the input array.
 
     If the input array is not supported, it returns None.
     */
+    std::cout << "Hello, world!" << std::endl;
+    // this is to allow grabbing the GIL before raising a python exception.
+    NPY_ALLOW_C_API_DEF;
     PyArrayObject *self = NULL;
     PyObject *res = NULL;
     if (!PyArg_ParseTuple(args, "O&", PyArray_Converter, &self))
@@ -143,31 +150,49 @@ PyArray_Unique(PyObject *NPY_UNUSED(dummy), PyObject *args)
     auto self_decref = finally([&]() { Py_XDECREF(self); });
 
 
-    /* Handle zero-sized arrays specially */
-    if (PyArray_SIZE(self) == 0) {
-        return PyArray_NewLikeArray(
-            self,
-            NPY_ANYORDER,
-            NULL, // descr (use prototype's descr)
-            1 // subok
-        );
-    }
+    try {
+        /* Handle zero-sized arrays specially */
+        if (PyArray_SIZE(self) == 0) {
+            return PyArray_NewLikeArray(
+                self,
+                NPY_ANYORDER,
+                NULL, // descr (use prototype's descr)
+                1 // subok
+            );
+        }
 
-    auto type = PyArray_TYPE(self);
-    // we only support data types present in our unique_funcs map
-    if (unique_funcs.find(type) == unique_funcs.end()) {
-        return Py_None;
-    }
+        auto type = PyArray_TYPE(self);
+        // we only support data types present in our unique_funcs map
+        if (unique_funcs.find(type) == unique_funcs.end()) {
+            std::cout << "Type not found" << std::endl;
+            return Py_None;
+        }
 
-    res = unique_funcs[type](self);
-    return res;
+        res = unique_funcs[type](self);
+        std::cout << "Bye, world!" << std::endl;
+        return res;
+    }
+    catch (const std::bad_alloc &e) {
+        // Grab GIL before raising an exception
+        NPY_ALLOW_C_API;
+        PyErr_NoMemory();
+        NPY_DISABLE_C_API;
+        return NULL;
+    }
+    catch (const std::exception &e) {
+        // Grab GIL before raising an exception
+        NPY_ALLOW_C_API;
+        PyErr_SetString(PyExc_RuntimeError, e.what());
+        NPY_DISABLE_C_API;
+        return NULL;
+    }
 }
 
 
 
 // The following is to expose the unique function to Python
 static PyMethodDef UniqueMethods[] = {
-    {"unique_hash",  PyArray_Unique, METH_VARARGS,
+    {"unique_hash",  array_unique, METH_VARARGS,
      "Collect unique values via a hash map."},
     {NULL, NULL, 0, NULL}
 };
@@ -190,6 +215,8 @@ PyInit__unique(void)
     if (m == NULL) {
         return m;
     }
+    // TODO: with no-gil we probably want to put a lock on the array so that it's
+    // not changed before we calculate the unique values???
 #if Py_GIL_DISABLED
     // signal this module supports running with the GIL disabled
     PyUnstable_Module_SetGIL(m, Py_MOD_GIL_NOT_USED);

diff --git a/numpy/_core/src/multiarray/unique.h b/numpy/_core/src/multiarray/unique.h
@@ -7,7 +7,7 @@
 extern "C" {
 #endif
 
-PyObject* PyArray_Unique(PyObject *NPY_UNUSED(dummy), PyObject *args);
+PyObject* array_unique(PyObject *NPY_UNUSED(dummy), PyObject *args);
 
 #ifdef __cplusplus
 }

diff --git a/numpy/lib/_arraysetops_impl.py b/numpy/lib/_arraysetops_impl.py
@@ -368,13 +368,19 @@ def _unique1d(ar, return_index=False, return_inverse=False,
     if not optional_indices and not return_counts and not np.ma.is_masked(ar):
         # from numpy._core._multiarray_umath import unique_hash
         from numpy._core._unique import unique_hash
-        hash_unique = unique_hash(ar)
+        # First we convert the array to a numpy array, later we wrap it back
+        # in case it was a subclass of numpy.ndarray.
+        conv = _array_converter(ar)
+        ar_, = conv
+
+        hash_unique = unique_hash(ar_)
         if hash_unique is not None:
             if sorted:
                 hash_unique.sort()
-            return (hash_unique,)
-
+            # We wrap the result back in case it was a subclass of numpy.ndarray.
+            return (conv.wrap(hash_unique),)
 
+    # If we don't use the hash map, we use the slower sorting method.
     if optional_indices:
         perm = ar.argsort(kind='mergesort' if return_index else 'quicksort')
         aux = ar[perm]

diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py
@@ -628,55 +628,71 @@ def test_manyways(self):
 
 class TestUnique:
 
-    def test_unique_1d(self):
+    def check_all(self, a, b, i1, i2, c, dt):
+        base_msg = 'check {0} failed for type {1}'
+
+        msg = base_msg.format('values', dt)
+        v = unique(a)
+        assert_array_equal(v, b, msg)
+        assert type(v) == type(b)
+
+        msg = base_msg.format('return_index', dt)
+        v, j = unique(a, True, False, False)
+        assert_array_equal(v, b, msg)
+        assert_array_equal(j, i1, msg)
+        assert type(v) == type(b)
+
+        msg = base_msg.format('return_inverse', dt)
+        v, j = unique(a, False, True, False)
+        assert_array_equal(v, b, msg)
+        assert_array_equal(j, i2, msg)
+        assert type(v) == type(b)
+
+        msg = base_msg.format('return_counts', dt)
+        v, j = unique(a, False, False, True)
+        assert_array_equal(v, b, msg)
+        assert_array_equal(j, c, msg)
+        assert type(v) == type(b)
+
+        msg = base_msg.format('return_index and return_inverse', dt)
+        v, j1, j2 = unique(a, True, True, False)
+        assert_array_equal(v, b, msg)
+        assert_array_equal(j1, i1, msg)
+        assert_array_equal(j2, i2, msg)
+        assert type(v) == type(b)
+
+        msg = base_msg.format('return_index and return_counts', dt)
+        v, j1, j2 = unique(a, True, False, True)
+        assert_array_equal(v, b, msg)
+        assert_array_equal(j1, i1, msg)
+        assert_array_equal(j2, c, msg)
+        assert type(v) == type(b)
+
+        msg = base_msg.format('return_inverse and return_counts', dt)
+        v, j1, j2 = unique(a, False, True, True)
+        assert_array_equal(v, b, msg)
+        assert_array_equal(j1, i2, msg)
+        assert_array_equal(j2, c, msg)
+        assert type(v) == type(b)
+
+        msg = base_msg.format(('return_index, return_inverse '
+                                'and return_counts'), dt)
+        v, j1, j2, j3 = unique(a, True, True, True)
+        assert_array_equal(v, b, msg)
+        assert_array_equal(j1, i1, msg)
+        assert_array_equal(j2, i2, msg)
+        assert_array_equal(j3, c, msg)
+        assert type(v) == type(b)
+
+    def get_types(self):
+        types = []
+        types.extend(np.typecodes['AllInteger'])
+        types.extend(np.typecodes['AllFloat'])
+        types.append('datetime64[D]')
+        types.append('timedelta64[D]')
+        return types
 
-        def check_all(a, b, i1, i2, c, dt):
-            base_msg = 'check {0} failed for type {1}'
-
-            msg = base_msg.format('values', dt)
-            v = unique(a)
-            assert_array_equal(v, b, msg)
-
-            msg = base_msg.format('return_index', dt)
-            v, j = unique(a, True, False, False)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j, i1, msg)
-
-            msg = base_msg.format('return_inverse', dt)
-            v, j = unique(a, False, True, False)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j, i2, msg)
-
-            msg = base_msg.format('return_counts', dt)
-            v, j = unique(a, False, False, True)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j, c, msg)
-
-            msg = base_msg.format('return_index and return_inverse', dt)
-            v, j1, j2 = unique(a, True, True, False)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j1, i1, msg)
-            assert_array_equal(j2, i2, msg)
-
-            msg = base_msg.format('return_index and return_counts', dt)
-            v, j1, j2 = unique(a, True, False, True)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j1, i1, msg)
-            assert_array_equal(j2, c, msg)
-
-            msg = base_msg.format('return_inverse and return_counts', dt)
-            v, j1, j2 = unique(a, False, True, True)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j1, i2, msg)
-            assert_array_equal(j2, c, msg)
-
-            msg = base_msg.format(('return_index, return_inverse '
-                                   'and return_counts'), dt)
-            v, j1, j2, j3 = unique(a, True, True, True)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j1, i1, msg)
-            assert_array_equal(j2, i2, msg)
-            assert_array_equal(j3, c, msg)
+    def test_unique_1d(self):
 
         a = [5, 7, 1, 2, 1, 5, 7] * 10
         b = [1, 2, 5, 7]
@@ -685,29 +701,25 @@ def check_all(a, b, i1, i2, c, dt):
         c = np.multiply([2, 1, 2, 2], 10)
 
         # test for numeric arrays
-        types = []
-        types.extend(np.typecodes['AllInteger'])
-        types.extend(np.typecodes['AllFloat'])
-        types.append('datetime64[D]')
-        types.append('timedelta64[D]')
+        types = self.get_types()
         for dt in types:
             aa = np.array(a, dt)
             bb = np.array(b, dt)
-            check_all(aa, bb, i1, i2, c, dt)
+            self.check_all(aa, bb, i1, i2, c, dt)
 
         # test for object arrays
         dt = 'O'
         aa = np.empty(len(a), dt)
         aa[:] = a
         bb = np.empty(len(b), dt)
         bb[:] = b
-        check_all(aa, bb, i1, i2, c, dt)
+        self.check_all(aa, bb, i1, i2, c, dt)
 
         # test for structured arrays
         dt = [('', 'i'), ('', 'i')]
         aa = np.array(list(zip(a, a)), dt)
         bb = np.array(list(zip(b, b)), dt)
-        check_all(aa, bb, i1, i2, c, dt)
+        self.check_all(aa, bb, i1, i2, c, dt)
 
         # test for ticket #2799
         aa = [1. + 0.j, 1 - 1.j, 1]
@@ -797,6 +809,33 @@ def check_all(a, b, i1, i2, c, dt):
         assert_equal(np.unique(all_nans, return_inverse=True), (ua, ua_inv))
         assert_equal(np.unique(all_nans, return_counts=True), (ua, ua_cnt))
 
+    def test_unique_zero_sized(self):
+        # test for zero-sized arrays
+        for dt in self.get_types():
+            a = np.array([], dt)
+            b = np.array([], dt)
+            i1 = np.array([], np.int64)
+            i2 = np.array([], np.int64)
+            c = np.array([], np.int64)
+            self.check_all(a, b, i1, i2, c, dt)
+
+    def test_unique_subclass(self):
+        class Subclass(np.ndarray):
+            pass
+
+        i1 = [2, 3, 0, 1]
+        i2 = [2, 3, 0, 1, 0, 2, 3] * 10
+        c = np.multiply([2, 1, 2, 2], 10)
+
+        # test for numeric arrays
+        types = self.get_types()
+        for dt in types:
+            a = np.array([5, 7, 1, 2, 1, 5, 7] * 10, dtype=dt)
+            b = np.array([1, 2, 5, 7], dtype=dt)
+            aa = Subclass(a.shape, dtype=dt, buffer=a)
+            bb = Subclass(b.shape, dtype=dt, buffer=b)
+            self.check_all(aa, bb, i1, i2, c, dt)
+
     def test_unique_axis_errors(self):
         assert_raises(TypeError, self._run_axis_tests, object)
         assert_raises(TypeError, self._run_axis_tests,