diff --git a/python/xorbits/_mars/dataframe/datasource/index.py b/python/xorbits/_mars/dataframe/datasource/index.py index 30d97d6c4..e823d0bb7 100644 --- a/python/xorbits/_mars/dataframe/datasource/index.py +++ b/python/xorbits/_mars/dataframe/datasource/index.py @@ -65,7 +65,9 @@ def __call__(self, shape=None, chunk_size=None, inp=None, name=None, names=None) None, shape=shape, dtype=self.dtype, - index_value=parse_index(self.data, store_data=self.store_data), + # `store_data=True` to make sure + # when the `to_pandas()` is called, we can get the actual pandas value. + index_value=parse_index(self.data, store_data=True), name=name, names=names, raw_chunk_size=chunk_size, diff --git a/python/xorbits/core/tests/test_execution.py b/python/xorbits/core/tests/test_execution.py index 499d2db3a..010d9ea43 100644 --- a/python/xorbits/core/tests/test_execution.py +++ b/python/xorbits/core/tests/test_execution.py @@ -339,6 +339,36 @@ def test_getitem(setup): pd.testing.assert_frame_equal(result.to_pandas(), expected) +def test_column_index_setitem(setup): + import pandas as pd + + from ... import pandas as xpd + + data = {"a": [1, 2, 3], "b": [4, 5, 6]} + df = pd.DataFrame(data) + xdf = xpd.DataFrame(data) + + xdf.columns = xpd.Index(["c1", "d1"]) + df.columns = pd.Index(["c1", "d1"]) + pd.testing.assert_frame_equal(xdf.to_pandas(), df) + + xdf.columns = ["c2", "d2"] + df.columns = ["c2", "d2"] + pd.testing.assert_frame_equal(xdf.to_pandas(), df) + + xdf.columns = pd.Index(["c3", "d3"]) + df.columns = pd.Index(["c3", "d3"]) + pd.testing.assert_frame_equal(xdf.to_pandas(), df) + + xdf.index = ["x1", "y1", "z1"] + df.index = ["x1", "y1", "z1"] + pd.testing.assert_frame_equal(xdf.to_pandas(), df) + + xdf.index = xpd.Index(["x2", "y2", "z2"]) + df.index = pd.Index(["x2", "y2", "z2"]) + pd.testing.assert_frame_equal(xdf.to_pandas(), df) + + def test_execution_with_process_exit_message(mocker): import numpy as np from xoscar.errors import ServerClosed