Re-implement df.index based on new structure (#853)

densmirn · web-flow · commit 02d56d06b6b9 · 2020-05-22T17:28:30.000+03:00
diff --git a/sdc/datatypes/hpat_pandas_dataframe_functions.py b/sdc/datatypes/hpat_pandas_dataframe_functions.py
@@ -100,7 +100,7 @@ def hpat_pandas_dataframe_index(df):
         empty_df = not df.columns
 
         def hpat_pandas_df_index_none_impl(df):
-            df_len = len(df._data[0]) if empty_df == False else 0  # noqa
+            df_len = len(df._data[0][0]) if empty_df == False else 0  # noqa
 
             return numpy.arange(df_len)
 
diff --git a/sdc/tests/test_dataframe.py b/sdc/tests/test_dataframe.py
@@ -741,7 +741,7 @@ def test_impl(df):
         np.testing.assert_array_equal(sdc_func(df), test_impl(df))
 
     @skip_sdc_jit
-    @dfRefactoringNotImplemented
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame unboxing
     def test_index_attribute(self):
         index_to_test = [[1, 2, 3, 4, 5],
                          [.1, .2, .3, .4, .5],
@@ -757,7 +757,7 @@ def test_index_attribute(self):
                 self._test_df_index(df)
 
     @skip_sdc_jit
-    @dfRefactoringNotImplemented
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame unboxing
     def test_index_attribute_empty(self):
         n = 5
         np.random.seed(0)
@@ -768,11 +768,45 @@ def test_index_attribute_empty(self):
         self._test_df_index(df)
 
     @skip_sdc_jit
-    @dfRefactoringNotImplemented
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame unboxing
     def test_index_attribute_empty_df(self):
         df = pd.DataFrame()
         self._test_df_index(df)
 
+    def test_index_attribute_no_unboxing(self):
+        def test_impl(n, index):
+            np.random.seed(0)
+            df = pd.DataFrame({
+                'A': np.ones(n),
+                'B': np.random.ranf(n)
+            }, index=index)
+            return df.index
+
+        sdc_impl = self.jit(test_impl)
+        index_to_test = [
+            [1, 2, 3, 4, 5],
+            [.1, .2, .3, .4, .5],
+            ['a', 'b', 'c', 'd', 'e']
+        ]
+        for index in index_to_test:
+            with self.subTest(index=index):
+                n = len(index)
+                jit_result = sdc_impl(n, index)
+                ref_result = test_impl(n, index)
+                np.testing.assert_array_equal(jit_result, ref_result)
+
+    def test_index_attribute_default_no_unboxing(self):
+        def test_impl(n):
+            np.random.seed(0)
+            df = pd.DataFrame({
+                'A': np.ones(n),
+                'B': np.random.ranf(n)
+            })
+            return df.index
+
+        sdc_impl = self.jit(test_impl)
+        np.testing.assert_array_equal(sdc_impl(10), test_impl(10))
+
     @skip_sdc_jit
     @skip_numba_jit
     def test_df_apply(self):
