Repository: dask/dask

Issue:
Renaming columns fails on multi-level columns
Taken from https://github.com/dask/dask/pull/1678#issuecomment-255799493 by @chmp

The rename is not related to my code. dask does not seem to handle multilevel columns too well (for example also df['A', 'mean] does not work). A failing test for the rename functionality could be:

``` python
def test_renames_multilevelcolumns():
    pdf = pd.DataFrame({('A', '0') : [1,2,2], ('B', 1) : [1, 2, 3]})
    ddf = dd.from_pandas(pdf, npartitions=3)

    pdf.columns = ['foo', 'bar']
    ddf.columns = ['foo', 'bar']

    assert_eq(pdf, ddf)

# E       AssertionError: Index are different
# E       
# E       Index classes are not equivalent
# E       [left]:  Index(['foo', 'bar'], dtype='object')
# E       [right]: MultiIndex(levels=[['A', 'B'], [1, '0']],
# E                  labels=[[0, 1], [1, 0]])
```

The culprit is the `dask.dataframe.core._rename` implementation:

``` python
def test_renames_pandas():
    pdf = pd.DataFrame({('A', '0') : [1,2,2], ('B', '1') : [1, 2, 3]})

    expected = pdf.copy()
    expected.columns = ['foo', 'bar']

    # df.rename(columns=dict(zip(df.columns, columns))) in core._rename
    actual = pdf.rename(columns=dict(zip(pdf.columns, ['foo', 'bar'])))

    assert_eq(expected, actual)

# result:
# E       AssertionError: DataFrame.columns are different
# E       
# E       DataFrame.columns classes are not equivalent
# E       [left]:  Index(['foo', 'bar'], dtype='object')
# E       [right]: MultiIndex(levels=[['A', 'B'], ['0', '1']],
# E                  labels=[[0, 1], [0, 1]])
```

However, as of now, I have no good idea how to fix this.



Discussion:

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1920,12 +1920,10 @@ def columns(self):
 
     @columns.setter
     def columns(self, columns):
-        # if length mismatches, error is raised from pandas
-        self._meta.columns = columns
         renamed = _rename_dask(self, columns)
-        # update myself
-        self.dask.update(renamed.dask)
+        self._meta = renamed._meta
         self._name = renamed._name
+        self.dask.update(renamed.dask)
 
     def __getitem__(self, key):
         name = 'getitem-%s' % tokenize(self, key)
@@ -2311,13 +2309,20 @@ def meth(self, other, axis='columns', level=None, fill_value=None):
 
             axis = self._validate_axis(axis)
 
-            if axis == 1:
-                # when axis=1, series will be added to each row
-                # it not supported for dd.Series.
-                # dd.DataFrame is not affected as op is applied elemwise
+            if axis in (1, 'columns'):
+                # When axis=1 and other is a series, `other` is transposed
+                # and the operator is applied broadcast across rows. This
+                # isn't supported with dd.Series.
                 if isinstance(other, Series):
                     msg = 'Unable to {0} dd.Series with axis=1'.format(name)
                     raise ValueError(msg)
+                elif isinstance(other, pd.Series):
+                    # Special case for pd.Series to avoid unwanted partitioning
+                    # of other. We pass it in as a kwarg to prevent this.
+                    meta = _emulate(op, self, other=other, axis=axis,
+                                    fill_value=fill_value)
+                    return map_partitions(op, self, other=other, meta=meta,
+                                          axis=axis, fill_value=fill_value)
 
             meta = _emulate(op, self, other, axis=axis, fill_value=fill_value)
             return map_partitions(op, self, other, meta=meta,
@@ -2788,28 +2793,32 @@ def map_partitions(func, *args, **kwargs):
         # If `meta` is not a pandas object, the concatenated results will be a
         # different type
         meta = _concat([meta])
-
-    if isinstance(meta, pd.DataFrame):
-        columns = meta.columns
-    elif isinstance(meta, (pd.Series, pd.Index)):
-        columns = meta.name
-    else:
-        columns = None
+    meta = make_meta(meta)
 
     dfs = [df for df in args if isinstance(df, _Frame)]
     dsk = {}
     for i in range(dfs[0].npartitions):
         values = [(arg._name, i if isinstance(arg, _Frame) else 0)
                   if isinstance(arg, (_Frame, Scalar)) else arg for arg in args]
-        values = (apply, func, (tuple, values), kwargs)
-        if columns is not None:
-            values = (_rename, columns, values)
-        dsk[(name, i)] = values
+        dsk[(name, i)] = (apply_and_enforce, func, values, kwargs, meta)
 
     dasks = [arg.dask for arg in args if isinstance(arg, (_Frame, Scalar))]
     return new_dd_object(merge(dsk, *dasks), name, meta, args[0].divisions)
 
 
+def apply_and_enforce(func, args, kwargs, meta):
+    """Apply a function, and enforce the output to match meta
+
+    Ensures the output has the same columns, even if empty."""
+    df = func(*args, **kwargs)
+    if isinstance(df, (pd.DataFrame, pd.Series, pd.Index)):
+        if len(df) == 0:
+            return meta
+        c = meta.columns if isinstance(df, pd.DataFrame) else meta.name
+        return _rename(c, df)
+    return df
+
+
 def _rename(columns, df):
     """
     Rename columns of pd.DataFrame or name of pd.Series.
@@ -2817,7 +2826,6 @@ def _rename(columns, df):
 
     Parameters
     ----------
-
     columns : tuple, string, pd.DataFrame or pd.Series
         Column names, Series name or pandas instance which has the
         target column names / name.
@@ -2826,12 +2834,12 @@ def _rename(columns, df):
     """
     assert not isinstance(df, _Frame)
 
-    if isinstance(columns, Iterator):
-        columns = list(columns)
-
     if columns is no_default:
         return df
 
+    if isinstance(columns, Iterator):
+        columns = list(columns)
+
     if isinstance(df, pd.DataFrame):
         if isinstance(columns, pd.DataFrame):
             columns = columns.columns
@@ -2840,14 +2848,16 @@ def _rename(columns, df):
             if columns.equals(df.columns):
                 # if target is identical, rename is not necessary
                 return df
-            # each functions must be pure op, do not use df.columns = columns
-            return df.rename(columns=dict(zip(df.columns, columns)))
+        # deep=False doesn't doesn't copy any data/indices, so this is cheap
+        df = df.copy(deep=False)
+        df.columns = columns
+        return df
     elif isinstance(df, (pd.Series, pd.Index)):
         if isinstance(columns, (pd.Series, pd.Index)):
             columns = columns.name
-        if name == columns:
+        if df.name == columns:
             return df
-        return pd.Series(df, name=columns)
+        return df.rename(columns)
     # map_partition may pass other types
     return df
 
@@ -2862,7 +2872,6 @@ def _rename_dask(df, names):
 
     Parameters
     ----------
-
     df : dd.DataFrame or dd.Series
         target DataFrame / Series to be renamed
     names : tuple, string
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -728,8 +728,7 @@ def apply(self, func, meta=no_default, columns=no_default):
 
             with raise_on_meta_error("groupby.apply({0})".format(funcname(func))):
                 meta = self._meta_nonempty.apply(func)
-        else:
-            meta = make_meta(meta)
+        meta = make_meta(meta)
 
         df = self.obj
         if isinstance(self.index, DataFrame):  # add index columns to dataframe


Test patch:
diff --git a/dask/dataframe/tests/test_arithmetics_reduction.py b/dask/dataframe/tests/test_arithmetics_reduction.py
--- a/dask/dataframe/tests/test_arithmetics_reduction.py
+++ b/dask/dataframe/tests/test_arithmetics_reduction.py
@@ -514,6 +514,7 @@ def test_frame_series_arithmetic_methods():
                         index=list('abcdefghjk'), columns=list('ABCX'))
     ps1 = pdf1.A
     ps2 = pdf2.A
+    ps3 = pd.Series(np.random.randn(10), index=list('ABCDXabcde'))
 
     ddf1 = dd.from_pandas(pdf1, 2)
     ddf2 = dd.from_pandas(pdf2, 2)
@@ -577,7 +578,7 @@ def test_frame_series_arithmetic_methods():
 
         pytest.raises(ValueError, lambda: l.add(r, axis=1))
 
-    for l, r, el, er in [(ddf1, pdf2, pdf1, pdf2), (ddf1, ps2, pdf1, ps2)]:
+    for l, r, el, er in [(ddf1, pdf2, pdf1, pdf2), (ddf1, ps3, pdf1, ps3)]:
         assert_eq(l, el)
         assert_eq(r, er)
 
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -308,6 +308,16 @@ def test_rename_columns():
     with tm.assertRaisesRegexp(ValueError, msg):
         ddf.columns = [1, 2, 3, 4]
 
+    # Multi-index columns
+    df = pd.DataFrame({('A', '0') : [1, 2, 2, 3], ('B', 1) : [1, 2, 3, 4]})
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    df.columns = ['x', 'y']
+    ddf.columns = ['x', 'y']
+    tm.assert_index_equal(ddf.columns, pd.Index(['x', 'y']))
+    tm.assert_index_equal(ddf._meta.columns, pd.Index(['x', 'y']))
+    assert_eq(ddf, df)
+
 
 def test_rename_series():
     # GH 819
@@ -319,6 +329,13 @@ def test_rename_series():
     assert s.name == 'renamed'
     assert_eq(ds, s)
 
+    ind = s.index
+    dind = ds.index
+    ind.name = 'renamed'
+    dind.name = 'renamed'
+    assert ind.name == 'renamed'
+    assert_eq(dind, ind)
+
 
 def test_describe():
     # prepare test case which approx quantiles will be the same as actuals
@@ -1536,8 +1553,8 @@ def test_datetime_accessor():
     assert_eq(a.x.dt.date, df.x.dt.date, check_names=False)
     assert (a.x.dt.to_pydatetime().compute() == df.x.dt.to_pydatetime()).all()
 
-    assert a.x.dt.date.dask == a.x.dt.date.dask
-    assert a.x.dt.to_pydatetime().dask == a.x.dt.to_pydatetime().dask
+    assert set(a.x.dt.date.dask) == set(a.x.dt.date.dask)
+    assert set(a.x.dt.to_pydatetime().dask) == set(a.x.dt.to_pydatetime().dask)
 
 
 def test_str_accessor():
@@ -1547,27 +1564,27 @@ def test_str_accessor():
 
     assert 'upper' in dir(a.x.str)
     assert_eq(a.x.str.upper(), df.x.str.upper())
-    assert a.x.str.upper().dask == a.x.str.upper().dask
+    assert set(a.x.str.upper().dask) == set(a.x.str.upper().dask)
 
     assert 'upper' in dir(a.index.str)
     assert_eq(a.index.str.upper(), df.index.str.upper())
-    assert a.index.str.upper().dask == a.index.str.upper().dask
+    assert set(a.index.str.upper().dask) == set(a.index.str.upper().dask)
 
     # make sure to pass thru args & kwargs
     assert 'contains' in dir(a.x.str)
     assert_eq(a.x.str.contains('a'), df.x.str.contains('a'))
-    assert a.x.str.contains('a').dask == a.x.str.contains('a').dask
+    assert set(a.x.str.contains('a').dask) == set(a.x.str.contains('a').dask)
 
     assert_eq(a.x.str.contains('d', case=False), df.x.str.contains('d', case=False))
-    assert a.x.str.contains('d', case=False).dask == a.x.str.contains('d', case=False).dask
+    assert set(a.x.str.contains('d', case=False).dask) == set(a.x.str.contains('d', case=False).dask)
 
     for na in [True, False]:
         assert_eq(a.x.str.contains('a', na=na), df.x.str.contains('a', na=na))
-        assert a.x.str.contains('a', na=na).dask == a.x.str.contains('a', na=na).dask
+        assert set(a.x.str.contains('a', na=na).dask) == set(a.x.str.contains('a', na=na).dask)
 
     for regex in [True, False]:
         assert_eq(a.x.str.contains('a', regex=regex), df.x.str.contains('a', regex=regex))
-        assert a.x.str.contains('a', regex=regex).dask == a.x.str.contains('a', regex=regex).dask
+        assert set(a.x.str.contains('a', regex=regex).dask) == set(a.x.str.contains('a', regex=regex).dask)
 
 
 def test_empty_max():