Adds Int64Index type and updates Series and DF methods to use it (#950)

* Adds Int64Index type and updates Series and DF methods to use it

Motivation: as part of the work on supporting common pandas indexes
a new type (Int64IndexType) representing pandas.Int64Index is added.
Boxing/unboxing of Series and DataFrames as well as common numpy-like
functions are changed accordingly to handle it.

* Fixing DateTime tests and PEP remarks

* Fixing review comments #1

Author: kozlov-alexey

sdc/__init__.py

@@ -48,6 +48,7 @@
 import sdc.datatypes.series.init
 
 import sdc.extensions.indexes.range_index_ext
+import sdc.extensions.indexes.int64_index_ext
 
 from ._version import get_versions
 

sdc/datatypes/common_functions.py

@@ -48,14 +48,17 @@
 from sdc.functions import numpy_like
 from sdc.str_arr_type import string_array_type, StringArrayType
 from sdc.datatypes.range_index_type import RangeIndexType
+from sdc.datatypes.int64_index_type import Int64IndexType
 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)
 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 (find_common_dtype_from_numpy_dtypes,
-                                            TypeChecker)
+from sdc.utilities.sdc_typing_utils import (
+                            find_common_dtype_from_numpy_dtypes,
+                            TypeChecker)
+from sdc.utilities.sdc_typing_utils import sdc_pandas_index_types
 
 
 class SDCLimitation(Exception):
@@ -71,18 +74,20 @@ def hpat_arrays_append(A, B):
 def hpat_arrays_append_overload(A, B):
     """Function for appending underlying arrays (A and B) or list/tuple of arrays B to an array A"""
 
-    A_is_range_index = isinstance(A, RangeIndexType)
-    B_is_range_index = isinstance(B, RangeIndexType)
-    if isinstance(A, (types.Array, RangeIndexType)):
-        if isinstance(B, (types.Array, RangeIndexType)):
+    use_A_array = isinstance(A, (RangeIndexType, Int64IndexType))
+    use_B_array = isinstance(B, (RangeIndexType, Int64IndexType))
+    if isinstance(A, (types.Array, RangeIndexType, Int64IndexType)):
+        if isinstance(B, (types.Array, RangeIndexType, Int64IndexType)):
             def _append_single_numeric_impl(A, B):
-                _A = A.values if A_is_range_index == True else A  # noqa
-                _B = B.values if B_is_range_index == True else B  # noqa
+                _A = A.values if use_A_array == True else A  # noqa
+                _B = B.values if use_B_array == True else B  # noqa
                 return numpy.concatenate((_A, _B,))
 
             return _append_single_numeric_impl
-        elif isinstance(B, (types.UniTuple, types.List)) and isinstance(B.dtype, (types.Array, RangeIndexType)):
-            B_dtype_is_range_index = isinstance(B.dtype, RangeIndexType)
+
+        elif (isinstance(B, (types.UniTuple, types.List))
+              and isinstance(B.dtype, (types.Array, RangeIndexType, Int64IndexType))):
+            B_dtype_is_index = isinstance(B.dtype, (RangeIndexType, Int64IndexType))
             numba_common_dtype = find_common_dtype_from_numpy_dtypes([A.dtype, B.dtype.dtype], [])
 
             # TODO: refactor to use numpy.concatenate when Numba supports building a tuple at runtime
@@ -92,10 +97,10 @@ def _append_list_numeric_impl(A, B):
                 new_data = numpy.empty(total_length, numba_common_dtype)
 
                 stop = len(A)
-                _A = numpy.array(A) if A_is_range_index == True else A  # noqa
+                _A = numpy.array(A) if use_A_array == True else A  # noqa
                 new_data[:stop] = _A
                 for arr in B:
-                    _arr = numpy.array(arr) if B_dtype_is_range_index == True else arr  # noqa
+                    _arr = arr.values if B_dtype_is_index == True else arr  # noqa
                     start = stop
                     stop = start + len(_arr)
                     new_data[start:stop] = _arr
@@ -218,12 +223,13 @@ def sdc_join_series_indexes_overload(left, right):
     """Function for joining arrays left and right in a way similar to pandas.join 'outer' algorithm"""
 
     # check that both operands are of types used for representing Pandas indexes
-    if not (isinstance(left, (types.Array, StringArrayType, RangeIndexType))
-            and isinstance(right, (types.Array, StringArrayType, RangeIndexType))):
+    if not (isinstance(left, sdc_pandas_index_types) and isinstance(right, sdc_pandas_index_types)
+            and not isinstance(left, types.NoneType)
+            and not isinstance(right, types.NoneType)):
         return None
 
