Fixing bug in stability of mergesort impl for StringArray (#961)

Motivation: for StringArray type legacy implementation of stable sort computed result when sorting with ascending=False by reversing the result of argsorting with ascending=True, which produces wrong order in groups of elements with the same value. Implemented solution adds new function argument 'ascening' and uses it when calling native function impl via serial stable_sort.
IntelPython · kozlov-alexey · Feb 19, 2021 · Jan 29, 2021 · Feb 18, 2021 · Feb 18, 2021
commit 4926b0d4ea8d75d95f4ca16be9f6fd50dad8a256
diff --git a/sdc/_str_ext.cpp b/sdc/_str_ext.cpp
@@ -31,6 +31,7 @@
 #include <string>
 #include <vector>
 #include <cmath>
+#include <algorithm>
 
 #include "_str_decode.cpp"
 
@@ -129,6 +130,7 @@ extern "C"
     npy_intp array_size(PyArrayObject* arr);
     void* array_getptr1(PyArrayObject* arr, npy_intp ind);
     void array_setitem(PyArrayObject* arr, char* p, PyObject* s);
+    void stable_argsort(char* data_ptr, uint32_t* in_offsets, int64_t len, int8_t ascending, uint64_t* result);
 
     PyMODINIT_FUNC PyInit_hstr_ext(void)
     {
@@ -201,6 +203,7 @@ extern "C"
         PyObject_SetAttrString(m, "array_setitem", PyLong_FromVoidPtr((void*)(&array_setitem)));
         PyObject_SetAttrString(m, "decode_utf8", PyLong_FromVoidPtr((void*)(&decode_utf8)));
         PyObject_SetAttrString(m, "get_utf8_size", PyLong_FromVoidPtr((void*)(&get_utf8_size)));
+        PyObject_SetAttrString(m, "stable_argsort", PyLong_FromVoidPtr((void*)(&stable_argsort)));
         return m;
     }
 
@@ -871,4 +874,35 @@ extern "C"
         return;
     }
 
+    void stable_argsort(char* data_ptr, uint32_t* in_offsets, int64_t len, int8_t ascending, uint64_t* result)
+    {
+        using str_index_pair_type = std::pair<std::string, int64_t>;
+        std::vector<str_index_pair_type> str_arr_indexed;
+        str_arr_indexed.reserve(len);
+
+        for (int64_t i=0; i < len; ++i)
+        {
+            uint32_t start = in_offsets[i];
+            uint32_t size = in_offsets[i + 1] - in_offsets[i];
+            str_arr_indexed.emplace_back(
+                    std::move(std::string(&data_ptr[start], size)),
+                    i
+            );
+        }
+
+        std::stable_sort(str_arr_indexed.begin(),
+                         str_arr_indexed.end(),
+                         [=](const str_index_pair_type& left, const str_index_pair_type& right){
+                            if (ascending)
+                                return left.first < right.first;
+                            else
+                                return left.first > right.first;
+                         }
+        );
+
+        for (int64_t i=0; i < len; ++i)
+            result[i] = str_arr_indexed[i].second;
+    }
+
+
 } // extern "C"
