Add pipe method to dask.dataframe

dask/dataframe/core.py

@@ -531,6 +531,20 @@ def reduction(self, chunk, aggregate=None, meta=no_default,
                    meta=meta, token=token, chunk_kwargs=chunk_kwargs,
                    aggregate_kwargs=aggregate_kwargs, **kwargs)
 
+    @derived_from(pd.DataFrame)
+    def pipe(self, func, *args, **kwargs):
+        # Taken from pandas:
+        # https://github.com/pydata/pandas/blob/master/pandas/core/generic.py#L2698-L2707
+        if isinstance(func, tuple):
+            func, target = func
+            if target in kwargs:
+                raise ValueError('%s is both the pipe target and a keyword '
+                                 'argument' % target)
+            kwargs[target] = self
+            return func(*args, **kwargs)
+        else:
+            return func(self, *args, **kwargs)
+
     def random_split(self, p, random_state=None):
         """ Pseudorandomly split dataframe into different pieces row-wise
 

dask/dataframe/tests/test_dataframe.py

@@ -1541,6 +1541,17 @@ def sum_and_count(x):
     assert eq(res, pd.DataFrame({'sum': df.sum(), 'count': df.count()}))
 
 
+def test_pipe():
+    df = pd.DataFrame({'x': range(50), 'y': range(50, 100)})
+    ddf = dd.from_pandas(df, npartitions=4)
+
+    def f(x, y, z=0):
+        return x + y + z
+
+    assert eq(ddf.pipe(f, 1, z=2), f(ddf, 1, z=2))
+    assert eq(ddf.x.pipe(f, 1, z=2), f(ddf.x, 1, z=2))
+
+
 def test_gh_517():
     arr = np.random.randn(100, 2)
     df = pd.DataFrame(arr, columns=['a', 'b'])