IntelPython
diff --git a/‎sdc/datatypes/hpat_pandas_dataframe_functions.py‎
Lines changed: 290 additions & 301 deletions b/‎sdc/datatypes/hpat_pandas_dataframe_functions.py‎
Lines changed: 290 additions & 301 deletions
diff --git a/‎sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py‎
Lines changed: 19 additions & 8 deletions b/‎sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py‎
Lines changed: 19 additions & 8 deletions
diff --git a/‎sdc/datatypes/hpat_pandas_groupby_functions.py‎
Lines changed: 26 additions & 19 deletions b/‎sdc/datatypes/hpat_pandas_groupby_functions.py‎
Lines changed: 26 additions & 19 deletions
diff --git a/‎sdc/hiframes/boxing.py‎
Lines changed: 74 additions & 36 deletions b/‎sdc/hiframes/boxing.py‎
Lines changed: 74 additions & 36 deletions
@@ -123,19 +123,25 @@ def df_rolling_method_other_df_codegen(method_name, self, other, args=None, kws=
             f'    raise ValueError("Method rolling.{method_name}(). The object pairwise\\n expected: False, None")'
         ]
 
-    data_length = 'len(self._data._data[0])' if data_columns else '0'
-    other_length = 'len(other._data[0])' if other_columns else '0'
+    data_length = 'len(self._data._data[0][0])' if data_columns else '0'
+    other_length = 'len(other._data[0][0])' if other_columns else '0'
     func_lines += [f'  length = max([{data_length}, {other_length}])']
 
     for col in all_columns:
         res_data = f'result_data_{col}'
         if col in common_columns:
+            col_loc = self.data.column_loc[col]
+            type_id, col_id = col_loc.type_id, col_loc.col_id
+            other_col_loc = other.column_loc[col]
+            other_type_id = other_col_loc.type_id
+            other_col_id = other_col_loc.col_id
+
             other_series = f'other_series_{col}'
             method_kws['other'] = other_series
             method_params = ', '.join(args + kwsparams2list(method_kws))
             func_lines += [
-                f'  data_{col} = self._data._data[{data_columns[col]}]',
-                f'  other_data_{col} = other._data[{other_columns[col]}]',
+                f'  data_{col} = self._data._data[{type_id}][{col_id}]',
+                f'  other_data_{col} = other._data[{other_type_id}][{other_col_id}]',
                 f'  series_{col} = pandas.Series(data_{col})',
                 f'  {other_series} = pandas.Series(other_data_{col})',
                 f'  rolling_{col} = series_{col}.rolling({rolling_params})',
@@ -158,16 +164,18 @@ def df_rolling_method_other_df_codegen(method_name, self, other, args=None, kws=
     return func_text, global_vars
 
 
-def df_rolling_method_main_codegen(method_params, df_columns, method_name):
+def df_rolling_method_main_codegen(method_params, df_columns, column_loc, method_name):
     rolling_params = df_rolling_params_codegen()
     method_params_as_str = ', '.join(method_params)
 
     results = []
     func_lines = []
     for idx, col in enumerate(df_columns):
+        col_loc = column_loc[col]
+        type_id, col_id = col_loc.type_id, col_loc.col_id
         res_data = f'result_data_{col}'
         func_lines += [
-            f'  data_{col} = self._data._data[{idx}]',
+            f'  data_{col} = self._data._data[{type_id}][{col_id}]',
             f'  series_{col} = pandas.Series(data_{col})',
             f'  rolling_{col} = series_{col}.rolling({rolling_params})',
             f'  result_{col} = rolling_{col}.{method_name}({method_params_as_str})',
@@ -204,7 +212,9 @@ def df_rolling_method_other_none_codegen(method_name, self, args=None, kws=None)
                 f'    raise ValueError("Method rolling.{_method_name}(). The object pairwise\\n expected: False")'
             ]
         method_params = args + ['{}={}'.format(k, k) for k in kwargs if k != 'other']
-        func_lines += df_rolling_method_main_codegen(method_params, self.data.columns, method_name)
+        func_lines += df_rolling_method_main_codegen(method_params, self.data.columns, self.data.column_loc,
+                                                     method_name)
+
         func_text = '\n'.join(func_lines)
 
         global_vars = {'pandas': pandas}
@@ -229,7 +239,8 @@ def df_rolling_method_codegen(method_name, self, args=None, kws=None):
     func_lines = [f'def {impl_name}({impl_params_as_str}):']
 
     method_params = args + ['{}={}'.format(k, k) for k in kwargs]
-    func_lines += df_rolling_method_main_codegen(method_params, self.data.columns, method_name)
+    func_lines += df_rolling_method_main_codegen(method_params, self.data.columns,
+                                                 self.data.column_loc, method_name)
     func_text = '\n'.join(func_lines)
 
     global_vars = {'pandas': pandas}
 
@@ -38,16 +38,13 @@
 from numba.core import cgutils
 from numba.extending import intrinsic
 from numba.core.registry import cpu_target
-from numba.typed import List, Dict
 from numba.core.typing import signature
 from numba import literally
 
 from sdc.datatypes.common_functions import sdc_arrays_argsort, _sdc_asarray, _sdc_take
 from sdc.datatypes.hpat_pandas_groupby_types import DataFrameGroupByType, SeriesGroupByType
 from sdc.utilities.sdc_typing_utils import TypeChecker, kwsparams2list, sigparams2list
-from sdc.utilities.utils import (sdc_overload, sdc_overload_method, sdc_register_jitable,
-                                 sdc_register_jitable)
-from sdc.hiframes.pd_dataframe_ext import get_dataframe_data
+from sdc.utilities.utils import (sdc_overload, sdc_overload_method, sdc_register_jitable)
 from sdc.hiframes.pd_series_type import SeriesType
 from sdc.str_ext import string_type
 
@@ -155,7 +152,14 @@ def sdc_pandas_dataframe_getitem(self, idx):
             and all(isinstance(a, types.StringLiteral) for a in idx))):
 
         by_col_id_literal = self.col_id.literal_value
