Add a multiply ufunc

Author: peytondmurray

stringdtype/stringdtype/src/umath.c

@@ -13,6 +13,62 @@
 #include "string.h"
 #include "umath.h"
 
+static NPY_CASTING
+multiply_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];
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[2] = given_descrs[0];
+
+    return NPY_NO_CASTING;
+}
+
+static int
+multiply_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];
+
+    ss *s1 = NULL, *os = NULL;
+
+    while (N--) {
+        s1 = (ss *)in1;
+        npy_int64 factor = *(npy_int64 *)in2;
+        os = (ss *)out;
+        npy_int64 newlen = (s1->len) * factor;
+
+        ssfree(os);
+        if (ssnewemptylen(newlen, os) < 0) {
+            return -1;
+        }
+
+        for (int i = 0; i < factor; i++) {
+            memcpy(os->buf + i * s1->len, s1->buf, s1->len);
+        }
+        os->buf[newlen] = '\0';
+
+        in1 += in1_stride;
+        in2 += in2_stride;
+        out += out_stride;
+    }
+    return 0;
+}
+
 static NPY_CASTING
 binary_resolve_descriptors(struct PyArrayMethodObject_tag *NPY_UNUSED(method),
                            PyArray_DTypeMeta *NPY_UNUSED(dtypes[]),
@@ -733,6 +789,18 @@ init_ufuncs(void)
         goto error;
     }
 
+    PyArray_DTypeMeta *multiply_types[] = {
+            (PyArray_DTypeMeta *)&StringDType,
+            &PyArray_Int64DType,
+            (PyArray_DTypeMeta *)&StringDType
+    };
+    
+    if (init_ufunc(numpy, "multiply", multiply_types,
+                   &multiply_resolve_descriptors, &multiply_strided_loop,
+                   "string_multiply", 2, 1, NPY_NO_CASTING, 0) < 0) {
+        goto error;
+    }
+
     Py_DECREF(numpy);
     return 0;
 

stringdtype/tests/test_stringdtype.py

@@ -450,19 +450,23 @@ def test_ufuncs_minmax(dtype, string_list, ufunc, func):
 
 
 @pytest.mark.parametrize(
-    "other_strings",
+    "ufunc,other,func",
     [
-        ["abc", "def", "ghi", "🤣", "📵", "😰"],
-        ["🚜", "🙃", "😾", "😹", "🚠", "🚌"],
-        ["🥦", "¨", "⨯", "∰ ", "⨌ ", "⎶ "],
+        ("add", "asdf", lambda arr, b: [x + b for x in arr]),
+        (
+            "add",
+            np.array(["a", "b", "c", "d", "e", "f"], dtype=StringDType()),
+            lambda arr1, arr2: [x + y for x, y in zip(arr1, arr2)],
+        ),
+        ("multiply", 3, lambda arr, b: [x * b for x in arr]),
     ],
 )
-def test_ufunc_add(dtype, string_list, other_strings):
-    arr1 = np.array(string_list, dtype=dtype)
-    arr2 = np.array(other_strings, dtype=dtype)
+def test_binary_ufuncs(dtype, string_list, ufunc, other, func):
+    """Test the two-argument ufuncs match python builtin behavior."""
+    arr = np.array(string_list, dtype=StringDType())
     np.testing.assert_array_equal(
-        np.add(arr1, arr2),
-        np.array([a + b for a, b in zip(arr1, arr2)], dtype=dtype),
+        getattr(np, ufunc)(arr, other),
+        np.array(func(string_list, other), dtype=StringDType()),
     )