promote string to object in comparisons

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -449,124 +449,6 @@ static PyType_Slot s2b_slots[] = {
 
 static char *s2b_name = "cast_StringDType_to_Bool";
 
-// object to string
-
-typedef struct {
-    NpyAuxData base;
-    PyArray_Descr *descr;
-    int move_references;
-} _object_to_string_auxdata;
-
-static void
-_object_to_string_auxdata_free(NpyAuxData *auxdata)
-{
-    _object_to_string_auxdata *data = (_object_to_string_auxdata *)auxdata;
-    Py_DECREF(data->descr);
-    PyMem_Free(data);
-}
-
-static NpyAuxData *
-_object_to_string_auxdata_clone(NpyAuxData *data)
-{
-    _object_to_string_auxdata *res = PyMem_Malloc(sizeof(*res));
-    if (res == NULL) {
-        return NULL;
-    }
-    memcpy(res, data, sizeof(*res));
-    Py_INCREF(res->descr);
-    return (NpyAuxData *)res;
-}
-
-static int
-object_to_string_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
-                              char *const *args, const npy_intp *dimensions,
-                              const npy_intp *strides, NpyAuxData *auxdata)
-{
-    npy_intp N = dimensions[0];
-    char *src = args[0], *dst = args[1];
-    npy_intp src_stride = strides[0], dst_stride = strides[1];
-    _object_to_string_auxdata *data = (_object_to_string_auxdata *)auxdata;
-
-    PyObject *src_ref;
-
-    while (N > 0) {
-        memcpy(&src_ref, src, sizeof(src_ref));
-        if (stringdtype_setitem((StringDTypeObject *)(data->descr),
-                                src_ref ? src_ref : Py_None,
-                                (void *)dst) < 0) {
-            return -1;
-        }
-
-        if (data->move_references && src_ref != NULL) {
-            Py_DECREF(src_ref);
-            memset(src, 0, sizeof(src_ref));
-        }
-
-        N--;
-        dst += dst_stride;
-        src += src_stride;
-    }
-    return 0;
-}
-
-NPY_NO_EXPORT int
-object_to_string_get_loop(PyArrayMethod_Context *context,
-                          int NPY_UNUSED(aligned), int move_references,
-                          const npy_intp *NPY_UNUSED(strides),
-                          PyArrayMethod_StridedLoop **out_loop,
-                          NpyAuxData **out_transferdata,
-                          NPY_ARRAYMETHOD_FLAGS *flags)
-{
-    *flags = NPY_METH_REQUIRES_PYAPI;
-
-    /* NOTE: auxdata is only really necessary to flag `move_references` */
-    _object_to_string_auxdata *data = PyMem_Malloc(sizeof(*data));
-    if (data == NULL) {
-        return -1;
-    }
-    data->base.free = &_object_to_string_auxdata_free;
-    data->base.clone = &_object_to_string_auxdata_clone;
-
-    Py_INCREF(context->descriptors[1]);
-    data->descr = context->descriptors[1];
-    data->move_references = move_references;
-    *out_transferdata = (NpyAuxData *)data;
-    *out_loop = &object_to_string_strided_loop;
-    return 0;
-}
-
-static NPY_CASTING
-object_to_string_resolve_descriptors(PyArrayMethodObject *NPY_UNUSED(self),
-                                     PyArray_DTypeMeta *dtypes[2],
-                                     PyArray_Descr *given_descrs[2],
-                                     PyArray_Descr *loop_descrs[2],
-                                     npy_intp *NPY_UNUSED(view_offset))
-{
-    if (given_descrs[1] == NULL) {
-        loop_descrs[1] = (PyArray_Descr *)new_stringdtype_instance(
-                (PyTypeObject *)dtypes[1]);
-        if (loop_descrs[1] == NULL) {
-            return -1;
-        }
-    }
-    else {
-        Py_INCREF(given_descrs[1]);
-        loop_descrs[1] = given_descrs[1];
-    }
-
-    Py_INCREF(given_descrs[0]);
-    loop_descrs[0] = given_descrs[0];
-
-    return NPY_SAFE_CASTING;
-}
-
-static PyType_Slot o2s_slots[] = {
-        {NPY_METH_resolve_descriptors, &object_to_string_resolve_descriptors},
-        {_NPY_METH_get_loop, &object_to_string_get_loop},
-        {0, NULL}};
-
-static char *o2s_name = "cast_object_to_StringDType";
-
 PyArrayMethod_Spec *
 get_cast_spec(const char *name, NPY_CASTING casting,
               NPY_ARRAYMETHOD_FLAGS flags, PyArray_DTypeMeta **dtypes,
@@ -619,10 +501,10 @@ get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other)
 
     int is_pandas = (this == (PyArray_DTypeMeta *)&PandasStringDType);
 
