diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index ea5136f6f5..eb5ed997a1 100644 --- a/bigframes/dataframe.py +++ b/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 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000000..a9b26afeef --- /dev/null +++ b/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() diff --git a/tests/unit/core/test_groupby.py b/tests/unit/core/test_groupby.py index 8df0e5344e..f3d9218123 100644 --- a/tests/unit/core/test_groupby.py +++ b/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) diff --git a/tests/unit/test_dataframe.py b/tests/unit/test_dataframe.py index d630380e7a..6aaccd644e 100644 --- a/tests/unit/test_dataframe.py +++ b/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, ): diff --git a/tests/unit/test_local_engine.py b/tests/unit/test_local_engine.py index 509bc6ade2..7d3d532d88 100644 --- a/tests/unit/test_local_engine.py +++ b/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(