feat: support inplace=True in rename and rename_axis

bigframes/core/compile/polars/compiler.py

@@ -92,6 +92,8 @@ def _(
                 return args[0] < args[1]
             if isinstance(op, ops.eq_op.__class__):
                 return args[0] == args[1]
+            if isinstance(op, ops.ne_op.__class__):
+                return args[0] != args[1]
             if isinstance(op, ops.mod_op.__class__):
                 return args[0] % args[1]
             if isinstance(op, ops.coalesce_op.__class__):
@@ -101,6 +103,9 @@ def _(
                 for pred, result in zip(args[2::2], args[3::2]):
                     return expr.when(pred).then(result)
                 return expr
+            if isinstance(op, ops.where_op.__class__):
+                original, condition, otherwise = args
+                return pl.when(condition).then(original).otherwise(otherwise)
             raise NotImplementedError(f"Polars compiler hasn't implemented {op}")
 
     @dataclasses.dataclass(frozen=True)

bigframes/core/indexes/base.py

@@ -145,12 +145,7 @@ def names(self) -> typing.Sequence[blocks.Label]:
 
     @names.setter
     def names(self, values: typing.Sequence[blocks.Label]):
-        new_block = self._block.with_index_labels(values)
-        if self._linked_frame is not None:
-            self._linked_frame._set_block(
-                self._linked_frame._block.with_index_labels(values)
-            )
-        self._block = new_block
+        self.rename(values, inplace=True)
 
     @property
     def nlevels(self) -> int:
@@ -411,11 +406,52 @@ def fillna(self, value=None) -> Index:
             ops.fillna_op.as_expr(ex.free_var("arg"), ex.const(value))
         )
 
-    def rename(self, name: Union[str, Sequence[str]]) -> Index:
-        names = [name] if isinstance(name, str) else list(name)
+    @overload
+    def rename(
+        self,
+        name: Union[blocks.Label, Sequence[blocks.Label]],
+    ) -> Index:
+        ...
+
+    @overload
+    def rename(
+        self,
+        name: Union[blocks.Label, Sequence[blocks.Label]],
+        *,
+        inplace: Literal[False],
+    ) -> Index:
+        ...
+
+    @overload
+    def rename(
+        self,
+        name: Union[blocks.Label, Sequence[blocks.Label]],
+        *,
+        inplace: Literal[True],
+    ) -> None:
+        ...
+
+    def rename(
+        self,
+        name: Union[blocks.Label, Sequence[blocks.Label]],
+        *,
+        inplace: bool = False,
+    ) -> Optional[Index]:
+        names = [name] if isinstance(name, blocks.Label) else list(name)
         if len(names) != self.nlevels:
             raise ValueError("'name' must be same length as levels")
-        return Index(self._block.with_index_labels(names))
+
+        new_block = self._block.with_index_labels(names)
+
+        if inplace:
+            if self._linked_frame is not None:
+                self._linked_frame._set_block(
+                    self._linked_frame._block.with_index_labels(names)
+                )
+            self._block = new_block
+            return None
+        else:
+            return Index(new_block)
 
     def drop(
         self,

bigframes/dataframe.py

@@ -2081,15 +2081,67 @@ def reorder_levels(self, order: LevelsType, axis: int | str = 0):
     def _resolve_levels(self, level: LevelsType) -> typing.Sequence[str]:
         return self._block.index.resolve_level(level)
 
+    @overload
     def rename(self, *, columns: Mapping[blocks.Label, blocks.Label]) -> DataFrame:
+        ...
+
+    @overload
+    def rename(
+        self, *, columns: Mapping[blocks.Label, blocks.Label], inplace: Literal[False]
+    ) -> DataFrame:
+        ...
+
+    @overload
+    def rename(
+        self, *, columns: Mapping[blocks.Label, blocks.Label], inplace: Literal[True]
+    ) -> None:
+        ...
+
+    def rename(
+        self, *, columns: Mapping[blocks.Label, blocks.Label], inplace: bool = False
+    ) -> Optional[DataFrame]:
         block = self._block.rename(columns=columns)
-        return DataFrame(block)
 
+        if inplace:
+            self._block = block
+            return None
+        else:
+            return DataFrame(block)
+
+    @overload
+    def rename_axis(
+        self,
+        mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+    ) -> DataFrame:
+        ...
+
+    @overload
     def rename_axis(
         self,
         mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+        *,
+        inplace: Literal[False],
         **kwargs,
     ) -> DataFrame:
+        ...
+
+    @overload
+    def rename_axis(
+        self,
+        mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+        *,
+        inplace: Literal[True],
+        **kwargs,
+    ) -> None:
+        ...
+
+    def rename_axis(
+        self,
+        mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+        *,
+        inplace: bool = False,
+        **kwargs,
+    ) -> Optional[DataFrame]:
         if len(kwargs) != 0:
             raise NotImplementedError(
                 f"rename_axis does not currently support any keyword arguments. {constants.FEEDBACK_LINK}"
@@ -2099,7 +2151,14 @@ def rename_axis(
             labels = mapper
         else:
             labels = [mapper]
-        return DataFrame(self._block.with_index_labels(labels))
+
+        block = self._block.with_index_labels(labels)
+
+        if inplace:
+            self._block = block
+            return None
+        else:
+            return DataFrame(block)
 
     @validations.requires_ordering()
     def equals(self, other: typing.Union[bigframes.series.Series, DataFrame]) -> bool:

bigframes/series.py

@@ -31,6 +31,7 @@
     Literal,
     Mapping,
     Optional,
+    overload,
     Sequence,
     Tuple,
     Union,
@@ -94,6 +95,10 @@ class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Ser
     # Must be above 5000 for pandas to delegate to bigframes for binops
     __pandas_priority__ = 13000
 
+    # Ensure mypy can more robustly determine the type of self._block since it
+    # gets set in various places.
+    _block: blocks.Block
+
     def __init__(self, *args, **kwargs):
         self._query_job: Optional[bigquery.QueryJob] = None
         super().__init__(*args, **kwargs)
@@ -253,22 +258,45 @@ def __iter__(self) -> typing.Iterator:
     def copy(self) -> Series:
         return Series(self._block)
 
+    @overload
     def rename(
-        self, index: Union[blocks.Label, Mapping[Any, Any]] = None, **kwargs
+        self,
+        index: Union[blocks.Label, Mapping[Any, Any]] = None,
+    ) -> Series:
+        ...
+
+    @overload
+    def rename(
+        self,
+        index: Union[blocks.Label, Mapping[Any, Any]] = None,
+        *,
+        inplace: Literal[False],
+        **kwargs,
     ) -> Series:
+        ...
+
+    @overload
+    def rename(
+        self,
+        index: Union[blocks.Label, Mapping[Any, Any]] = None,
+        *,
+        inplace: Literal[True],
+        **kwargs,
+    ) -> None:
+        ...
+
+    def rename(
+        self,
+        index: Union[blocks.Label, Mapping[Any, Any]] = None,
+        *,
+        inplace: bool = False,
+        **kwargs,
+    ) -> Optional[Series]:
         if len(kwargs) != 0:
             raise NotImplementedError(
                 f"rename does not currently support any keyword arguments. {constants.FEEDBACK_LINK}"
             )
 
-        # rename the Series name
-        if index is None or isinstance(
-            index, str
-        ):  # Python 3.9 doesn't allow isinstance of Optional
-            index = typing.cast(Optional[str], index)
-            block = self._block.with_column_labels([index])
-            return Series(block)
-
         # rename the index
         if isinstance(index, Mapping):
             index = typing.cast(Mapping[Any, Any], index)
@@ -293,22 +321,61 @@ def rename(
 
                 block = block.set_index(new_idx_ids, index_labels=block.index.names)
 
-            return Series(block)
+            if inplace:
+                self._block = block
+                return None
+            else:
+                return Series(block)
 
         # rename the Series name
         if isinstance(index, typing.Hashable):
+            # Python 3.9 doesn't allow isinstance of Optional
             index = typing.cast(Optional[str], index)
             block = self._block.with_column_labels([index])
-            return Series(block)
+
+            if inplace:
+                self._block = block
+                return None
+            else:
+                return Series(block)
 
         raise ValueError(f"Unsupported type of parameter index: {type(index)}")
 
-    @validations.requires_index
+    @overload
+    def rename_axis(
+        self,
+        mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+    ) -> Series:
+        ...
+
+    @overload
     def rename_axis(
         self,
         mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+        *,
+        inplace: Literal[False],
         **kwargs,
     ) -> Series:
+        ...
+
+    @overload
+    def rename_axis(
+        self,
+        mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+        *,
+        inplace: Literal[True],
+        **kwargs,
+    ) -> None:
+        ...
+
+    @validations.requires_index
+    def rename_axis(
+        self,
+        mapper: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
+        *,
+        inplace: bool = False,
+        **kwargs,
+    ) -> Optional[Series]:
         if len(kwargs) != 0:
             raise NotImplementedError(
                 f"rename_axis does not currently support any keyword arguments. {constants.FEEDBACK_LINK}"
@@ -318,7 +385,13 @@ def rename_axis(
             labels = mapper
         else:
             labels = [mapper]
-        return Series(self._block.with_index_labels(labels))
+
+        block = self._block.with_index_labels(labels)
+        if inplace:
+            self._block = block
+            return None
+        else:
+            return Series(block)
 
     def equals(
         self, other: typing.Union[Series, bigframes.dataframe.DataFrame]

bigframes/testing/mocks.py

@@ -14,7 +14,7 @@
 
 import copy
 import datetime
-from typing import Optional, Sequence
+from typing import Any, Dict, Optional, Sequence
 import unittest.mock as mock
 
 import google.auth.credentials
@@ -23,12 +23,9 @@
 
 import bigframes
 import bigframes.clients
-import bigframes.core.ordering
+import bigframes.core.global_session
 import bigframes.dataframe
-import bigframes.series
 import bigframes.session.clients
-import bigframes.session.executor
-import bigframes.session.metrics
 
 """Utilities for creating test resources."""
 
@@ -129,7 +126,10 @@ def query_and_wait_mock(query, *args, job_config=None, **kwargs):
 
 
 def create_dataframe(
-    monkeypatch: pytest.MonkeyPatch, *, session: Optional[bigframes.Session] = None
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    session: Optional[bigframes.Session] = None,
+    data: Optional[Dict[str, Sequence[Any]]] = None,
 ) -> bigframes.dataframe.DataFrame:
     """[Experimental] Create a mock DataFrame that avoids making Google Cloud API calls.
 
@@ -138,8 +138,11 @@ def create_dataframe(
     if session is None:
         session = create_bigquery_session()
 
+    if data is None:
+        data = {"col": []}
+
     # Since this may create a ReadLocalNode, the session we explicitly pass in
     # might not actually be used. Mock out the global session, too.
     monkeypatch.setattr(bigframes.core.global_session, "_global_session", session)
     bigframes.options.bigquery._session_started = True
-    return bigframes.dataframe.DataFrame({"col": []}, session=session)
+    return bigframes.dataframe.DataFrame(data, session=session)

tests/unit/polars_session.py

@@ -16,6 +16,7 @@
 from typing import Optional, Union
 import weakref
 
+import pandas
 import polars
 
 import bigframes
@@ -87,5 +88,7 @@ def __init__(self):
 
     def read_pandas(self, pandas_dataframe, write_engine="default"):
         # override read_pandas to always keep data local-only
+        if isinstance(pandas_dataframe, pandas.Series):
+            pandas_dataframe = pandas_dataframe.to_frame()
         local_block = bigframes.core.blocks.Block.from_local(pandas_dataframe, self)
         return bigframes.dataframe.DataFrame(local_block)