-    int num_casts = 6;
+    int num_casts = 5;
 
     if (is_pandas) {
-        num_casts = 8;
+        num_casts = 7;
 
         PyArray_DTypeMeta **t2o_dtypes = get_dtypes(this, other);
 
@@ -655,12 +537,6 @@ get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other)
             s2b_name, NPY_UNSAFE_CASTING, NPY_METH_NO_FLOATINGPOINT_ERRORS,
             s2b_dtypes, s2b_slots);
 
-    PyArray_DTypeMeta **o2s_dtypes = get_dtypes(&PyArray_ObjectDType, this);
-
-    PyArrayMethod_Spec *ObjectToStringCastSpec =
-            get_cast_spec(o2s_name, NPY_SAFE_CASTING, NPY_METH_REQUIRES_PYAPI,
-                          o2s_dtypes, o2s_slots);
-
     PyArrayMethod_Spec **casts = NULL;
 
     casts = malloc(num_casts * sizeof(PyArrayMethod_Spec *));
@@ -669,14 +545,13 @@ get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other)
     casts[1] = UnicodeToStringCastSpec;
     casts[2] = StringToUnicodeCastSpec;
     casts[3] = StringToBoolCastSpec;
-    casts[4] = ObjectToStringCastSpec;
     if (is_pandas) {
-        casts[5] = ThisToOtherCastSpec;
-        casts[6] = OtherToThisCastSpec;
-        casts[7] = NULL;
+        casts[4] = ThisToOtherCastSpec;
+        casts[5] = OtherToThisCastSpec;
+        casts[6] = NULL;
     }
     else {
-        casts[5] = NULL;
+        casts[4] = NULL;
     }
 
     return casts;

stringdtype/stringdtype/src/umath.c

@@ -456,7 +456,7 @@ ufunc_promoter_internal(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
     }
     Py_XDECREF(common);
 
