Fix incorrect MutliIndex.names (#3355)

Author: ericpai

mars/dataframe/core.py

@@ -378,6 +378,10 @@ def min_max(self):
     def name(self):
         return getattr(self._index_value, "_name", None)
 
+    @property
+    def names(self):
+        return getattr(self._index_value, "_names", [self.name])
+
     @property
     def inferred_type(self):
         return self._index_value.inferred_type

mars/dataframe/datasource/index.py

@@ -68,7 +68,7 @@ def __call__(self, shape=None, chunk_size=None, inp=None, name=None, names=None)
         elif hasattr(inp, "index_value"):
             # get index from Mars DataFrame, Series or Index
             name = name if name is not None else inp.index_value.name
-            names = names if names is not None else [name]
+            names = names if names is not None else inp.index_value.names
             if inp.index_value.has_value():
                 self.data = data = inp.index_value.to_pandas()
                 return self.new_index(

mars/dataframe/datasource/tests/test_datasource.py

@@ -27,7 +27,7 @@
 from ....core import tile
 from ....tests.core import require_ray
 from ....utils import lazy_import
-from ...core import IndexValue, DatetimeIndex, Int64Index, Float64Index
+from ...core import IndexValue, DatetimeIndex, Int64Index, Float64Index, MultiIndex
 from ..core import merge_small_files
 from ..dataframe import from_pandas as from_pandas_df
 from ..date_range import date_range
@@ -156,6 +156,16 @@ def test_from_pandas_dataframe():
         assert len([ns for ns in df.nsplits[1] if ns == 0]) == 0
 
 
+def test_from_pandas_dataframe_with_multi_index():
+    index = pd.MultiIndex.from_tuples([("k1", "v1")], names=["X", "Y"])
+    data = np.random.randint(0, 100, size=(1, 3))
+    pdf = pd.DataFrame(data, columns=["A", "B", "C"], index=index)
+    df = from_pandas_df(pdf, chunk_size=4)
+    assert isinstance(df.index, MultiIndex)
+    assert df.index.names == ["X", "Y"]
+    assert df.index.name is None
+
+
 def test_from_pandas_series():
     data = pd.Series(np.random.rand(10), name="a")
     series = from_pandas_series(data, chunk_size=4)