Repository: mars-project/mars

Issue:
Support {DataFrame,Series}.add_prefix
Support {DataFrame,Series}.add_prefix


Discussion:

Resolving patch:
diff --git a/mars/dataframe/indexing/__init__.py b/mars/dataframe/indexing/__init__.py
--- a/mars/dataframe/indexing/__init__.py
+++ b/mars/dataframe/indexing/__init__.py
@@ -31,6 +31,7 @@ def _install():
     from .where import mask, where
     from .set_axis import df_set_axis, series_set_axis
     from .sample import sample
+    from .add_preifx import df_add_prefix, series_add_prefix
 
     for cls in DATAFRAME_TYPE + SERIES_TYPE:
         setattr(cls, 'iloc', cache_readonly(iloc))
@@ -54,12 +55,14 @@ def _install():
         setattr(cls, 'reset_index', df_reset_index)
         setattr(cls, 'rename', df_rename)
         setattr(cls, 'set_axis', df_set_axis)
+        setattr(cls, 'add_prefix', df_add_prefix)
 
     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)
+        setattr(cls, 'add_prefix', series_add_prefix)
 
     for cls in INDEX_TYPE:
         setattr(cls, '__getitem__', index_getitem)
diff --git a/mars/dataframe/indexing/add_preifx.py b/mars/dataframe/indexing/add_preifx.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/indexing/add_preifx.py
@@ -0,0 +1,93 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partial
+
+
+def df_add_prefix(df, prefix):
+    """
+    Prefix labels with string `prefix`.
+
+    For DataFrame, the column labels are prefixed.
+
+    Parameters
+    ----------
+    prefix : str
+        The string to add before each label.
+
+    Returns
+    -------
+    DataFrame
+        New DataFrame with updated labels.
+
+    Examples
+    --------
+    >>> import mars.dataframe as md
+    >>> df = md.DataFrame({'A': [1, 2, 3, 4], 'B': [3, 4, 5, 6]})
+    >>> df.execute()
+        A  B
+    0  1  3
+    1  2  4
+    2  3  5
+    3  4  6
+
+    >>> df.add_prefix('col_').execute()
+            col_A  col_B
+    0       1       3
+    1       2       4
+    2       3       5
+    3       4       6
+    """
+    f = partial("{prefix}{}".format, prefix=prefix)
+
+    return df.rename(columns=f)
+
+
+def series_add_prefix(series, prefix):
+    """
+    Prefix labels with string `prefix`.
+
+    For Series, the row labels are prefixed.
+
+    Parameters
+    ----------
+    prefix : str
+        The string to add before each label.
+
+    Returns
+    -------
+    Series
+        New Series with updated labels.
+
+    Examples
+    --------
+    >>> import mars.dataframe as md
+    >>> s = md.Series([1, 2, 3, 4])
+    >>> s.execute()
+    0    1
+    1    2
+    2    3
+    3    4
+    dtype: int64
+
+    >>> s.add_prefix('item_').execute()
+    item_0    1
+    item_1    2
+    item_2    3
+    item_3    4
+    dtype: int64
+    """
+    f = partial("{prefix}{}".format, prefix=prefix)
+
+    return series.rename(index=f)


Test patch:
diff --git a/mars/dataframe/indexing/tests/test_indexing_execution.py b/mars/dataframe/indexing/tests/test_indexing_execution.py
--- a/mars/dataframe/indexing/tests/test_indexing_execution.py
+++ b/mars/dataframe/indexing/tests/test_indexing_execution.py
@@ -1501,3 +1501,19 @@ def testSampleExecution(self):
         r2 = s[:].sample(frac=0.1, weights=weights, random_state=rs)
         pd.testing.assert_series_equal(self.executor.execute_dataframe(r1, concat=True)[0],
                                        self.executor.execute_dataframe(r2, concat=True)[0])
+
+    def testAddPrefix(self):
+        rs = np.random.RandomState(0)
+        raw = pd.DataFrame(rs.rand(10, 4), columns=['A', 'B', 'C', 'D'])
+        df = md.DataFrame(raw, chunk_size=3)
+
+        r = df.add_prefix('col_')
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                      raw.add_prefix('col_'))
+
+        raw = pd.Series(rs.rand(10), name='series')
+        series = md.Series(raw, chunk_size=3)
+
+        r = series.add_prefix('item_')
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw.add_prefix('item_'))