Implement parallel argsort (#910)

Author: AlexanderKalistratov

sdc/datatypes/common_functions.py

@@ -277,8 +277,8 @@ def sdc_join_series_indexes_impl(left, right):
                         right_nan.append(i)
 
                 # sort arrays saving the old positions
-                sorted_left = numpy.argsort(left, kind='mergesort')
-                sorted_right = numpy.argsort(right, kind='mergesort')
+                sorted_left = numpy_like.argsort(left, kind='mergesort')
+                sorted_right = numpy_like.argsort(right, kind='mergesort')
                 # put the position of the nans in an increasing sequence
                 sorted_left[lsize-len(left_nan):] = left_nan
                 sorted_right[rsize-len(right_nan):] = right_nan
@@ -523,7 +523,7 @@ def sdc_arrays_argsort_overload(A, kind='quicksort'):
     if isinstance(A, types.Array):
         def _sdc_arrays_argsort_array_impl(A, kind='quicksort'):
             _kind = 'quicksort' if kind_is_default == True else kind  # noqa
-            return numpy.argsort(A, kind=_kind)
+            return numpy_like.argsort(A, kind=_kind)
 
         return _sdc_arrays_argsort_array_impl
 

sdc/datatypes/hpat_pandas_series_functions.py

@@ -1137,7 +1137,7 @@ def hpat_pandas_series_nsmallest_impl(self, n=5, keep='first'):
             raise ValueError("Method nsmallest(). Unsupported parameter. Given 'keep' != 'first'")
 
         # mergesort is used for stable sorting of repeated values
-        indices = self._data.argsort(kind='mergesort')[:max(n, 0)]
+        indices = numpy_like.argsort(self._data, kind='mergesort')[:max(n, 0)]
 
         return self.take(indices)
 
@@ -1207,7 +1207,7 @@ def hpat_pandas_series_nlargest_impl(self, n=5, keep='first'):
         # data: [0, 1, -1, 1, 0] -> [1, 1, 0, 0, -1]
         # index: [0, 1,  2, 3, 4] -> [1, 3, 0, 4,  2] (not [3, 1, 4, 0, 2])
         # subtract 1 to ensure reverse ordering at boundaries
-        indices = (-self._data - 1).argsort(kind='mergesort')[:max(n, 0)]
+        indices = numpy_like.argsort(-self._data - 1, kind='mergesort')[:max(n, 0)]
 
         return self.take(indices)
 
@@ -1457,7 +1457,7 @@ def hpat_pandas_series_value_counts_str_impl(
             counts = numpy.asarray(counts_as_list, dtype=numpy.intp)
             indexes_order = numpy.arange(values_len)
             if sort:
-                indexes_order = counts.argsort()
+                indexes_order = numpy_like.argsort(counts)
                 if not ascending:
                     indexes_order = indexes_order[::-1]
 
@@ -1521,7 +1521,7 @@ def hpat_pandas_series_value_counts_number_impl(
 
             indexes_order = numpy.arange(len(value_counts))
             if sort:
-                indexes_order = value_counts.argsort()
+                indexes_order = numpy_like.argsort(value_counts)
                 if not ascending:
                     indexes_order = indexes_order[::-1]
 
@@ -3808,16 +3808,16 @@ def hpat_pandas_series_argsort_idx_impl(self, axis=0, kind='quicksort', order=No
                 raise ValueError("Method argsort(). Unsupported parameter. Given 'kind' != 'quicksort' or 'mergesort'")
             if kind == 'mergesort':
                 #It is impossible to use numpy.argsort(self._data, kind=kind) since numba gives typing error
-                sort = numpy.argsort(self._data, kind='mergesort')
+                sort = numpy_like.argsort(self._data, kind='mergesort')
             else:
-                sort = numpy.argsort(self._data)
+                sort = numpy_like.argsort(self._data)
             na = self.isna().sum()
             result = numpy.empty(len(self._data), dtype=numpy.int64)
             na_data_arr = sdc.hiframes.api.get_nan_mask(self._data)
             if kind == 'mergesort':
-                sort_nona = numpy.argsort(self._data[~na_data_arr], kind='mergesort')
+                sort_nona = numpy_like.argsort(self._data[~na_data_arr], kind='mergesort')
             else:
-                sort_nona = numpy.argsort(self._data[~na_data_arr])
+                sort_nona = numpy_like.argsort(self._data[~na_data_arr])
             q = 0
             for id, i in enumerate(sort):
                 if id in set(sort[len(self._data) - na:]):
@@ -3835,16 +3835,16 @@ def hpat_pandas_series_argsort_noidx_impl(self, axis=0, kind='quicksort', order=
         if kind != 'quicksort' and kind != 'mergesort':
             raise ValueError("Method argsort(). Unsupported parameter. Given 'kind' != 'quicksort' or 'mergesort'")
         if kind == 'mergesort':
-            sort = numpy.argsort(self._data, kind='mergesort')
+            sort = numpy_like.argsort(self._data, kind='mergesort')
         else:
-            sort = numpy.argsort(self._data)
+            sort = numpy_like.argsort(self._data)
         na = self.isna().sum()
         result = numpy.empty(len(self._data), dtype=numpy.int64)
         na_data_arr = sdc.hiframes.api.get_nan_mask(self._data)
         if kind == 'mergesort':
-            sort_nona = numpy.argsort(self._data[~na_data_arr], kind='mergesort')
+            sort_nona = numpy_like.argsort(self._data[~na_data_arr], kind='mergesort')
         else:
-            sort_nona = numpy.argsort(self._data[~na_data_arr])
+            sort_nona = numpy_like.argsort(self._data[~na_data_arr])
         q = 0
         for id, i in enumerate(sort):
             if id in set(sort[len(self._data) - na:]):

sdc/functions/numpy_like.py

@@ -47,7 +47,7 @@
 from sdc.functions.statistics import skew_formula
 from sdc.hiframes.api import isna
 from sdc.datatypes.range_index_type import RangeIndexType
-from sdc.utilities.sdc_typing_utils import TypeChecker
+from sdc.utilities.sdc_typing_utils import TypeChecker, is_default
 from sdc.utilities.utils import (sdc_overload, sdc_register_jitable,
                                  min_dtype_int_val, max_dtype_int_val, min_dtype_float_val,
                                  max_dtype_float_val)
@@ -56,6 +56,7 @@
                              num_total_chars, str_arr_is_na)
 from sdc.utilities.prange_utils import parallel_chunks
 from sdc.utilities.sdc_typing_utils import check_types_comparable
+from sdc.functions.sort import parallel_sort, parallel_stable_sort, parallel_argsort, parallel_stable_argsort
 
 def astype(self, dtype):
     pass
@@ -1154,3 +1155,110 @@ def f(A):
                 i += 1
             return arr
         return f
+
+
+def sort(a, axis=-1, kind=None, order=None):
+    """
+    Sort input array inplace.
+
+    Parameters
+    -----------
+    a: :obj:`Array`
+        Input array
+    axis: Unsupported
+    kind: {'quicksort', 'mergesort'}, optional
+        Sorting algorithm. Default is 'quicksort'.
+        In fact sorting algorithm is never niether 'quicksort' nor 'mergesort'.
+        It is just either nonstable or stable sort.
+    order: Unsupported
+
+    Returns
+    -------
+    None
+    """
+
+    pass
+
+
+@sdc_overload(sort)
+def sort_overload(a, axis=-1, kind=None, order=None):
+    _func_name = 'sort'
+    ty_checker = TypeChecker(_func_name)
+
+    ty_checker.check(a, types.Array)
+
+    if not is_default(axis, -1):
+        raise TypingError(f'{_func_name} Unsupported parameter axis')
+
+    if not is_default(order, None):
+        raise TypingError(f'{_func_name} Unsupported parameter order')
+
+    def sort_impl(a, axis=-1, kind=None, order=None):
+        _kind = 'quicksort'
+        if kind is not None:
+            _kind = kind
+
+        if _kind == 'quicksort':
+            return parallel_sort(a)
+        elif _kind == 'mergesort':
+            return parallel_stable_sort(a)
+        else:
+            raise ValueError("Unsupported value of 'kind' parameter")
+
+    return sort_impl
+
+
+def argsort(a, axis=-1, kind=None, order=None):
+    """
+    Returns the indices that would sort an array.
+
+    Perform an indirect sort along the given axis using the algorithm specified
+    by the `kind` keyword. It returns an array of indices of the same shape as
+    `a` that index data along the given axis in sorted order.
+
+    Parameters
+    ----------
+    a : :obj:`Array`
+        Array to sort.
+    axis: Unsupported
+    kind: {'quicksort', 'mergesort'}, optional
+        Sorting algorithm. Default is 'quicksort'.
+        In fact sorting algorithm is never niether 'quicksort' nor 'mergesort'.
+        It is just either nonstable or stable sort.
+    order: Unsupported
+
+    Returns
+    -------
+    index_array : ndarray, int
+        Array of indices that sort `a`.
+    """
+
+    pass
+
+
+@sdc_overload(argsort)
+def argsort_overload(a, axis=-1, kind=None, order=None):
+    _func_name = 'argsort'
+    ty_checker = TypeChecker(_func_name)
+
+    ty_checker.check(a, types.Array)
+
+    if not is_default(axis, -1):
+        raise TypingError(f'{_func_name} Unsupported parameter axis')
+
+    if not is_default(order, None):
+        raise TypingError(f'{_func_name} Unsupported parameter order')
+
+    def argsort_impl(a, axis=-1, kind=None, order=None):
+        _kind = 'quicksort'
+        if kind is not None:
+            _kind = kind
+
+        if _kind == 'quicksort':
+            return parallel_argsort(a)
+        elif _kind == 'mergesort':
+            return parallel_stable_argsort(a)
+        else:
+            raise ValueError("Unsupported value of 'kind' parameter")
+
+    return argsort_impl