feat: implement inplace parameter for DataFrame.drop

bigframes/dataframe.py

@@ -2006,6 +2006,7 @@ def insert(
 
         self._set_block(block)
 
+    @overload
     def drop(
         self,
         labels: typing.Any = None,
@@ -2014,7 +2015,33 @@ def drop(
         index: typing.Any = None,
         columns: Union[blocks.Label, Sequence[blocks.Label]] = None,
         level: typing.Optional[LevelType] = None,
+        inplace: Literal[False] = False,
     ) -> DataFrame:
+        ...
+
+    @overload
+    def drop(
+        self,
+        labels: typing.Any = None,
+        *,
+        axis: typing.Union[int, str] = 0,
+        index: typing.Any = None,
+        columns: Union[blocks.Label, Sequence[blocks.Label]] = None,
+        level: typing.Optional[LevelType] = None,
+        inplace: Literal[True],
+    ) -> None:
+        ...
+
+    def drop(
+        self,
+        labels: typing.Any = None,
+        *,
+        axis: typing.Union[int, str] = 0,
+        index: typing.Any = None,
+        columns: Union[blocks.Label, Sequence[blocks.Label]] = None,
+        level: typing.Optional[LevelType] = None,
+        inplace: bool = False,
+    ) -> Optional[DataFrame]:
         if labels:
             if index or columns:
                 raise ValueError("Cannot specify both 'labels' and 'index'/'columns")
@@ -2056,7 +2083,11 @@ def drop(
                         inverse_condition_id, ops.invert_op
                     )
             elif isinstance(index, indexes.Index):
-                return self._drop_by_index(index)
+                dropped_block = self._drop_by_index(index)._get_block()
+                if inplace:
+                    self._set_block(dropped_block)
+                    return None
+                return DataFrame(dropped_block)
             else:
                 block, condition_id = block.project_expr(
                     ops.ne_op.as_expr(level_id, ex.const(index))
@@ -2068,7 +2099,12 @@ def drop(
             block = block.drop_columns(self._sql_names(columns))
         if index is None and not columns:
             raise ValueError("Must specify 'labels' or 'index'/'columns")
-        return DataFrame(block)
+
+        if inplace:
+            self._set_block(block)
+            return None
+        else:
+            return DataFrame(block)
 
     def _drop_by_index(self, index: indexes.Index) -> DataFrame:
         block = index._block

tests/unit/test_dataframe.py

@@ -129,6 +129,35 @@ def test_dataframe_rename_axis_inplace_returns_none(monkeypatch: pytest.MonkeyPa
     assert list(dataframe.index.names) == ["a", "b"]
 
 
+def test_dataframe_drop_columns_inplace_returns_none(monkeypatch: pytest.MonkeyPatch):
+    dataframe = mocks.create_dataframe(
+        monkeypatch, data={"col1": [1], "col2": [2], "col3": [3]}
+    )
+    assert dataframe.columns.to_list() == ["col1", "col2", "col3"]
+    assert dataframe.drop(columns=["col1", "col3"], inplace=True) is None
+    assert dataframe.columns.to_list() == ["col2"]
+
+
+def test_dataframe_drop_index_inplace_returns_none(monkeypatch: pytest.MonkeyPatch):
+    dataframe = mocks.create_dataframe(
+        monkeypatch, data={"col1": [1, 2, 3], "index_col": [0, 1, 2]}
+    ).set_index("index_col")
+    # TODO(swast): Support dataframe.index.to_list()
+    # assert dataframe.index.to_list() == [0, 1, 2]
+    assert dataframe.drop(index=[0, 2], inplace=True) is None
+    # assert dataframe.index.to_list() == [1]
+
+
+def test_dataframe_drop_columns_returns_new_dataframe(monkeypatch: pytest.MonkeyPatch):
+    dataframe = mocks.create_dataframe(
+        monkeypatch, data={"col1": [1], "col2": [2], "col3": [3]}
+    )
+    assert dataframe.columns.to_list() == ["col1", "col2", "col3"]
+    new_dataframe = dataframe.drop(columns=["col1", "col3"])
+    assert dataframe.columns.to_list() == ["col1", "col2", "col3"]
+    assert new_dataframe.columns.to_list() == ["col2"]
+
+
 def test_dataframe_semantics_property_future_warning(
     monkeypatch: pytest.MonkeyPatch,
 ):