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/conftest.py

@@ -0,0 +1,24 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+
+@pytest.fixture(scope="session")
+def polars_session():
+    pytest.importorskip("polars")
+
+    from bigframes.testing import polars_session
+
+    return polars_session.TestSession()

tests/unit/core/test_groupby.py

@@ -23,14 +23,6 @@
 pytest.importorskip("pandas", minversion="2.0.0")
 
 
-# All tests in this file require polars to be installed to pass.
-@pytest.fixture(scope="module")
-def polars_session():
-    from bigframes.testing import polars_session
-
-    return polars_session.TestSession()
-
-
 def test_groupby_df_iter_by_key_singular(polars_session):
     pd_df = pd.DataFrame({"colA": ["a", "a", "b", "c", "c"], "colB": [1, 2, 3, 4, 5]})
     bf_df = bpd.DataFrame(pd_df, session=polars_session)

tests/unit/test_dataframe.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 
 import google.cloud.bigquery
+import pandas as pd
 import pytest
 
 import bigframes.dataframe
+import bigframes.session
 from bigframes.testing import mocks
 
 
@@ -129,6 +131,38 @@ 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(
+    # Drop index depends on the actual data, not just metadata, so use the
+    # local engine for more robust testing.
+    polars_session: bigframes.session.Session,
+):
+    dataframe = polars_session.read_pandas(
+        pd.DataFrame({"col1": [1, 2, 3], "index_col": [0, 1, 2]}).set_index("index_col")
+    )
+    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,
 ):

tests/unit/test_local_engine.py

@@ -24,14 +24,6 @@
 pytest.importorskip("pandas", minversion="2.0.0")
 
 
-# All tests in this file require polars to be installed to pass.
-@pytest.fixture(scope="module")
-def polars_session():
-    from bigframes.testing import polars_session
-
-    return polars_session.TestSession()
-
-
 @pytest.fixture(scope="module")
 def small_inline_frame() -> pd.DataFrame:
     df = pd.DataFrame(