Closes #2897 DataFrame.corr returns dataframe without index (#3012)

Co-authored-by: Amanda Potts <ajpotts@users.noreply.github.com>
Bears-R-Us · Mar 1, 2024 · 3b9f8f7 · 3b9f8f7
1 parent 99245b2
commit 3b9f8f7
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 8 deletions.
diff --git a/PROTO_tests/tests/dataframe_test.py b/PROTO_tests/tests/dataframe_test.py
@@ -872,6 +872,18 @@ def test_isin(self):
         assert test_df["col_A"].to_list() == [True, True]
         assert test_df["col_B"].to_list() == [False, False]
 
+    def test_corr(self):
+        df = ak.DataFrame({'col1': [1, 2], 'col2': [-1, -2]})
+        corr = df.corr()
+        pd_corr = df.to_pandas().corr()
+        assert_frame_equal(corr.to_pandas(retain_index=True), pd_corr)
+
+        for i in range(5):
+            df = ak.DataFrame({'col1': ak.randint(0, 10, 10), 'col2': ak.randint(0, 10, 10)})
+            corr = df.corr()
+            pd_corr = df.to_pandas().corr()
+            assert_frame_equal(corr.to_pandas(retain_index=True), pd_corr)
+
     def test_multiindex_compat(self):
         # Added for testing Issue #1505
         df = ak.DataFrame({"a": ak.arange(10), "b": ak.arange(10), "c": ak.arange(10)})

diff --git a/arkouda/dataframe.py b/arkouda/dataframe.py
@@ -3937,13 +3937,13 @@ def corr(self) -> DataFrame:
 
         >>> corr = df.corr()
 
-        +----+--------+--------+
-        |    |   col1 |   col2 |
-        +====+========+========+
-        |  0 |      1 |     -1 |
-        +----+--------+--------+
-        |  1 |     -1 |      1 |
-        +----+--------+--------+
+        +------+--------+--------+
+        |      |   col1 |   col2 |
+        +======+========+========+
+        | col1 |      1 |     -1 |
+        +------+--------+--------+
+        | col2 |     -1 |      1 |
+        +------+--------+--------+
 
         """
 
@@ -3959,7 +3959,10 @@ def numeric_help(d):
         }
 
         ret_dict = json.loads(generic_msg(cmd="corrMatrix", args=args))
-        return DataFrame({c: create_pdarray(ret_dict[c]) for c in self.columns.values})
+        return DataFrame(
+            {c: create_pdarray(ret_dict[c]) for c in self.columns.values},
+            index=array(self.columns.values),
+        )
 
     @typechecked
     def merge(

diff --git a/tests/dataframe_test.py b/tests/dataframe_test.py
@@ -943,6 +943,18 @@ def test_isin(self):
         self.assertListEqual(test_df["col_A"].to_list(), [True, True])
         self.assertListEqual(test_df["col_B"].to_list(), [False, False])
 
+    def test_corr(self):
+        df = ak.DataFrame({'col1': [1, 2], 'col2': [-1, -2]})
+        corr = df.corr()
+        pd_corr = df.to_pandas().corr()
+        assert_frame_equal(corr.to_pandas(retain_index=True), pd_corr)
+
+        for i in range(5):
+            df = ak.DataFrame({'col1': ak.randint(0, 10, 10), 'col2': ak.randint(0, 10, 10)})
+            corr = df.corr()
+            pd_corr = df.to_pandas().corr()
+            assert_frame_equal(corr.to_pandas(retain_index=True), pd_corr)
+
     def test_multiindex_compat(self):
         # Added for testing Issue #1505
         df = ak.DataFrame({"a": ak.arange(10), "b": ak.arange(10), "c": ak.arange(10)})