-    /* Otherwise, set all input operands to StringDType */
+    /* Otherwise, set all input operands to final_dtype */
     for (int i = 0; i < ufunc->nargs; i++) {
         PyArray_DTypeMeta *tmp = final_dtype;
         if (signature[i]) {
@@ -474,21 +474,32 @@ ufunc_promoter_internal(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
 }
 
 static int
-string_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
-                      PyArray_DTypeMeta *signature[],
-                      PyArray_DTypeMeta *new_op_dtypes[])
+string_object_promoter(PyObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
+                       PyArray_DTypeMeta *signature[],
+                       PyArray_DTypeMeta *new_op_dtypes[])
 {
-    return ufunc_promoter_internal(ufunc, op_dtypes, signature, new_op_dtypes,
+    return ufunc_promoter_internal((PyUFuncObject *)ufunc, op_dtypes,
+                                   signature, new_op_dtypes,
+                                   (PyArray_DTypeMeta *)&PyArray_ObjectDType);
+}
+
+static int
+string_unicode_promoter(PyObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
+                        PyArray_DTypeMeta *signature[],
+                        PyArray_DTypeMeta *new_op_dtypes[])
+{
+    return ufunc_promoter_internal((PyUFuncObject *)ufunc, op_dtypes,
+                                   signature, new_op_dtypes,
                                    (PyArray_DTypeMeta *)&StringDType);
 }
 
 static int
-pandas_string_ufunc_promoter(PyUFuncObject *ufunc,
-                             PyArray_DTypeMeta *op_dtypes[],
-                             PyArray_DTypeMeta *signature[],
-                             PyArray_DTypeMeta *new_op_dtypes[])
+pandas_string_unicode_promoter(PyObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
+                               PyArray_DTypeMeta *signature[],
+                               PyArray_DTypeMeta *new_op_dtypes[])
 {
-    return ufunc_promoter_internal(ufunc, op_dtypes, signature, new_op_dtypes,
+    return ufunc_promoter_internal((PyUFuncObject *)ufunc, op_dtypes,
+                                   signature, new_op_dtypes,
                                    (PyArray_DTypeMeta *)&PandasStringDType);
 }
 
@@ -538,7 +549,7 @@ init_ufunc(PyObject *numpy, const char *ufunc_name, PyArray_DTypeMeta **dtypes,
 int
 add_promoter(PyObject *numpy, const char *ufunc_name,
              PyArray_DTypeMeta *ldtype, PyArray_DTypeMeta *rdtype,
-             PyArray_DTypeMeta *edtype, int is_pandas)
+             PyArray_DTypeMeta *edtype, promoter_function *promoter_impl)
 {
     PyObject *ufunc = PyObject_GetAttrString(numpy, ufunc_name);
 
@@ -553,16 +564,8 @@ add_promoter(PyObject *numpy, const char *ufunc_name,
         return -1;
     }
 
-    PyObject *promoter_capsule = NULL;
-
-    if (is_pandas == 0) {
-        promoter_capsule = PyCapsule_New((void *)&string_ufunc_promoter,
-                                         "numpy._ufunc_promoter", NULL);
-    }
-    else {
-        promoter_capsule = PyCapsule_New((void *)&pandas_string_ufunc_promoter,
-                                         "numpy._ufunc_promoter", NULL);
-    }
+    PyObject *promoter_capsule = PyCapsule_New((void *)promoter_impl,
+                                               "numpy._ufunc_promoter", NULL);
 
     if (promoter_capsule == NULL) {
         Py_DECREF(ufunc);
@@ -592,21 +595,31 @@ init_ufuncs(void)
         return -1;
     }
 
-    StringDType_type **dtype_classes = NULL;
     int num_dtypes;
 
     if (PANDAS_AVAILABLE) {
-        dtype_classes = malloc(sizeof(StringDType_type *) * 2);
-        dtype_classes[0] = &StringDType;
-        dtype_classes[1] = &PandasStringDType;
         num_dtypes = 2;
     }
     else {
-        dtype_classes = malloc(sizeof(StringDType_type *) * 1);
-        dtype_classes[0] = &StringDType;
         num_dtypes = 1;
     }
 
+    StringDType_type **dtype_classes =
+            malloc(sizeof(StringDType_type *) * num_dtypes);
+    promoter_function **unicode_promoters =
+            malloc(sizeof(promoter_function *) * num_dtypes);
+    dtype_classes[0] = &StringDType;
+    unicode_promoters[0] = &string_unicode_promoter;
+
+    if (PANDAS_AVAILABLE) {
+        dtype_classes[1] = &PandasStringDType;
+        unicode_promoters[1] = &pandas_string_unicode_promoter;
+    }
+
+    static char *comparison_ufunc_names[6] = {"equal",   "not_equal",
+                                              "greater", "greater_equal",
+                                              "less",    "less_equal"};
+
     for (int di = 0; di < num_dtypes; di++) {
         PyArray_DTypeMeta *comparison_dtypes[] = {
                 (PyArray_DTypeMeta *)dtype_classes[di],
@@ -654,34 +667,32 @@ init_ufuncs(void)
             goto error;
         }
 
-        static char *ufunc_names[6] = {"equal",   "not_equal",
-                                       "greater", "greater_equal",
-                                       "less",    "less_equal"};
-
         for (int i = 0; i < 6; i++) {
-            if (add_promoter(numpy, ufunc_names[i],
+            if (add_promoter(numpy, comparison_ufunc_names[i],
                              (PyArray_DTypeMeta *)dtype_classes[di],
                              &PyArray_UnicodeDType, &PyArray_BoolDType,
-                             0) < 0) {
+                             unicode_promoters[di]) < 0) {
                 goto error;
             }
 
-            if (add_promoter(numpy, ufunc_names[i], &PyArray_UnicodeDType,
+            if (add_promoter(numpy, comparison_ufunc_names[i],
+                             &PyArray_UnicodeDType,
                              (PyArray_DTypeMeta *)dtype_classes[di],
-                             &PyArray_BoolDType, 0) < 0) {
+                             &PyArray_BoolDType, unicode_promoters[di]) < 0) {
                 goto error;
             }
 
-            if (add_promoter(numpy, ufunc_names[i], &PyArray_ObjectDType,
-                             (PyArray_DTypeMeta *)dtype_classes[di],
-                             &PyArray_BoolDType, 0) < 0) {
+            if (add_promoter(
+                        numpy, comparison_ufunc_names[i], &PyArray_ObjectDType,
+                        (PyArray_DTypeMeta *)dtype_classes[di],
+                        &PyArray_BoolDType, &string_object_promoter) < 0) {
                 goto error;
             }
 
-            if (add_promoter(numpy, ufunc_names[i],
+            if (add_promoter(numpy, comparison_ufunc_names[i],
                              (PyArray_DTypeMeta *)dtype_classes[di],
                              &PyArray_ObjectDType, &PyArray_BoolDType,
-                             0) < 0) {
+                             &string_object_promoter) < 0) {
                 goto error;
             }
         }
@@ -720,10 +731,36 @@ init_ufuncs(void)
         }
     }
 
+    // add promoters for all ufuncs so comparison operations mixing StringDType
+    // and PandasStringDType work correctly.
+
+    if (PANDAS_AVAILABLE) {
+        for (int i = 0; i < 6; i++) {
+            if (add_promoter(numpy, comparison_ufunc_names[i],
+                             (PyArray_DTypeMeta *)&StringDType,
+                             (PyArray_DTypeMeta *)&PandasStringDType,
+                             &PyArray_BoolDType,
+                             string_unicode_promoter) < 0) {
+                goto error;
+            }
+
+            if (add_promoter(numpy, comparison_ufunc_names[i],
+                             (PyArray_DTypeMeta *)&PandasStringDType,
+                             (PyArray_DTypeMeta *)&StringDType,
+                             &PyArray_BoolDType,
+                             string_unicode_promoter) < 0) {
+                goto error;
+            }
+        }
+    }
+    free(dtype_classes);
+    free(unicode_promoters);
     Py_DECREF(numpy);
     return 0;
 
 error:
+    free(dtype_classes);
+    free(unicode_promoters);
     Py_DECREF(numpy);
     return -1;
 }