-        target_col_id_literal = self.parent.columns.index(idx.literal_value) if idx_is_literal_str else None
+        by_col_loc = self.parent.column_loc[self.parent.columns[by_col_id_literal]]
+        by_type_id, by_col_id = by_col_loc.type_id, by_col_loc.col_id
+
+        if idx_is_literal_str:
+            target_col_id_literal = self.parent.columns.index(idx.literal_value)
+            target_col_loc = self.parent.column_loc[self.parent.columns[target_col_id_literal]]
+            target_type_id, target_col_id = target_col_loc.type_id, target_col_loc.col_id
+
         def sdc_pandas_dataframe_getitem_common_impl(self, idx):
 
             # calling getitem twice raises IndexError, just as in pandas
@@ -165,10 +169,10 @@ def sdc_pandas_dataframe_getitem_common_impl(self, idx):
             if idx_is_literal_str == True:  # noqa
                 # no need to pass index into this series, as we group by array
                 target_series = pandas.Series(
-                    data=self._parent._data[target_col_id_literal],
+                    data=self._parent._data[target_type_id][target_col_id],
                     name=self._parent._columns[target_col_id_literal]
                 )
-                by_arr_data = self._parent._data[by_col_id_literal]
+                by_arr_data = self._parent._data[by_type_id][by_col_id]
                 return init_series_groupby(target_series, by_arr_data, self._data, self._sort)
             else:
                 return init_dataframe_groupby(self._parent, by_col_id_literal, self._data, self._sort, idx)
@@ -184,8 +188,8 @@ def sdc_pandas_dataframe_getitem_idx_unicode_str_impl(self, idx):
     return None
 
 