-    convert_left = isinstance(left, RangeIndexType)
-    convert_right = isinstance(right, RangeIndexType)
+    convert_left = isinstance(left, (RangeIndexType, Int64IndexType))
+    convert_right = isinstance(right, (RangeIndexType, Int64IndexType))
 
     def _convert_to_arrays_impl(left, right):
         _left = left.values if convert_left == True else left  # noqa
@@ -243,10 +249,9 @@ def sdc_join_range_indexes_impl(left, right):
 
         return sdc_join_range_indexes_impl
 
-    elif isinstance(left, RangeIndexType) and isinstance(right, types.Array):
-        return _convert_to_arrays_impl
-
-    elif isinstance(left, types.Array) and isinstance(right, RangeIndexType):
+    elif (isinstance(left, (RangeIndexType, Int64IndexType, types.Array))
+          and isinstance(right, (RangeIndexType, Int64IndexType, types.Array))
+          and not (isinstance(left, types.Array) and isinstance(right, types.Array))):
         return _convert_to_arrays_impl
 
     # TODO: remove code duplication below and merge numeric and StringArray impls into one
@@ -618,13 +623,16 @@ def _sdc_take(data, indexes):
 @sdc_overload(_sdc_take)
 def _sdc_take_overload(data, indexes):
 
-    if not isinstance(data, (types.Array, StringArrayType, RangeIndexType)):
+    valid_data_types = (types.Array,) + sdc_pandas_index_types
+    if not (isinstance(data, valid_data_types) and not isinstance(data, types.NoneType)):
         return None
-    if not (isinstance(indexes, (types.Array, types.List))
+
+    if not (isinstance(indexes, (types.Array, types.List, Int64IndexType))
             and isinstance(indexes.dtype, (types.Integer, types.ListType))):
         return None
 
-    if isinstance(indexes.dtype, types.ListType) and isinstance(data, (types.Array, types.List, RangeIndexType)):
+    if (isinstance(indexes.dtype, types.ListType)
+            and isinstance(data, (types.Array, types.List, RangeIndexType, Int64IndexType))):
         arr_dtype = data.dtype
 
         def _sdc_take_list_impl(data, indexes):
@@ -677,7 +685,7 @@ def _sdc_take_list_str_impl(data, indexes):
 
         return _sdc_take_list_str_impl
 
-    elif isinstance(data, (types.Array, RangeIndexType)):
+    elif isinstance(data, (types.Array, RangeIndexType, Int64IndexType)):
         arr_dtype = data.dtype
 
         def _sdc_take_array_impl(data, indexes):
@@ -740,6 +748,7 @@ def sdc_reindex_series_overload(arr, index, name, by_index):
     """ Reindexes series data by new index following the logic of pandas.core.indexing.check_bool_indexer """
 
     range_indexes = isinstance(index, RangeIndexType) and isinstance(by_index, RangeIndexType)
+    int64_indexes = isinstance(index, Int64IndexType) and isinstance(by_index, Int64IndexType)
     data_dtype, index_dtype = arr.dtype, index.dtype
     data_is_str_arr = isinstance(arr.dtype, types.UnicodeType)
 
@@ -748,6 +757,8 @@ def sdc_reindex_series_impl(arr, index, name, by_index):
         # no reindexing is needed if indexes are equal
         if range_indexes == True:  # noqa
             equal_indexes = numpy_like.array_equal(index, by_index)
+        elif int64_indexes == True:  # noqa
+            equal_indexes = numpy_like.array_equal(index, by_index)
         else:
             equal_indexes = False
         if (index is by_index or equal_indexes):
@@ -772,10 +783,10 @@ def sdc_reindex_series_impl(arr, index, name, by_index):
                 map_index_to_position[value] = i
 
         index_mismatch = 0
-        # FIXME: TypingError in parfor step (wrong promotion to float64?) if prange is used
-        for i in numpy.arange(len(by_index)):
-            if by_index[i] in map_index_to_position:
-                pos_in_self = map_index_to_position[by_index[i]]
+        for i in numba.prange(len(by_index)):
+            val = by_index[i]
+            if val in map_index_to_position:
+                pos_in_self = map_index_to_position[val]
                 _res_data[i] = arr[pos_in_self]
                 if data_is_str_arr == True:  # noqa
                     res_data_nan_mask[i] = isna(arr, i)

sdc/datatypes/hpat_pandas_dataframe_functions.py

@@ -50,6 +50,7 @@
                                             gen_impl_generator, find_common_dtype_from_numpy_dtypes)
 from sdc.str_arr_ext import StringArrayType
 from sdc.datatypes.range_index_type import RangeIndexType
