fix comparisons of unequal length strings

Author: ngoldbaum

stringdtype/stringdtype/src/umath.c

@@ -226,8 +226,7 @@ string_greater_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
             // s1 or s2 is NA
             *out = (npy_bool)0;
         }
-        else if (s1->len == s2->len &&
-                 strncmp(s1->buf, s2->buf, s1->len) > 0) {
+        else if (strcmp(s1->buf, s2->buf) > 0) {
             *out = (npy_bool)1;
         }
         else {
@@ -266,8 +265,7 @@ string_greater_equal_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
             // s1 or s2 is NA
             *out = (npy_bool)0;
         }
-        else if (s1->len == s2->len &&
-                 strncmp(s1->buf, s2->buf, s1->len) >= 0) {
+        else if (strcmp(s1->buf, s2->buf) >= 0) {
             *out = (npy_bool)1;
         }
         else {
@@ -305,8 +303,7 @@ string_less_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
             // s1 or s2 is NA
             *out = (npy_bool)0;
         }
-        else if (s1->len == s2->len &&
-                 strncmp(s1->buf, s2->buf, s1->len) < 0) {
+        else if (strcmp(s1->buf, s2->buf) < 0) {
             *out = (npy_bool)1;
         }
         else {
@@ -344,8 +341,7 @@ string_less_equal_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
             // s1 or s2 is NA
             *out = (npy_bool)0;
         }
-        else if (s1->len == s2->len &&
-                 strncmp(s1->buf, s2->buf, s1->len) <= 0) {
+        else if (strcmp(s1->buf, s2->buf) <= 0) {
             *out = (npy_bool)1;
         }
         else {

stringdtype/tests/test_stringdtype.py

@@ -154,21 +154,35 @@ def test_insert_scalar(dtype, scalar, string_list):
 
 @pytest.mark.parametrize("op", comparison_operators)
 @pytest.mark.parametrize("o_dtype", [np.str_, object])
-def test_comparison(string_list, dtype, op, o_dtype):
+def test_comparisons(string_list, dtype, op, o_dtype):
     sarr = np.array(string_list, dtype=dtype)
     oarr = np.array(string_list, dtype=o_dtype)
 
     # test that comparison operators work
     res = op(sarr, sarr)
     ores = op(oarr, oarr)
-    # test that promotion on the operator works as well
+    # test that promotion works as well
     orres = op(sarr, oarr)
     olres = op(oarr, sarr)
 
     np.testing.assert_array_equal(res, ores)
     np.testing.assert_array_equal(res, orres)
     np.testing.assert_array_equal(res, olres)
 
+    # test we get the correct answer for unequal length strings
+    sarr2 = np.array([s + "2" for s in string_list], dtype=dtype)
+    oarr2 = np.array([s + "2" for s in string_list], dtype=o_dtype)
+
+    res = op(sarr, sarr2)
+    ores = op(oarr, oarr2)
+
+    np.testing.assert_array_equal(res, ores)
+
+    res = op(sarr2, sarr)
+    ores = op(oarr2, oarr)
+
+    np.testing.assert_array_equal(res, ores)
+
 
 def test_isnan(dtype, string_list):
     sarr = np.array(string_list + [dtype.na_object], dtype=dtype)