diff --git a/sdc/datatypes/common_functions.py b/sdc/datatypes/common_functions.py
@@ -52,7 +52,7 @@
 from sdc.str_arr_ext import (num_total_chars, append_string_array_to,
                              str_arr_is_na, pre_alloc_string_array, str_arr_set_na, string_array_type,
                              cp_str_list_to_array, create_str_arr_from_list, get_utf8_size,
-                             str_arr_set_na_by_mask)
+                             str_arr_set_na_by_mask, str_arr_stable_argosort)
 from sdc.utilities.prange_utils import parallel_chunks
 from sdc.utilities.utils import sdc_overload, sdc_register_jitable
 from sdc.utilities.sdc_typing_utils import (
@@ -518,41 +518,39 @@ def sdc_arrays_argsort(A, kind='quicksort'):
 
 
 @sdc_overload(sdc_arrays_argsort, jit_options={'parallel': False})
-def sdc_arrays_argsort_overload(A, kind='quicksort'):
+def sdc_arrays_argsort_overload(A, kind='quicksort', ascending=True):
     """Function providing pandas argsort implementation for different 1D array types"""
 
     # kind is not known at compile time, so get this function here and use in impl if needed
     quicksort_func = quicksort.make_jit_quicksort().run_quicksort
 
     kind_is_default = isinstance(kind, str)
     if isinstance(A, types.Array):
-        def _sdc_arrays_argsort_array_impl(A, kind='quicksort'):
+        def _sdc_arrays_argsort_array_impl(A, kind='quicksort', ascending=True):
             _kind = 'quicksort' if kind_is_default == True else kind  # noqa
-            return numpy_like.argsort(A, kind=_kind)
+            return numpy_like.argsort(A, kind=_kind, ascending=ascending)
 
         return _sdc_arrays_argsort_array_impl
 
     elif A == string_array_type:
-        def _sdc_arrays_argsort_str_arr_impl(A, kind='quicksort'):
+        def _sdc_arrays_argsort_str_arr_impl(A, kind='quicksort', ascending=True):
 
-            nan_mask = sdc.hiframes.api.get_nan_mask(A)
-            idx = numpy.arange(len(A))
-            old_nan_positions = idx[nan_mask]
-
-            data = A[~nan_mask]
-            keys = idx[~nan_mask]
             if kind == 'quicksort':
-                zipped = list(zip(list(data), list(keys)))
-                zipped = quicksort_func(zipped)
-                argsorted = [zipped[i][1] for i in numpy.arange(len(data))]
+                indexes = numpy.arange(len(A))
+                data_index_pairs = list(zip(list(A), list(indexes)))
+                zipped = quicksort_func(data_index_pairs)
+                argsorted = [zipped[i][1] for i in indexes]
+                res = numpy.array(argsorted, dtype=numpy.int64)
+                # for non-stable sort the order within groups does not matter
+                # so just reverse the result when sorting in descending order
+                if not ascending:
+                    res = res[::-1]
             elif kind == 'mergesort':
-                sdc.hiframes.sort.local_sort((data, ), (keys, ))
-                argsorted = list(keys)
+                res = str_arr_stable_argosort(A, ascending=ascending)
             else:
                 raise ValueError("Unrecognized kind of sort in sdc_arrays_argsort")
 
-            argsorted.extend(old_nan_positions)
-            return numpy.asarray(argsorted, dtype=numpy.int32)
+            return res
 
         return _sdc_arrays_argsort_str_arr_impl
 

diff --git a/sdc/datatypes/hpat_pandas_series_functions.py b/sdc/datatypes/hpat_pandas_series_functions.py
@@ -3950,11 +3950,9 @@ def _sdc_pandas_series_sort_values_impl(
         good = ~data_nan_mask
 
         if kind_is_none_or_default == True:  # noqa
-            argsort_res = sdc_arrays_argsort(self._data[good], kind='quicksort')
+            argsort_res = sdc_arrays_argsort(self._data[good], kind='quicksort', ascending=ascending)
         else:
-            argsort_res = sdc_arrays_argsort(self._data[good], kind=kind)
-        if not ascending:
-            argsort_res = argsort_res[::-1]
+            argsort_res = sdc_arrays_argsort(self._data[good], kind=kind, ascending=ascending)
 
         idx = numpy.arange(len(self), dtype=numpy.int32)
         sorted_index = numpy.empty(len(self), dtype=numpy.int32)

diff --git a/sdc/functions/numpy_like.py b/sdc/functions/numpy_like.py
@@ -1225,7 +1225,7 @@ def sort_impl(a, axis=-1, kind=None, order=None):
     return sort_impl
 
 
-def argsort(a, axis=-1, kind=None, order=None):
+def argsort(a, axis=-1, kind=None, order=None, ascending=True):
     """
     Returns the indices that would sort an array.
 
@@ -1254,7 +1254,7 @@ def argsort(a, axis=-1, kind=None, order=None):
 
 
 @sdc_overload(argsort)
-def argsort_overload(a, axis=-1, kind=None, order=None):
+def argsort_overload(a, axis=-1, kind=None, order=None, ascending=True):
     _func_name = 'argsort'
     ty_checker = TypeChecker(_func_name)
 
@@ -1266,15 +1266,15 @@ def argsort_overload(a, axis=-1, kind=None, order=None):
     if not is_default(order, None):
         raise TypingError(f'{_func_name} Unsupported parameter order')
 
-    def argsort_impl(a, axis=-1, kind=None, order=None):
+    def argsort_impl(a, axis=-1, kind=None, order=None, ascending=True):
         _kind = 'quicksort'
         if kind is not None:
             _kind = kind
 
         if _kind == 'quicksort':
-            return parallel_argsort(a)
+            return parallel_argsort(a, ascending)
         elif _kind == 'mergesort':
-            return parallel_stable_argsort(a)
+            return parallel_stable_argsort(a, ascending)
         else:
             raise ValueError("Unsupported value of 'kind' parameter")
 

diff --git a/sdc/functions/sort.py b/sdc/functions/sort.py
@@ -47,7 +47,7 @@ def bind(sym, sig):
 parallel_sort_sig = ct.CFUNCTYPE(None, ct.c_void_p, ct.c_uint64,
                                  ct.c_uint64, ct.c_void_p,)
 
-parallel_argsort_arithm_sig = ct.CFUNCTYPE(None, ct.c_void_p, ct.c_void_p, ct.c_uint64)
+parallel_argsort_arithm_sig = ct.CFUNCTYPE(None, ct.c_void_p, ct.c_void_p, ct.c_uint64, ct.c_uint8)
 
 parallel_argsort_sig = ct.CFUNCTYPE(None, ct.c_void_p, ct.c_void_p, ct.c_uint64,
                                     ct.c_uint64, ct.c_void_p,)
@@ -66,7 +66,7 @@ def bind(sym, sig):
 
 parallel_sort_t_sig = ct.CFUNCTYPE(None, ct.c_void_p, ct.c_uint64)
 
-parallel_argsort_t_sig = ct.CFUNCTYPE(None, ct.c_void_p, ct.c_void_p, ct.c_uint64)
+parallel_argsort_t_sig = ct.CFUNCTYPE(None, ct.c_void_p, ct.c_void_p, ct.c_uint64, ct.c_uint8)
 
 set_threads_count_sig = ct.CFUNCTYPE(None, ct.c_uint64)
 set_threads_count_sym = bind('set_number_of_threads', set_threads_count_sig)
@@ -290,30 +290,32 @@ def parallel_xargsort_overload_impl(dt, xargsort_map, xargsort_sym):
     if dt in types_to_postfix.keys():
         sort_f = xargsort_map[dt]
 
-        def parallel_xargsort_arithm_impl(arr):
+        def parallel_xargsort_arithm_impl(arr, ascending=True):
             index = numpy.empty(shape=len(arr), dtype=numpy.int64)
-            sort_f(index.ctypes, arr.ctypes, len(arr))
+            sort_f(index.ctypes, arr.ctypes, len(arr), types.uint8(ascending))
 
             return index
 
         return parallel_xargsort_arithm_impl
 
-    def parallel_xargsort_impl(arr):
+    # TO-DO: add/change adaptor to handle case of ascending=False
+    def parallel_xargsort_impl(arr, ascending=True):
         item_size = itemsize(arr)
         index = numpy.empty(shape=len(arr), dtype=numpy.int64)
+
         xargsort_sym(index.ctypes, arr.ctypes, len(arr), item_size, adaptor(arr[0], arr[0]))
 
         return index
 
     return parallel_xargsort_impl
 
 
-def parallel_argsort(arr):
+def parallel_argsort(arr, ascending=True):
     pass
 
 
 @overload(parallel_argsort)
-def parallel_argsort_overload(arr):
+def parallel_argsort_overload(arr, ascending=True):
 
     if not isinstance(arr, types.Array):
         raise NotImplementedError
@@ -323,12 +325,12 @@ def parallel_argsort_overload(arr):
     return parallel_xargsort_overload_impl(dt, argsort_map, parallel_argsort_sym)
 
 
-def parallel_stable_argsort(arr):
+def parallel_stable_argsort(arr, ascending=True):
     pass
 
 
 @overload(parallel_stable_argsort)
-def parallel_argsort_overload(arr):
+def parallel_stable_argsort_overload(arr, ascending=True):
 
     if not isinstance(arr, types.Array):
         raise NotImplementedError

diff --git a/sdc/native/module.cpp b/sdc/native/module.cpp
@@ -60,31 +60,31 @@ extern "C"
 
     void parallel_argsort_u64v(void* index, void* begin, uint64_t len, uint64_t size, void* compare);
 
-    void parallel_argsort_u64i8(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64u8(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64i16(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64u16(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64i32(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64u32(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64i64(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64u64(void* index, void* begin, uint64_t len);
-
-    void parallel_argsort_u64f32(void* index, void* begin, uint64_t len);
-    void parallel_argsort_u64f64(void* index, void* begin, uint64_t len);
+    void parallel_argsort_u64i8(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64u8(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64i16(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64u16(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64i32(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64u32(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64i64(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64u64(void* index, void* begin, uint64_t len, uint8_t ascending);
+
+    void parallel_argsort_u64f32(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_argsort_u64f64(void* index, void* begin, uint64_t len, uint8_t ascending);
 
     void parallel_stable_argsort_u64v(void* index, void* begin, uint64_t len, uint64_t size, void* compare);
 
-    void parallel_stable_argsort_u64i8(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64u8(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64i16(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64u16(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64i32(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64u32(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64i64(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64u64(void* index, void* begin, uint64_t len);
-
-    void parallel_stable_argsort_u64f32(void* index, void* begin, uint64_t len);
-    void parallel_stable_argsort_u64f64(void* index, void* begin, uint64_t len);
+    void parallel_stable_argsort_u64i8(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64u8(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64i16(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64u16(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64i32(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64u32(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64i64(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64u64(void* index, void* begin, uint64_t len, uint8_t ascending);
+
+    void parallel_stable_argsort_u64f32(void* index, void* begin, uint64_t len, uint8_t ascending);
+    void parallel_stable_argsort_u64f64(void* index, void* begin, uint64_t len, uint8_t ascending);
 
     void set_number_of_threads(uint64_t threads)
     {

diff --git a/sdc/native/sort.cpp b/sdc/native/sort.cpp
@@ -92,8 +92,16 @@ void parallel_argsort_(I* index, void* data, uint64_t len, uint64_t size, compar
 } // namespace
 
 #define declare_single_argsort(index_prefix, type_prefix, ity, ty) \
-void parallel_argsort_##index_prefix##type_prefix(void* index, void* begin, uint64_t len) \
-{ parallel_argsort_(reinterpret_cast<ity*>(index), reinterpret_cast<ty*>(begin), len); }
+void parallel_argsort_##index_prefix##type_prefix(void* index, void* begin, uint64_t len, uint8_t ascending) \
+{ \
+    if (ascending) { \
+        auto cmp = utils::less<ty>(); \
+        parallel_argsort_(reinterpret_cast<ity*>(index), reinterpret_cast<ty*>(begin), len, cmp); \
+    } else { \
+        auto cmp = utils::greater<ty>(); \
+        parallel_argsort_(reinterpret_cast<ity*>(index), reinterpret_cast<ty*>(begin), len, cmp); \
+    } \
+}
 
 #define declare_argsort(prefix, ty) \
 declare_single_argsort(u8,  prefix, uint8_t,  ty) \

diff --git a/sdc/native/stable_sort.cpp b/sdc/native/stable_sort.cpp
@@ -281,8 +281,16 @@ struct parallel_sort_fixed_size
 } // namespace
 
 #define declare_single_argsort(index_prefix, type_prefix, ity, ty) \
-void parallel_stable_argsort_##index_prefix##type_prefix(ity* index, void* begin, uint64_t len) \
-{ parallel_stable_argsort_(reinterpret_cast<ity*>(index), reinterpret_cast<ty*>(begin), len); }
+void parallel_stable_argsort_##index_prefix##type_prefix(ity* index, void* begin, uint64_t len, uint8_t ascending) \
+{ \
+    if (ascending) { \
+        auto cmp = utils::less<ty>(); \
+        parallel_stable_argsort_(reinterpret_cast<ity*>(index), reinterpret_cast<ty*>(begin), len, cmp); \
+    } else { \
+        auto cmp = utils::greater<ty>(); \
+        parallel_stable_argsort_(reinterpret_cast<ity*>(index), reinterpret_cast<ty*>(begin), len, cmp); \
+    } \
+}
 
 #define declare_argsort(prefix, ty) \
 declare_single_argsort(u8,  prefix, uint8_t,  ty) \
@@ -339,4 +347,4 @@ void parallel_stable_sort(void* begin, uint64_t len, uint64_t size, void* compar
 #undef declare_int_sort
 #undef declare_sort
 #undef declare_argsort
-#undef declare_single_argsort
+#undef declare_single_argsort
diff --git a/sdc/native/utils.cpp b/sdc/native/utils.cpp
@@ -169,4 +169,16 @@ bool nanless<double>(const double& left, const double& right)
     return std::less<double>()(left, right) || (std::isnan(right) && !std::isnan(left));
 }
 
+template<>
+bool nangreater<float>(const float& left, const float& right)
+{
+    return std::greater<float>()(left, right) || (std::isnan(right) && !std::isnan(left));
+}
+
+template<>
+bool nangreater<double>(const double& left, const double& right)
+{
+    return std::greater<double>()(left, right) || (std::isnan(right) && !std::isnan(left));
+}
+
 }
diff --git a/sdc/native/utils.hpp b/sdc/native/utils.hpp
@@ -266,6 +266,27 @@ struct less
     }
 };
 
+template<typename T>
+bool nangreater(const T& left, const T& right)
+{
+    return std::greater<T>()(left, right);
+}
+
+template<>
+bool nangreater<float>(const float& left, const float& right);
+
+template<>
+bool nangreater<double>(const double& left, const double& right);

+template<typename T>
+struct greater
+{
+    bool operator() (const T& left, const T& right) const
+    {
+        return nangreater<T>(left, right);
+    }
+};
+
 namespace tbb_control
 {
     void init();