+from sdc.datatypes.int64_index_type import Int64IndexType
 
 from sdc.hiframes.pd_dataframe_type import DataFrameType
 from sdc.hiframes.pd_dataframe_ext import init_dataframe_internal, get_structure_maps
@@ -2257,7 +2258,7 @@ def sdc_pandas_dataframe_accessor_getitem(self, idx):
 
     if accessor == 'at':
         num_idx = (isinstance(idx[0], types.Number)
-                   and isinstance(self.dataframe.index, (types.Array, types.NoneType, RangeIndexType)))
+                   and isinstance(self.dataframe.index, (types.NoneType, RangeIndexType, Int64IndexType)))
         str_idx = (isinstance(idx[0], (types.UnicodeType, types.StringLiteral))
                    and isinstance(self.dataframe.index, StringArrayType))
         if isinstance(idx, types.Tuple) and isinstance(idx[1], types.StringLiteral):

sdc/datatypes/hpat_pandas_series_functions.py

@@ -53,6 +53,7 @@
                                             find_common_dtype_from_numpy_dtypes, has_literal_value,
                                             has_python_value)
 from sdc.datatypes.range_index_type import RangeIndexType
+from sdc.datatypes.int64_index_type import Int64IndexType
 from sdc.datatypes.common_functions import (sdc_join_series_indexes, sdc_arrays_argsort, sdc_reindex_series)
 from sdc.datatypes.hpat_pandas_rolling_types import (
     gen_sdc_pandas_rolling_overload_body, sdc_pandas_rolling_docstring_tmpl)
@@ -618,7 +619,7 @@ def sdc_pandas_series_setitem(self, idx, value):
         def sdc_pandas_series_setitem_no_reindexing_impl(self, idx, value):
 
             if assign_via_idx_mask == True:  # noqa
-                # FIXME_Numba#5157: using asarray since eq impl for RangeIndexType returns list
+                # FIXME_Numba#5157: using asarray since eq impl for index types returns list
                 _idx = numpy.asarray(self._index == idx)
             elif assign_via_idx_data == True:  # noqa
                 _idx = idx._data
@@ -4034,7 +4035,8 @@ def hpat_pandas_series_dropna(self, axis=0, inplace=False):
         ty_checker.raise_exc(inplace, 'bool', 'inplace')
 
     if (isinstance(self.data.dtype, types.Number)
-            and isinstance(self.index, (types.Number, types.NoneType, RangeIndexType))):
+            and (isinstance(self.index, types.NoneType)
+                 or isinstance(self.index.dtype, types.Number))):
         def hpat_pandas_series_dropna_impl(self, axis=0, inplace=False):
             index = self.index
             return numpy_like.dropna(self._data, index, self._name)

sdc/datatypes/int64_index_type.py

@@ -0,0 +1,65 @@
+# -*- coding: utf-8 -*-
+# *****************************************************************************
+# Copyright (c) 2020, Intel Corporation All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#     Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
+# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+# *****************************************************************************
+
+from numba import types
+from numba.extending import (
+    models,
+    register_model,
+    make_attribute_wrapper
+)
+
+
+class Int64IndexType(types.IterableType):
+    dtype = types.int64
+
+    def __init__(self, data, is_named=False):
+        self.data = data
+        self.is_named = is_named
+        super(Int64IndexType, self).__init__(
+            name='Int64IndexType({}, {})'.format(data, is_named))
+
+    @property
+    def iterator_type(self):
+        res = self.data.iterator_type
+        return res
+
+
+@register_model(Int64IndexType)
+class Int64IndexModel(models.StructModel):
+    def __init__(self, dmm, fe_type):
+
+        data_type = fe_type.data
+        name_type = types.unicode_type if fe_type.is_named else types.none
+        members = [
+            ('data', data_type),
+            ('name', name_type),
+        ]
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+
+make_attribute_wrapper(Int64IndexType, 'data', '_data')
+make_attribute_wrapper(Int64IndexType, 'name', '_name')

sdc/extensions/indexes/indexes_generic.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+# *****************************************************************************
+# Copyright (c) 2019-2020, Intel Corporation All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#     Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
+# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+# *****************************************************************************
+
+import numba
+import numpy as np
+import pandas as pd
+
+from numba import types
+
+
+def _check_dtype_param_type(dtype):
+    """ Returns True is dtype is a valid type for dtype parameter and False otherwise.
+        Used in RangeIndex ctor and other methods that take dtype parameter. """
+
+    valid_dtype_types = (types.NoneType, types.Omitted, types.UnicodeType, types.NumberClass)
+    return isinstance(dtype, valid_dtype_types) or dtype is None