[BACKPORT] Implements {DataFrame,Series}.set_axis (#1950) (#1951)

Author: wjsi

.github/workflows/os-compat-ci.yml

@@ -57,10 +57,10 @@ jobs:
       shell: bash
       run: |
         source ./.github/workflows/reload-env.sh
-        export CYTHON_TRACE=1
         for cf in `ls .coveragerc*`; do
+          sed -i.bak "s/plugins *= *Cython\.Coverage//g" $cf;
           sed -i.bak -e '/*\.pxd/ a\
-          \ \ \ \ *.py \
+          \ \ \ \ *.pyx \
           ' $cf
         done
         retry python setup.py build_ext -i -j 2

docs/source/reference/dataframe/frame.rst

@@ -151,6 +151,7 @@ Reindexing / selection / label manipulation
    DataFrame.rename
    DataFrame.rename_axis
    DataFrame.reset_index
+   DataFrame.set_axis
    DataFrame.set_index
    DataFrame.tail
 

docs/source/reference/dataframe/series.rst

@@ -148,6 +148,7 @@ Reindexing / selection / label manipulation
    Series.rename
    Series.rename_axis
    Series.reset_index
+   Series.set_axis
    Series.tail
 
 Missing data handling

mars/dataframe/base/astype.py

@@ -20,9 +20,9 @@
 from ...serialize import AnyField, StringField, ListField
 from ...utils import recursive_tile
 from ...tensor.base import sort
-from ..utils import build_empty_df, build_empty_series
 from ..core import DATAFRAME_TYPE, SERIES_TYPE
 from ..operands import DataFrameOperand, DataFrameOperandMixin
+from ..utils import build_empty_df, build_empty_series, parse_index
 
 
 class DataFrameAstype(DataFrameOperand, DataFrameOperandMixin):
@@ -197,8 +197,10 @@ def __call__(self, df):
                 return self.new_series([df], shape=df.shape, dtype=dtype,
                                        name=df.name, index_value=df.index_value)
             else:
+                new_index = df.index_value.to_pandas().astype(self.dtype_values)
+                new_index_value = parse_index(new_index, store_data=df.index_value.has_value())
                 return self.new_index([df], shape=df.shape, dtype=dtype,
-                                      name=df.name, index_value=df.index_value)
+                                      name=df.name, index_value=new_index_value)
 
 
 def astype(df, dtype, copy=True, errors='raise'):

mars/dataframe/core.py

@@ -546,7 +546,7 @@ def to_tensor(self, dtype=None, extract_multi_index=False):
 
 
 class Index(HasShapeTileableEnity, _ToPandasMixin):
-    __slots__ = '_df_or_series', '_axis'
+    __slots__ = '_df_or_series', '_parent_key', '_axis'
     _allow_data_type_ = (IndexData,)
 
     def __new__(cls, data: Union[pd.Index, IndexData], **_):
@@ -572,6 +572,7 @@ def _get_df_or_series(self):
 
     def _set_df_or_series(self, df_or_series, axis):
         self._df_or_series = weakref.ref(df_or_series)
+        self._parent_key = df_or_series.key
         self._axis = axis
 
     @property
@@ -581,7 +582,7 @@ def name(self):
     @name.setter
     def name(self, value):
         df_or_series = self._get_df_or_series()
-        if df_or_series is not None:
+        if df_or_series is not None and df_or_series.key == self._parent_key:
             df_or_series.rename_axis(value, axis=self._axis, inplace=True)
             self.data = df_or_series.axes[self._axis].data
         else:
@@ -962,6 +963,10 @@ def index(self):
         idx._set_df_or_series(self, 0)
         return idx
 
+    @index.setter
+    def index(self, new_index):
+        self.set_axis(new_index, axis=0, inplace=True)
+
     @property
     def name(self):
         return self._data.name
@@ -1265,7 +1270,7 @@ def index(self):
     def columns(self):
         from .datasource.index import from_pandas as from_pandas_index
 
-        return from_pandas_index(self.dtypes.index)
+        return from_pandas_index(self.dtypes.index, store_data=True)
 
     @property
     def axes(self):
@@ -1423,6 +1428,10 @@ def index(self):
         idx._set_df_or_series(self, 0)
         return idx
 
+    @index.setter
+    def index(self, new_index):
+        self.set_axis(new_index, axis=0, inplace=True)
+
     @property
     def columns(self):
         col = self._data.columns
@@ -1431,11 +1440,7 @@ def columns(self):
 
     @columns.setter
     def columns(self, new_columns):
