Support duplicate column cum-functions (#9685)

dask · Nov 30, 2022 · 8ef438b · 8ef438b
1 parent 3ac3b8d
commit 8ef438b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
@@ -7214,7 +7214,8 @@ def _last_valid(s):
             if a.empty:
                 return series_typ([], dtype="float")
             return series_typ(
-                {col: _last_valid(a[col]) for col in a.columns}, index=a.columns
+                [_last_valid(a.iloc[:, i]) for i in range(len(a.columns))],
+                index=a.columns,
             )
         else:
             return _last_valid(a)

diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
@@ -717,6 +717,30 @@ def test_cumulative():
     assert_eq(df.cummax(axis=1, skipna=False), ddf.cummax(axis=1, skipna=False))
     assert_eq(df.cumprod(axis=1, skipna=False), ddf.cumprod(axis=1, skipna=False))
 
+    # With duplicate columns
+    df = pd.DataFrame(np.random.randn(100, 3), columns=list("abb"))
+    ddf = dd.from_pandas(df, 3)
+
+    assert_eq(df.cumsum(), ddf.cumsum())
+    assert_eq(df.cummin(), ddf.cummin())
+    assert_eq(df.cummax(), ddf.cummax())
+    assert_eq(df.cumprod(), ddf.cumprod())
+
+    assert_eq(df.cumsum(skipna=False), ddf.cumsum(skipna=False))
+    assert_eq(df.cummin(skipna=False), ddf.cummin(skipna=False))
+    assert_eq(df.cummax(skipna=False), ddf.cummax(skipna=False))
+    assert_eq(df.cumprod(skipna=False), ddf.cumprod(skipna=False))
+
+    assert_eq(df.cumsum(axis=1), ddf.cumsum(axis=1))
+    assert_eq(df.cummin(axis=1), ddf.cummin(axis=1))
+    assert_eq(df.cummax(axis=1), ddf.cummax(axis=1))
+    assert_eq(df.cumprod(axis=1), ddf.cumprod(axis=1))
+
+    assert_eq(df.cumsum(axis=1, skipna=False), ddf.cumsum(axis=1, skipna=False))
+    assert_eq(df.cummin(axis=1, skipna=False), ddf.cummin(axis=1, skipna=False))
+    assert_eq(df.cummax(axis=1, skipna=False), ddf.cummax(axis=1, skipna=False))
+    assert_eq(df.cumprod(axis=1, skipna=False), ddf.cumprod(axis=1, skipna=False))
+
 
 @pytest.mark.parametrize(
     "func",