Merge pull request #54 from peytondmurray/ufunc-minmax

ngoldbaum · web-flow · commit 718fea92cd40 · 2023-04-05T08:59:08.000-06:00
Add min and max ufuncs
diff --git a/stringdtype/stringdtype/src/casts.c b/stringdtype/stringdtype/src/casts.c
@@ -52,13 +52,16 @@ string_to_string(PyArrayMethod_Context *NPY_UNUSED(context),
     npy_intp in_stride = strides[0];
     npy_intp out_stride = strides[1];
 
-    ss *s = NULL, *os = NULL;
+    ss *s = NULL;
 
     while (N--) {
-        load_string(in, &s);
-        os = (ss *)out;
-        ssfree(os);
-        if (ssdup(s, os) < 0) {
+        // *out* may be reallocated later; *in->buf* may point to a statically
+        // allocated empty ss struct, so we need to load the string into an
+        // intermediate buffer *s* to avoid the possibility of freeing static
+        // data later on.
+        load_string(in, (ss **)&s);
+        ssfree((ss *)out);
+        if (ssdup((ss *)s, (ss *)out) < 0) {
             gil_error(PyExc_MemoryError, "ssdup failed");
             return -1;
         }
diff --git a/stringdtype/stringdtype/src/dtype.h b/stringdtype/stringdtype/src/dtype.h
@@ -26,6 +26,9 @@ new_stringdtype_instance(void);
 int
 init_string_dtype(void);
 
+int
+compare(void *, void *, void *);
+
 // from dtypemeta.h, not public in numpy
 #define NPY_DTYPE(descr) ((PyArray_DTypeMeta *)Py_TYPE(descr))
 
diff --git a/stringdtype/stringdtype/src/umath.c b/stringdtype/stringdtype/src/umath.c
@@ -13,6 +13,101 @@
 #include "string.h"
 #include "umath.h"
 
+static int
+minmax_resolve_descriptors(struct PyArrayMethodObject_tag *NPY_UNUSED(method),
+                           PyArray_DTypeMeta *NPY_UNUSED(dtypes[]),
+                           PyArray_Descr *given_descrs[],
+                           PyArray_Descr *loop_descrs[],
+                           npy_intp *NPY_UNUSED(view_offset))
+{
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+    Py_INCREF(given_descrs[1]);
+    loop_descrs[1] = given_descrs[1];
+
+    StringDTypeObject *new = new_stringdtype_instance();
+    if (new == NULL) {
+        return -1;
+    }
+    loop_descrs[2] = (PyArray_Descr *)new;
+
+    return NPY_NO_CASTING;
+}
+
+static int
+maximum_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
+                     char *const data[], npy_intp const dimensions[],
+                     npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
+{
+    npy_intp N = dimensions[0];
+    char *in1 = data[0];
+    char *in2 = data[1];
+    char *out = data[2];
+    npy_intp in1_stride = strides[0];
+    npy_intp in2_stride = strides[1];
+    npy_intp out_stride = strides[2];
+
+    while (N--) {
+        if (compare(in1, in2, NULL) > 0) {
+            // Only copy *out* to *in1* if they point to different locations;
+            // for *arr.max()* they point to the same address.
+            if (in1 != out) {
+                ssfree((ss *)out);
+                if (ssdup((ss *)in1, (ss *)out) < 0) {
+                    return -1;
+                }
+            }
+        }
+        else {
+            ssfree((ss *)out);
+            if (ssdup((ss *)in2, (ss *)out) < 0) {
+                return -1;
+            }
+        }
+        in1 += in1_stride;
+        in2 += in2_stride;
+        out += out_stride;
+    }
+
+    return 0;
+}
+
+static int
+minimum_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
+                     char *const data[], npy_intp const dimensions[],
+                     npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
+{
+    npy_intp N = dimensions[0];
+    char *in1 = data[0];
+    char *in2 = data[1];
+    char *out = data[2];
+    npy_intp in1_stride = strides[0];
+    npy_intp in2_stride = strides[1];
+    npy_intp out_stride = strides[2];
+
+    while (N--) {
+        if (compare(in1, in2, NULL) < 0) {
+            if (in1 != out) {
+                ssfree((ss *)out);
+                if (ssdup((ss *)in1, (ss *)out) < 0) {
+                    return -1;
+                }
+            }
+        }
+        else {
+            ssfree((ss *)out);
+            if (ssdup((ss *)in2, (ss *)out) < 0) {
+                return -1;
+            }
+        }
+        in1 += in1_stride;
+        in2 += in2_stride;
+        out += out_stride;
+    }
+
+    return 0;
+}
+
 static int
 string_equal_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
                           char *const data[], npy_intp const dimensions[],
@@ -270,6 +365,19 @@ init_ufuncs(void)
         goto error;
     }
 
+    PyArray_DTypeMeta *minmax_dtypes[] = {&StringDType, &StringDType,
+                                          &StringDType};
+    if (init_ufunc(numpy, "maximum", minmax_dtypes,
+                   &minmax_resolve_descriptors, &maximum_strided_loop,
+                   "string_maximum", 2, 1, NPY_NO_CASTING, 0) < 0) {
+        goto error;
+    }
+    if (init_ufunc(numpy, "minimum", minmax_dtypes,
+                   &minmax_resolve_descriptors, &minimum_strided_loop,
+                   "string_minimum", 2, 1, NPY_NO_CASTING, 0) < 0) {
+        goto error;
+    }
+
     Py_DECREF(numpy);
     return 0;
 
diff --git a/stringdtype/tests/test_stringdtype.py b/stringdtype/tests/test_stringdtype.py
@@ -12,7 +12,7 @@
 
 @pytest.fixture
 def string_list():
-    return ["abc", "def", "ghi"]
+    return ["abc", "def", "ghi", "A¢☃€ 😊", "Abc", "DEF"]
 
 
 def test_scalar_creation():
@@ -44,13 +44,7 @@ def test_array_creation_utf8(data):
 
 
 def test_array_creation_scalars(string_list):
-    arr = np.array(
-        [
-            StringScalar("abc"),
-            StringScalar("def"),
-            StringScalar("ghi"),
-        ]
-    )
+    arr = np.array([StringScalar(s) for s in string_list])
     assert repr(arr) == repr(np.array(string_list, dtype=StringDType()))
 
 
@@ -69,29 +63,30 @@ def test_bad_scalars(data):
 
 
 @pytest.mark.parametrize(
-    ("string_list"),
+    ("strings"),
     [
         ["this", "is", "an", "array"],
         ["€", "", "😊"],
         ["A¢☃€ 😊", " A☃€¢😊", "☃€😊 A¢", "😊☃A¢ €"],
     ],
 )
-def test_unicode_casts(string_list):
-    arr = np.array(string_list, dtype=np.unicode_).astype(StringDType())
-    expected = np.array(string_list, dtype=StringDType())
+def test_unicode_casts(strings):
+    arr = np.array(strings, dtype=np.unicode_).astype(StringDType())
+    expected = np.array(strings, dtype=StringDType())
     np.testing.assert_array_equal(arr, expected)
 
-    arr = np.array(string_list, dtype=StringDType())
+    arr = np.array(strings, dtype=StringDType())
+
     np.testing.assert_array_equal(
-        arr.astype("U8"), np.array(string_list, dtype="U8")
+        arr.astype("U8"), np.array(strings, dtype="U8")
     )
     np.testing.assert_array_equal(arr.astype("U8").astype(StringDType()), arr)
     np.testing.assert_array_equal(
-        arr.astype("U3"), np.array(string_list, dtype="U3")
+        arr.astype("U3"), np.array(strings, dtype="U3")
     )
     np.testing.assert_array_equal(
         arr.astype("U3").astype(StringDType()),
-        np.array([s[:3] for s in string_list], dtype=StringDType()),
+        np.array([s[:3] for s in strings], dtype=StringDType()),
     )
 
 
@@ -107,10 +102,14 @@ def test_additional_unicode_cast(string_list):
 
 
 def test_insert_scalar(string_list):
+    """Test that inserting a scalar works."""
     dtype = StringDType()
     arr = np.array(string_list, dtype=dtype)
     arr[1] = StringScalar("what")
-    assert repr(arr) == repr(np.array(["abc", "what", "ghi"], dtype=dtype))
+    np.testing.assert_array_equal(
+        arr,
+        np.array(string_list[:1] + ["what"] + string_list[2:], dtype=dtype),
+    )
 
 
 def test_equality_promotion(string_list):
@@ -128,8 +127,8 @@ def test_isnan(string_list):
     )
 
 
-def test_memory_usage(string_list):
-    sarr = np.array(string_list, dtype=StringDType())
+def test_memory_usage():
+    sarr = np.array(["abc", "def", "ghi"], dtype=StringDType())
     # 4 bytes for each ASCII string buffer in string_list
     # (three characters and null terminator)
     # plus enough bytes for the size_t length
@@ -258,16 +257,16 @@ def test_arrfuncs_empty(arrfunc, expected):
 
 
 @pytest.mark.parametrize(
-    ("string_list", "cast_answer", "any_answer", "all_answer"),
+    ("strings", "cast_answer", "any_answer", "all_answer"),
     [
         [["hello", "world"], [True, True], True, True],
         [["", ""], [False, False], False, False],
         [["hello", ""], [True, False], True, False],
         [["", "world"], [False, True], True, False],
     ],
 )
-def test_bool_cast(string_list, cast_answer, any_answer, all_answer):
-    sarr = np.array(string_list, dtype=StringDType())
+def test_bool_cast(strings, cast_answer, any_answer, all_answer):
+    sarr = np.array(strings, dtype=StringDType())
     np.testing.assert_array_equal(sarr.astype("bool"), cast_answer)
 
     assert np.any(sarr) == any_answer
@@ -285,3 +284,18 @@ def test_take(string_list):
     out[0] = "hello"
     res = sarr.take(np.arange(len(string_list)), out=out)
     np.testing.assert_array_equal(res, out)
+
+
+@pytest.mark.parametrize(
+    "ufunc,func",
+    [
+        ("min", min),
+        ("max", max),
+    ],
+)
+def test_ufuncs_minmax(string_list, ufunc, func):
+    """Test that the min/max ufuncs match Python builtin min/max behavior."""
+    arr = np.array(string_list, dtype=StringDType())
+    np.testing.assert_array_equal(
+        getattr(arr, ufunc)(), np.array(func(string_list), dtype=StringDType())
+    )