-def _sdc_pandas_groupby_generic_func_codegen(func_name, columns, func_params, defaults, impl_params):
-
+def _sdc_pandas_groupby_generic_func_codegen(func_name, columns, column_loc,
+                                             func_params, defaults, impl_params):
     all_params_as_str = ', '.join(sigparams2list(func_params, defaults))
     extra_impl_params = ', '.join(kwsparams2list(impl_params))
 
@@ -204,15 +208,18 @@ def _sdc_pandas_groupby_generic_func_codegen(func_name, columns, func_params, de
     ]
 
     # TODO: remove conversion from Numba typed.List to reflected one while creating group_arr_{i}
-    func_lines.extend(['\n'.join([
-        f'  result_data_{i} = numpy.empty(res_index_len, dtype=res_arrays_dtypes[{i}])',
-        f'  column_data_{i} = {df}._data[{column_ids[i]}]',
-        f'  for j in numpy.arange(res_index_len):',
-        f'    idx = argsorted_index[j] if {groupby_param_sort} else j',
-        f'    group_arr_{i} = _sdc_take(column_data_{i}, list({groupby_dict}[group_keys[idx]]))',
-        f'    group_series_{i} = pandas.Series(group_arr_{i})',
-        f'    result_data_{i}[j] = group_series_{i}.{func_name}({extra_impl_params})',
-    ]) for i in range(len(columns))])
+    for i in range(len(columns)):
+        col_loc = column_loc[column_names[i]]
+        type_id, col_id = col_loc.type_id, col_loc.col_id
+        func_lines += [
+            f'  result_data_{i} = numpy.empty(res_index_len, dtype=res_arrays_dtypes[{i}])',
+            f'  column_data_{i} = {df}._data[{type_id}][{col_id}]',
+            f'  for j in numpy.arange(res_index_len):',
+            f'    idx = argsorted_index[j] if {groupby_param_sort} else j',
+            f'    group_arr_{i} = _sdc_take(column_data_{i}, list({groupby_dict}[group_keys[idx]]))',
+            f'    group_series_{i} = pandas.Series(group_arr_{i})',
+            f'    result_data_{i}[j] = group_series_{i}.{func_name}({extra_impl_params})',
+        ]
 
     data = ', '.join(f'\'{column_names[i]}\': result_data_{i}' for i in range(len(columns)))
     func_lines.extend(['\n'.join([
@@ -314,7 +321,7 @@ def sdc_pandas_dataframe_groupby_apply_func(self, func_name, func_args, defaults
 
     groupby_func_name = f'_dataframe_groupby_{func_name}_impl'
     func_text, global_vars = _sdc_pandas_groupby_generic_func_codegen(
-        func_name, subject_columns, func_args, defaults, impl_args)
+        func_name, subject_columns, self.parent.column_loc, func_args, defaults, impl_args)
 
     # capture result column types into generated func context
     global_vars['res_arrays_dtypes'] = res_arrays_dtypes
 
@@ -48,6 +48,8 @@
 from sdc.hiframes.pd_series_ext import SeriesType
 from sdc.hiframes.pd_series_type import _get_series_array_type
 
+from sdc.hiframes.pd_dataframe_ext import get_structure_maps
+
 from .. import hstr_ext
 import llvmlite.binding as ll
 from llvmlite import ir as lir
@@ -58,12 +60,14 @@
 
 @typeof_impl.register(pd.DataFrame)
 def typeof_pd_dataframe(val, c):
+
     col_names = tuple(val.columns.tolist())
     # TODO: support other types like string and timestamp
     col_types = get_hiframes_dtypes(val)
     index_type = _infer_index_type(val.index)
+    column_loc, _, _ = get_structure_maps(col_types, col_names)
 
-    return DataFrameType(col_types, index_type, col_names, True)
+    return DataFrameType(col_types, index_type, col_names, True, column_loc=column_loc)
 
 
 # register series types for import
@@ -86,21 +90,55 @@ def unbox_dataframe(typ, val, c):
     # create dataframe struct and store values
     dataframe = cgutils.create_struct_proxy(typ)(c.context, c.builder)
 
-    column_tup = c.context.make_tuple(
-        c.builder, types.UniTuple(string_type, n_cols), column_strs)
+    errorptr = cgutils.alloca_once_value(c.builder, cgutils.false_bit)
 
-    # this unboxes all DF columns so that no column unboxing occurs later
-    for col_ind in range(n_cols):
-        series_obj = c.pyapi.object_getattr_string(val, typ.columns[col_ind])
-        arr_obj = c.pyapi.object_getattr_string(series_obj, "values")
-        ty_series = typ.data[col_ind]
-        if isinstance(ty_series, types.Array):
-            native_val = unbox_array(typ.data[col_ind], arr_obj, c)
-        elif ty_series == string_array_type:
-            native_val = unbox_str_series(string_array_type, series_obj, c)
+    col_list_type = types.List(string_type)
+    ok, inst = listobj.ListInstance.allocate_ex(c.context, c.builder, col_list_type, n_cols)
 
-        dataframe.data = c.builder.insert_value(
-            dataframe.data, native_val.value, col_ind)
+    with c.builder.if_else(ok, likely=True) as (if_ok, if_not_ok):
+        with if_ok:
+            inst.size = c.context.get_constant(types.intp, n_cols)
+            for i, column_str in enumerate(column_strs):
+                inst.setitem(c.context.get_constant(types.intp, i), column_str, incref=False)
+            dataframe.columns = inst.value
+
+        with if_not_ok:
+            c.builder.store(cgutils.true_bit, errorptr)
+
+    # If an error occurred, drop the whole native list
+    with c.builder.if_then(c.builder.load(errorptr)):
+        c.context.nrt.decref(c.builder, col_list_type, inst.value)
+
+    _, data_typs_map, types_order = get_structure_maps(typ.data, typ.columns)
+
+    for col_typ in types_order:
+        type_id, col_indices = data_typs_map[col_typ]
+        n_type_cols = len(col_indices)
+        list_type = types.List(col_typ)
+        ok, inst = listobj.ListInstance.allocate_ex(c.context, c.builder, list_type, n_type_cols)
+
+        with c.builder.if_else(ok, likely=True) as (if_ok, if_not_ok):
+            with if_ok:
+                inst.size = c.context.get_constant(types.intp, n_type_cols)
+                for i, col_idx in enumerate(col_indices):
+                    series_obj = c.pyapi.object_getattr_string(val, typ.columns[col_idx])
+                    arr_obj = c.pyapi.object_getattr_string(series_obj, "values")
+                    ty_series = typ.data[col_idx]
+                    if isinstance(ty_series, types.Array):
+                        native_val = unbox_array(typ.data[col_idx], arr_obj, c)
+                    elif ty_series == string_array_type:
+                        native_val = unbox_str_series(string_array_type, series_obj, c)
+
+                    inst.setitem(c.context.get_constant(types.intp, i), native_val.value, incref=False)
+
+                dataframe.data = c.builder.insert_value(dataframe.data, inst.value, type_id)
+
+            with if_not_ok:
+                c.builder.store(cgutils.true_bit, errorptr)
+
+        # If an error occurred, drop the whole native list
+        with c.builder.if_then(c.builder.load(errorptr)):
+            c.context.nrt.decref(c.builder, list_type, inst.value)
 
     # TODO: support unboxing index
     if typ.index == types.none:
@@ -113,7 +151,6 @@ def unbox_dataframe(typ, val, c):
         index_data = c.pyapi.object_getattr_string(index_obj, "_data")
         dataframe.index = unbox_array(typ.index, index_data, c).value
 
-    dataframe.columns = column_tup
     dataframe.parent = val
 
     # increase refcount of stored values
@@ -122,7 +159,7 @@ def unbox_dataframe(typ, val, c):
         for var in column_strs:
             c.context.nrt.incref(c.builder, string_type, var)
 
-    return NativeValue(dataframe._getvalue())
+    return NativeValue(dataframe._getvalue(), is_error=c.builder.load(errorptr))
 
 
 def get_hiframes_dtypes(df):
@@ -202,15 +239,10 @@ def box_dataframe(typ, val, c):
     context = c.context
     builder = c.builder
 
-    n_cols = len(typ.columns)
     col_names = typ.columns
     arr_typs = typ.data
-    dtypes = [a.dtype for a in arr_typs]  # TODO: check Categorical
 
     dataframe = cgutils.create_struct_proxy(typ)(context, builder, value=val)
-    col_arrs = [builder.extract_value(dataframe.data, i) for i in range(n_cols)]
-    # df unboxed from Python
-    has_parent = cgutils.is_not_null(builder, dataframe.parent)
 
     pyapi = c.pyapi
     # gil_state = pyapi.gil_ensure()  # acquire GIL
@@ -219,28 +251,31 @@ def box_dataframe(typ, val, c):
     class_obj = pyapi.import_module_noblock(mod_name)
     df_dict = pyapi.dict_new()
 
-    for i, cname, arr, arr_typ, dtype in zip(range(n_cols), col_names, col_arrs, arr_typs, dtypes):
+    arrays_list_objs = {}
+    for cname, arr_typ in zip(col_names, arr_typs):
         # df['cname'] = boxed_arr
         # TODO: datetime.date, DatetimeIndex?
         name_str = context.insert_const_string(c.builder.module, cname)
         cname_obj = pyapi.string_from_string(name_str)
 
-        if dtype == string_type:
-            arr_obj = box_str_arr(arr_typ, arr, c)
-        elif isinstance(arr_typ, Categorical):
-            arr_obj = box_Categorical(arr_typ, arr, c)
-            # context.nrt.incref(builder, arr_typ, arr)
-        elif dtype == types.List(string_type):
-            arr_obj = box_list(list_string_array_type, arr, c)
-            # context.nrt.incref(builder, arr_typ, arr)  # TODO required?
-            # pyapi.print_object(arr_obj)
-        else:
-            arr_obj = box_array(arr_typ, arr, c)
-            # TODO: is incref required?
-            # context.nrt.incref(builder, arr_typ, arr)
+        col_loc = typ.column_loc[cname]
+        type_id, col_id = col_loc.type_id, col_loc.col_id
+
+        # dataframe.data looks like a tuple(list(array))
+        # e.g. ([array(int64, 1d, C), array(int64, 1d, C)], [array(float64, 1d, C)])
+        arrays_list_obj = arrays_list_objs.get(type_id)
+        if arrays_list_obj is None:
+            list_typ = types.List(arr_typ)
+            # extracting list from the tuple
+            list_val = builder.extract_value(dataframe.data, type_id)
+            # getting array from the list to box it then
+            arrays_list_obj = box_list(list_typ, list_val, c)
+            arrays_list_objs[type_id] = arrays_list_obj
+
+        # PyList_GetItem returns borrowed reference
+        arr_obj = pyapi.list_getitem(arrays_list_obj, col_id)
         pyapi.dict_setitem(df_dict, cname_obj, arr_obj)
 
-        pyapi.decref(arr_obj)
         pyapi.decref(cname_obj)
 
     df_obj = pyapi.call_method(class_obj, "DataFrame", (df_dict,))
@@ -252,6 +287,9 @@ def box_dataframe(typ, val, c):
         pyapi.object_setattr_string(df_obj, 'index', arr_obj)
         pyapi.decref(arr_obj)
 
+    for arrays_list_obj in arrays_list_objs.values():
+        pyapi.decref(arrays_list_obj)
+
     pyapi.decref(class_obj)
     # pyapi.gil_release(gil_state)    # release GIL
     return df_obj