Initial support of DataFrame.GroupBy by single literal column (#590)

* Initial support of DataFrame.GroupBy by single literal column

* Fixing PEP and skipping one failed test

* Applying review comments for tests

Author: kozlov-alexey

sdc/__init__.py

@@ -44,6 +44,7 @@
 import sdc.datatypes.hpat_pandas_series_rolling_functions
 import sdc.datatypes.hpat_pandas_seriesgroupby_functions
 import sdc.datatypes.hpat_pandas_stringmethods_functions
+import sdc.datatypes.hpat_pandas_groupby_functions
 
 from ._version import get_versions
 

sdc/datatypes/common_functions.py

@@ -41,11 +41,12 @@
 from numba import numpy_support
 
 import sdc
+from sdc.hiframes.api import isna
 from sdc.hiframes.pd_series_type import SeriesType
-from sdc.str_arr_ext import (
-    append_string_array_to, cp_str_list_to_array, num_total_chars,
-    pre_alloc_string_array, str_arr_is_na, str_arr_set_na, string_array_type
-)
+from sdc.str_arr_type import string_array_type
+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)
 from sdc.utilities.utils import sdc_overload, sdc_register_jitable
 from sdc.utilities.sdc_typing_utils import (find_common_dtype_from_numpy_dtypes,
                                             TypeChecker)
@@ -483,18 +484,21 @@ def sdc_arrays_argsort(A, kind='quicksort'):
 
 @sdc_overload(sdc_arrays_argsort, jit_options={'parallel': False})
 def sdc_arrays_argsort_overload(A, kind='quicksort'):
