From b48e2dc6208eaea990de056e03701f5dd7235b67 Mon Sep 17 00:00:00 2001 From: Weizheng Lu Date: Wed, 18 Dec 2024 18:59:16 +0800 Subject: [PATCH] BUG: set column name when using `xorbits.pandas.Index` (#836) Co-authored-by: ChengjieLi --- .../_mars/dataframe/datasource/index.py | 4 ++- python/xorbits/core/tests/test_execution.py | 31 ++++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/python/xorbits/_mars/dataframe/datasource/index.py b/python/xorbits/_mars/dataframe/datasource/index.py index 3820994e4..a6cbb3d3e 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 5ad57f1dd..45a9c438e 100644 --- a/python/xorbits/core/tests/test_execution.py +++ b/python/xorbits/core/tests/test_execution.py @@ -339,11 +339,40 @@ 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) + + # TODO: process exit cause hang # def test_execution_with_process_exit_message(mocker): # import numpy as np # from xoscar.errors import ServerClosed - # import xorbits # import xorbits.remote as xr