Respond to review comments

Author: ngoldbaum

stringdtype/stringdtype/src/dtype.c

@@ -152,6 +152,7 @@ stringdtype_setitem(StringDTypeObject *NPY_UNUSED(descr), PyObject *obj,
     }
 
     if (eq_res == 1) {
+        // NULL is the representation of NA in the array buffer
         sdata = NULL;
     }
     else {
@@ -233,17 +234,11 @@ compare(void *a, void *b, void *NPY_UNUSED(arr))
     ss *ss_b = (ss *)b;
     int a_is_null = ss_isnull(ss_a);
     int b_is_null = ss_isnull(ss_b);
-    if (a_is_null || b_is_null) {
-        // numpy sorts NaNs to the end of the array
-        // pandas sorts NAs to the end as well
-        // so we follow that behavior here
-        if (!b_is_null) {
-            return 1;
-        }
-        else if (!a_is_null) {
-            return -1;
-        }
-        return 0;
+    if (a_is_null) {
+        return 1;
+    }
+    else if (b_is_null) {
+        return -1;
     }
     return strcmp(ss_a->buf, ss_b->buf);
 }

stringdtype/tests/test_stringdtype.py

@@ -185,12 +185,24 @@ def test_pickle(string_list):
 )
 def test_sort(strings):
     """Test that sorting matches python's internal sorting."""
-    arr = np.array(strings, dtype=StringDType())
+
+    def test_sort(strings, arr_sorted):
+        arr = np.array(strings, dtype=StringDType())
+        np.random.default_rng().shuffle(arr)
+        arr.sort()
+        assert np.array_equal(arr, arr_sorted, equal_nan=True)
+
     arr_sorted = np.array(sorted(strings), dtype=StringDType())
+    test_sort(strings, arr_sorted)
+
+    # make sure NAs get sorted to the end of the array
+    strings.insert(0, NA)
+    strings.insert(2, NA)
+    # can't use append because doing that with NA converts
+    # the result to object dtype
+    arr_sorted = np.array(arr_sorted.tolist() + [NA, NA], dtype=StringDType())
 
-    np.random.default_rng().shuffle(arr)
-    arr.sort()
-    np.testing.assert_array_equal(arr, arr_sorted)
+    test_sort(strings, arr_sorted)
 
 
 @pytest.mark.parametrize(