Fixing perf issue with getitem by Boolean mask (#790)

Author: kozlov-alexey

sdc/datatypes/common_functions.py

@@ -629,6 +629,24 @@ def _sdc_take_str_arr_impl(data, indexes):
 
         return _sdc_take_str_arr_impl
 
+    elif (isinstance(data, types.RangeType) and isinstance(data.dtype, types.Integer)):
+        arr_dtype = data.dtype
+
+        def _sdc_take_array_impl(data, indexes):
+            res_size = len(indexes)
+            index_errors = 0
+            res_arr = numpy.empty(res_size, dtype=arr_dtype)
+            for i in numba.prange(res_size):
+                value = data.start + data.step * indexes[i]
+                if value >= data.stop:
+                    index_errors += 1
+                res_arr[i] = value
+            if index_errors:
+                raise IndexError("_sdc_take: index out-of-bounds")
+            return res_arr
+
+        return _sdc_take_array_impl
+
     return None
 
 

sdc/datatypes/hpat_pandas_dataframe_functions.py

@@ -1454,13 +1454,16 @@ def df_length_expr(self):
     return '0'
 
 
-def df_index_expr(self, length_expr=None):
+def df_index_expr(self, length_expr=None, as_range=False):
     """Generate expression to get or create index of DF"""
     if isinstance(self.index, types.NoneType):
         if length_expr is None:
             length_expr = df_length_expr(self)
 
-        return f'numpy.arange({length_expr})'
+        if as_range:
+            return f'range({length_expr})'
+        else:
+            return f'numpy.arange({length_expr})'
 
     return 'self._index'
 
@@ -1507,12 +1510,13 @@ def df_getitem_bool_series_idx_main_codelines(self, idx):
     # optimization for default indexes in df and idx when index alignment is trivial
     if (isinstance(self.index, types.NoneType) and isinstance(idx.index, types.NoneType)):
         func_lines = [f'  length = {df_length_expr(self)}',
+                      f'  self_index = {df_index_expr(self, as_range=True)}',
                       f'  if length > len(idx):',
                       f'    msg = "Unalignable boolean Series provided as indexer " + \\',
                       f'          "(index of the boolean Series and of the indexed object do not match)."',
                       f'    raise IndexingError(msg)',
                       f'  # do not trim idx._data to length as getitem_by_mask handles such case',
-                      f'  res_index = getitem_by_mask(self.index, idx._data)',
+                      f'  res_index = getitem_by_mask(self_index, idx._data)',
                       f'  # df index is default, same as positions so it can be used in take']
         results = []
         for i, col in enumerate(self.columns):
@@ -1532,7 +1536,7 @@ def df_getitem_bool_series_idx_main_codelines(self, idx):
                       f'  self_index = self.index',
                       f'  idx_reindexed = sdc_reindex_series(idx._data, idx.index, idx._name, self_index)',
                       f'  res_index = getitem_by_mask(self_index, idx_reindexed._data)',
-                      f'  selected_pos = getitem_by_mask(numpy.arange(length), idx_reindexed._data)']
+                      f'  selected_pos = getitem_by_mask(range(length), idx_reindexed._data)']
 
         results = []
         for i, col in enumerate(self.columns):
@@ -1553,11 +1557,13 @@ def df_getitem_bool_series_idx_main_codelines(self, idx):
 
 def df_getitem_bool_array_idx_main_codelines(self, idx):
     """Generate main code lines for df.getitem"""
+
     func_lines = [f'  length = {df_length_expr(self)}',
                   f'  if length != len(idx):',
                   f'    raise ValueError("Item wrong length.")',
-                  f'  taken_pos = getitem_by_mask(numpy.arange(length), idx)',
-                  f'  res_index = sdc_take(self.index, taken_pos)']
+                  f'  self_index = {df_index_expr(self, as_range=True)}',
+                  f'  taken_pos = getitem_by_mask(self_index, idx)',
+                  f'  res_index = sdc_take(self_index, taken_pos)']
     results = []
     for i, col in enumerate(self.columns):
         res_data = f'res_data_{i}'