-        from .indexing.set_label import DataFrameSetLabel
-
-        op = DataFrameSetLabel(axis=1, value=new_columns)
-        new_df = op(self)
-        self.data = new_df.data
+        self.set_axis(new_columns, axis=1, inplace=True)
 
     def keys(self):
         """

mars/dataframe/datasource/index.py

@@ -20,7 +20,7 @@
 from ... import opcodes as OperandDef
 from ...config import options
 from ...core import OutputType
-from ...serialize import IndexField, DataTypeField, KeyField
+from ...serialize import IndexField, DataTypeField, KeyField, BoolField
 from ...tensor.utils import get_chunk_slices
 from ..operands import DataFrameOperand, DataFrameOperandMixin
 from ..utils import parse_index, decide_series_chunk_size
@@ -36,11 +36,12 @@ class IndexDataSource(DataFrameOperand, DataFrameOperandMixin):
     _input = KeyField('input')
     _data = IndexField('data')
     _dtype = DataTypeField('dtype')
+    _store_data = BoolField('store_data')
 
     def __init__(self, input=None, data=None, dtype=None, gpu=None,  # pylint: disable=redefined-builtin
-                 sparse=None, **kw):
-        super().__init__(_input=input, _data=data, _dtype=dtype, _gpu=gpu,
-                         _sparse=sparse, _output_types=[OutputType.index], **kw)
+                 store_data=None, sparse=None, **kw):
+        super().__init__(_input=input, _data=data, _dtype=dtype, _gpu=gpu, _sparse=sparse,
+                         _store_data=store_data, _output_types=[OutputType.index], **kw)
 
     @property
     def input(self):
@@ -54,6 +55,10 @@ def data(self):
     def dtype(self):
         return self._dtype
 
+    @property
+    def store_data(self):
+        return self._store_data
+
     def _set_inputs(self, inputs):
         super()._set_inputs(inputs)
         if inputs is not None and len(inputs) > 0:
@@ -66,7 +71,7 @@ def __call__(self, shape=None, chunk_size=None, inp=None, name=None,
             name = name if name is not None else self._data.name
             names = names if names is not None else self._data.names
             return self.new_index(None, shape=shape, dtype=self._dtype,
-                                  index_value=parse_index(self._data),
+                                  index_value=parse_index(self._data, store_data=self.store_data),
                                   name=name, names=names, raw_chunk_size=chunk_size)
         elif hasattr(inp, 'index_value'):
             # get index from Mars DataFrame, Series or Index
@@ -75,8 +80,8 @@ def __call__(self, shape=None, chunk_size=None, inp=None, name=None,
             if inp.index_value.has_value():
                 self._data = data = inp.index_value.to_pandas()
                 return self.new_index(None, shape=(inp.shape[0],), dtype=data.dtype,
-                                      index_value=parse_index(data), name=name,
-                                      names=names, raw_chunk_size=chunk_size)
+                                      index_value=parse_index(data, store_data=self.store_data),
+                                      name=name, names=names, raw_chunk_size=chunk_size)
             else:
                 if self._dtype is None:
                     self._dtype = inp.index_value.to_pandas().dtype
@@ -92,7 +97,7 @@ def __call__(self, shape=None, chunk_size=None, inp=None, name=None,
             if self._dtype is None:
                 self._dtype = pd_index.dtype
             return self.new_index([inp], shape=inp.shape, dtype=self._dtype,
-                                  index_value=parse_index(pd_index, inp),
+                                  index_value=parse_index(pd_index, inp, store_data=self.store_data),
                                   name=name, names=names)
 
     @classmethod
@@ -111,9 +116,9 @@ def _tile_from_pandas(cls, op):
             chunk_op = op.copy().reset_key()
             slc = get_chunk_slices(chunk_size, chunk_index)
             chunk_op._data = chunk_data = raw_index[slc]
-            out_chunk = chunk_op.new_chunk(None, shape=chunk_shape, dtype=index.dtype,
-                                           index=chunk_index, name=index.name,
-                                           index_value=parse_index(chunk_data))
+            out_chunk = chunk_op.new_chunk(
+                None, shape=chunk_shape, dtype=index.dtype, index=chunk_index,
+                name=index.name, index_value=parse_index(chunk_data, store_data=op.store_data))
             out_chunks.append(out_chunk)
 
         new_op = op.copy()
@@ -165,7 +170,7 @@ def _tile_from_tensor(cls, op):
         out_chunks = []
         for c in inp.chunks:
             chunk_op = op.copy().reset_key()
-            index_value = parse_index(out.index_value.to_pandas(), c)
+            index_value = parse_index(out.index_value.to_pandas(), c, store_data=op.store_data)
             out_chunk = chunk_op.new_chunk([c], shape=c.shape,
                                            dtype=out.dtype, index=c.index,
                                            index_value=index_value,
@@ -206,8 +211,9 @@ def execute(cls, ctx, op):
                 ctx[out.key] = pd.Index(inp, dtype=dtype, name=out.name)
 
 
-def from_pandas(data, chunk_size=None, gpu=False, sparse=False):
-    op = IndexDataSource(data=data, gpu=gpu, sparse=sparse, dtype=data.dtype)
+def from_pandas(data, chunk_size=None, gpu=False, sparse=False, store_data=False):
+    op = IndexDataSource(data=data, gpu=gpu, sparse=sparse, dtype=data.dtype,
+                         store_data=store_data)
     return op(shape=data.shape, chunk_size=chunk_size)
 
 

mars/dataframe/indexing/__init__.py

@@ -29,6 +29,7 @@ def _install():
     from .setitem import dataframe_setitem
     from .reindex import reindex
     from .where import mask, where
+    from .set_axis import df_set_axis, series_set_axis
 
     for cls in DATAFRAME_TYPE + SERIES_TYPE:
         setattr(cls, 'iloc', cache_readonly(iloc))
@@ -49,20 +50,18 @@ def _install():
         setattr(cls, 'insert', df_insert)
         setattr(cls, 'reset_index', df_reset_index)
         setattr(cls, 'rename', df_rename)
+        setattr(cls, 'set_axis', df_set_axis)
 
     for cls in SERIES_TYPE:
         setattr(cls, '__getitem__', series_getitem)
         setattr(cls, 'reset_index', series_reset_index)
         setattr(cls, 'rename', series_rename)
+        setattr(cls, 'set_axis', series_set_axis)
 
     for cls in INDEX_TYPE:
         setattr(cls, 'rename', index_rename)
         setattr(cls, 'set_names', index_set_names)
 
-    # make sure operand is registered
-    from .set_label import DataFrameSetLabel
-    del DataFrameSetLabel
-
 
 _install()
 del _install