-    """Function overloading argsort for different 1D array types"""
+    """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_numeric_impl(A, kind='quicksort'):
-            return numpy.argsort(A, kind=kind)
-        return _sdc_arrays_argsort_numeric_impl
+        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 _sdc_arrays_argsort_array_impl
 
     elif A == string_array_type:
-        def _sdc_arrays_argsort_str_impl(A, kind='quicksort'):
+        def _sdc_arrays_argsort_str_arr_impl(A, kind='quicksort'):
 
             nan_mask = sdc.hiframes.api.get_nan_mask(A)
             idx = numpy.arange(len(A))
@@ -515,7 +519,10 @@ def _sdc_arrays_argsort_str_impl(A, kind='quicksort'):
             argsorted.extend(old_nan_positions)
             return numpy.asarray(argsorted, dtype=numpy.int32)
 
-        return _sdc_arrays_argsort_str_impl
+        return _sdc_arrays_argsort_str_arr_impl
+
+    elif isinstance(A, types.List):
+        return None
 
     return None
 
@@ -591,3 +598,78 @@ def _sdc_pandas_series_align_impl(series, other, size='max', finiteness=False):
         return aligned, aligned_other
 
     return _sdc_pandas_series_align_impl
+
+
+def _sdc_asarray(data):
+    pass
+
+
+@sdc_overload(_sdc_asarray, jit_options={'parallel': True})
+def _sdc_asarray_overload(data):
+
+    # TODO: extend with other types
+    if not isinstance(data, types.List):
+        return None
+
+    if isinstance(data.dtype, types.UnicodeType):
+        def _sdc_asarray_impl(data):
+            return create_str_arr_from_list(data)
+
+        return _sdc_asarray_impl
+
+    else:
+        result_dtype = data.dtype
+
+        def _sdc_asarray_impl(data):
+            # TODO: check if elementwise copy is needed at all
+            res_size = len(data)
+            res_arr = numpy.empty(res_size, dtype=result_dtype)
+            for i in numba.prange(res_size):
+                res_arr[i] = data[i]
+            return res_arr
+
+        return _sdc_asarray_impl
+
+    return None
+
+
+def _sdc_take(data, indexes):
+    pass
+
+
+@sdc_overload(_sdc_take, jit_options={'parallel': True})
+def _sdc_take_overload(data, indexes):
+
+    if isinstance(data, types.Array):
+        arr_dtype = data.dtype
+
+        def _sdc_take_array_impl(data, indexes):
+            res_size = len(indexes)
+            res_arr = numpy.empty(res_size, dtype=arr_dtype)
+            for i in numba.prange(res_size):
+                res_arr[i] = data[indexes[i]]
+            return res_arr
+
+        return _sdc_take_array_impl
+
+    elif data == string_array_type:
+        def _sdc_take_str_arr_impl(data, indexes):
+            res_size = len(indexes)
+            nan_mask = numpy.zeros(res_size, dtype=numpy.bool_)
+            num_total_bytes = 0
+            for i in numba.prange(res_size):
+                num_total_bytes += get_utf8_size(data[indexes[i]])
+                if isna(data, indexes[i]):
+                    nan_mask[i] = True
+
+            res_arr = pre_alloc_string_array(res_size, num_total_bytes)
+            for i in numpy.arange(res_size):
+                res_arr[i] = data[indexes[i]]
+                if nan_mask[i]:
+                    str_arr_set_na(res_arr, i)
+
+            return res_arr
+
+        return _sdc_take_str_arr_impl
+
+    return None

sdc/datatypes/hpat_pandas_dataframe_functions.py

@@ -38,6 +38,8 @@
 
 from numba import types
 from numba.special import literally
+from numba.typed import List, Dict
+
 from sdc.hiframes.pd_dataframe_ext import DataFrameType
 from sdc.hiframes.pd_series_type import SeriesType
 from sdc.utilities.sdc_typing_utils import (TypeChecker, check_index_is_numeric,
@@ -50,8 +52,10 @@
 from sdc.datatypes.hpat_pandas_dataframe_rolling_types import _hpat_pandas_df_rolling_init
 from sdc.datatypes.hpat_pandas_rolling_types import (
     gen_sdc_pandas_rolling_overload_body, sdc_pandas_rolling_docstring_tmpl)
+from sdc.datatypes.hpat_pandas_groupby_functions import init_dataframe_groupby
 from sdc.hiframes.pd_dataframe_ext import get_dataframe_data
 from sdc.utilities.utils import sdc_overload, sdc_overload_method, sdc_overload_attribute
+from sdc.hiframes.api import isna
 
 
 @sdc_overload_attribute(DataFrameType, 'index')
@@ -1337,3 +1341,32 @@ def pct_change_overload(df, periods=1, fill_method='pad', limit=None, freq=None)
     ser_par = {'periods': 'periods', 'fill_method': 'fill_method', 'limit': 'limit', 'freq': 'freq'}
 
     return sdc_pandas_dataframe_apply_columns(df, name, params, ser_par)
+
+
+@sdc_overload_method(DataFrameType, 'groupby')
+def sdc_pandas_dataframe_groupby(self, by=None, axis=0, level=None, as_index=True, sort=True,
+                                 group_keys=True, squeeze=False, observed=False):
+
+    if not isinstance(by, types.StringLiteral):
+        return None
+
+    column_id = self.columns.index(by.literal_value)
+    list_type = types.ListType(types.int64)
+    by_type = self.data[column_id].dtype
+
+    def sdc_pandas_dataframe_groupby_impl(self, by=None, axis=0, level=None, as_index=True, sort=True,
+                                          group_keys=True, squeeze=False, observed=False):
+
+        grouped = Dict.empty(by_type, list_type)
+        by_column_data = get_dataframe_data(self, column_id)
+        for i in numpy.arange(len(by_column_data)):
+            if isna(by_column_data, i):
+                continue
+            value = by_column_data[i]
+            group_list = grouped.get(value, List.empty_list(types.int64))
+            group_list.append(i)
+            grouped[value] = group_list
+
+        return init_dataframe_groupby(self, column_id, grouped, sort)
+
+    return sdc_pandas_dataframe_groupby_impl

sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py

@@ -28,7 +28,7 @@
 
 from numba.types import (float64, Boolean, Integer, Number, Omitted,
                          NoneType, StringLiteral, UnicodeType)
-from sdc.utilities.sdc_typing_utils import TypeChecker, params2list
+from sdc.utilities.sdc_typing_utils import TypeChecker, kwsparams2list
 from sdc.datatypes.hpat_pandas_dataframe_rolling_types import DataFrameRollingType
 from sdc.hiframes.pd_dataframe_ext import get_dataframe_data
 from sdc.hiframes.pd_dataframe_type import DataFrameType
@@ -95,7 +95,7 @@ def df_rolling_method_other_df_codegen(method_name, self, other, args=None, kws=
 
     rolling_params = df_rolling_params_codegen()
     method_kws = {k: k for k in kwargs}
-    impl_params = ['self'] + args + params2list(kwargs)
+    impl_params = ['self'] + args + kwsparams2list(kwargs)
     impl_params_as_str = ', '.join(impl_params)
 
     data_columns = {col: idx for idx, col in enumerate(self.data.columns)}
@@ -132,7 +132,7 @@ def df_rolling_method_other_df_codegen(method_name, self, other, args=None, kws=
         if col in common_columns:
             other_series = f'other_series_{col}'
             method_kws['other'] = other_series
-            method_params = ', '.join(args + params2list(method_kws))
+            method_params = ', '.join(args + kwsparams2list(method_kws))
             func_lines += [
                 f'  data_{col} = get_dataframe_data(self._data, {data_columns[col]})',
                 f'  other_data_{col} = get_dataframe_data(other, {other_columns[col]})',
@@ -189,7 +189,7 @@ def df_rolling_method_other_none_codegen(method_name, self, args=None, kws=None)
         args = args or []
         kwargs = kws or {}
 
-        impl_params = ['self'] + args + params2list(kwargs)
+        impl_params = ['self'] + args + kwsparams2list(kwargs)
         impl_params_as_str = ', '.join(impl_params)
 
         impl_name = f'_df_rolling_{_method_name}_other_none_impl'
@@ -223,7 +223,7 @@ def df_rolling_method_codegen(method_name, self, args=None, kws=None):
     args = args or []
     kwargs = kws or {}
 
-    impl_params = ['self'] + args + params2list(kwargs)
+    impl_params = ['self'] + args + kwsparams2list(kwargs)
     impl_params_as_str = ', '.join(impl_params)
 
     impl_name = f'_df_rolling_{method_name}_impl'