@@ -1635,12 +1641,13 @@ def df_getitem_bool_series_idx_codegen(self, idx):
     Example of generated implementation with provided index:
         def _df_getitem_bool_series_idx_impl(self, idx):
           length = len(self._data[0])
+          self_index = range(len(self._data[0]))
           if length > len(idx):
             msg = "Unalignable boolean Series provided as indexer " + \
                   "(index of the boolean Series and of the indexed object do not match)."
             raise IndexingError(msg)
           # do not trim idx._data to length as getitem_by_mask handles such case
-          res_index = getitem_by_mask(self.index, idx._data)
+          res_index = getitem_by_mask(self_index, idx._data)
           # df index is default, same as positions so it can be used in take
           data_0 = self._data[0]
           res_data_0 = sdc_take(data_0, res_index)
@@ -1667,8 +1674,9 @@ def _df_getitem_bool_array_idx_impl(self, idx):
           length = len(self._data[0])
           if length != len(idx):
             raise ValueError("Item wrong length.")
-          taken_pos = getitem_by_mask(numpy.arange(length), idx)
-          res_index = sdc_take(self.index, taken_pos)
+          self_index = range(len(self._data[0]))
+          taken_pos = getitem_by_mask(self_index, idx)
+          res_index = sdc_take(self_index, taken_pos)
           data_0 = self._data[0]
           res_data_0 = sdc_take(data_0, taken_pos)
           data_1 = self._data[1]

sdc/datatypes/hpat_pandas_series_functions.py

@@ -430,7 +430,7 @@ def hpat_pandas_series_getitem_idx_bool_indexer_impl(self, idx):
 
                 return pandas.Series(
                     data=numpy_like.getitem_by_mask(self._data, idx._data),
-                    index=numpy_like.getitem_by_mask(self.index, idx._data),
+                    index=numpy_like.getitem_by_mask(range(len(self)), idx._data),
                     name=self._name
                 )
             else:

sdc/functions/numpy_like.py

@@ -900,9 +900,30 @@ def getitem_by_mask(arr, idx):
 
 @sdc_overload(getitem_by_mask)
 def getitem_by_mask_overload(arr, idx):
-    dtype = arr.dtype
+    """
+    Creates a new array from arr by selecting elements indicated by Boolean mask idx.
+
+    Parameters
+    -----------
+    arr: :obj:`Array` or :obj:`Range`
+        Input array or range
+    idx: :obj:`Array` of dtype :class:`bool`
+        Boolean mask
+
+    Returns
+    -------
+    :obj:`Array` of the same dtype as arr
+        Array with only elements indicated by mask left
+
+    """
+
+    is_range = isinstance(arr, types.RangeType) and isinstance(arr.dtype, types.Integer)
     is_str_arr = arr == string_array_type
+    if not (isinstance(arr, types.Array) or is_str_arr or is_range):
+        return
 
+    res_dtype = arr.dtype
+    is_str_arr = arr == string_array_type
     def getitem_by_mask_impl(arr, idx):
         chunks = parallel_chunks(len(arr))
         arr_len = numpy.empty(len(chunks), dtype=numpy.int64)
@@ -921,15 +942,19 @@ def getitem_by_mask_impl(arr, idx):
             result_data = [''] * length
             result_nan_mask = numpy.empty(shape=length, dtype=types.bool_)
         else:
-            result_data = numpy.empty(shape=length, dtype=dtype)
+            result_data = numpy.empty(shape=length, dtype=res_dtype)
         for i in prange(len(chunks)):
             chunk = chunks[i]
             new_start = int(sum(arr_len[0:i]))
             current_pos = new_start
 
             for j in range(chunk.start, chunk.stop):
                 if idx[j]:
-                    result_data[current_pos] = arr[j]
+                    if is_range == True:  # noqa
+                        value = arr.start + arr.step * j
+                    else:
+                        value = arr[j]
+                    result_data[current_pos] = value
                     if is_str_arr == True:  # noqa
                         result_nan_mask[current_pos] = isna(arr, j)
                     current_pos += 1