feat: implement inplace parameter for DataFrame.drop (#2105)

* feat: implement inplace parameter for drop method This commit implements the `inplace` parameter for the `DataFrame.drop` method. When `inplace=True`, the DataFrame is modified in place and the method returns `None`. When `inplace=False` (the default), a new DataFrame is returned. Unit tests have been added to verify the functionality for both column and row dropping. * update drop index test --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: tswast

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(