add DataFrame.drop

See http://stackoverflow.com/questions/31867983/cant-drop-columns-or-slice-dataframe-using-dask
mrocklin · Aug 7, 2015 · f5cb23b · f5cb23b
1 parent ac3335e
commit f5cb23b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
@@ -774,6 +774,16 @@ def to_csv(self, filename, **kwargs):
     def _elemwise_cols(self):
         return self.columns
 
+    @wraps(pd.DataFrame.drop)
+    def drop(self, labels, axis=0):
+        if axis != 1:
+            raise NotImplementedError("Drop currently only works for axis=1")
+
+        columns = list(pd.DataFrame(columns=self.columns)
+                         .drop(labels, axis=axis)
+                         .columns)
+        return elemwise(pd.DataFrame.drop, self, labels, axis, columns=columns)
+
 
 def _assign(df, *pairs):
     kwargs = dict(partition(2, pairs))

diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
@@ -858,3 +858,10 @@ def test_gh_517():
 
     ddf2 = dd.from_pandas(pd.concat([df, df]), 5)
     assert ddf2.index.nunique().compute() == 100
+
+
+def test_drop_axis_1():
+    df = pd.DataFrame({'x': [1, 2, 3, 4], 'y': [5, 6, 7, 8]})
+    a = dd.from_pandas(df, npartitions=2)
+
+    assert eq(a.drop('y', axis=1), df.drop('y', axis=1))