feat: create session-scoped cut, DataFrame, MultiIndex, Index, Series, to_datetime, and to_timedelta methods (#2157)

tswast · web-flow · commit 5e1e8098ecf2 · 2025-10-09T13:52:28.000-05:00
* docs: remove import bigframes.pandas as bpd boilerplate from many samples Also, fixes several constructors that didn't take a session for compatibility with multi-session applications. * fix docs * fix unit tests * skip sklearn test * fix snapshot * plumb through session for from_tuples and from_arrays * add from_frame * make sure polars session isnt skipped on Kokoro * fix apply doctest * make doctest conftest available everywhere * add python version flexibility for to_dict * disambiguate explicit names * disambiguate explicit name none versus no name * fix for column name comparison in pandas bin op * avoid setting column labels in special case of Series(block) * revert doctest changes * revert doctest changes * revert df docstrings * add polars series unit tests * restore a test * Revert "restore a test" This reverts commit 765b678. * skip null * skip unsupported tests * revert more docs changes * revert more docs * revert more docs * fix unit tests python 3.13 * add test to reproduce name error * add tests for session scoped methods * fix mypy errors
diff --git a/bigframes/core/indexes/base.py b/bigframes/core/indexes/base.py
@@ -383,9 +383,16 @@ def to_series(
 
         name = self.name if name is None else name
         if index is None:
-            return bigframes.series.Series(data=self, index=self, name=name)
+            return bigframes.series.Series(
+                data=self, index=self, name=name, session=self._session
+            )
         else:
-            return bigframes.series.Series(data=self, index=Index(index), name=name)
+            return bigframes.series.Series(
+                data=self,
+                index=Index(index, session=self._session),
+                name=name,
+                session=self._session,
+            )
 
     def get_level_values(self, level) -> Index:
         level_n = level if isinstance(level, int) else self.names.index(level)
diff --git a/bigframes/core/indexes/multi.py b/bigframes/core/indexes/multi.py
@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import cast, Hashable, Iterable, Sequence
+from typing import cast, Hashable, Iterable, Optional, Sequence, TYPE_CHECKING
 
 import bigframes_vendored.pandas.core.indexes.multi as vendored_pandas_multindex
 import pandas
@@ -23,6 +23,9 @@
 from bigframes.core import expression as ex
 from bigframes.core.indexes.base import Index
 
+if TYPE_CHECKING:
+    import bigframes.session
+
 
 class MultiIndex(Index, vendored_pandas_multindex.MultiIndex):
     __doc__ = vendored_pandas_multindex.MultiIndex.__doc__
@@ -33,21 +36,25 @@ def from_tuples(
         tuples: Iterable[tuple[Hashable, ...]],
         sortorder: int | None = None,
         names: Sequence[Hashable] | Hashable | None = None,
+        *,
+        session: Optional[bigframes.session.Session] = None,
     ) -> MultiIndex:
         pd_index = pandas.MultiIndex.from_tuples(tuples, sortorder, names)
         # Index.__new__ should detect multiple levels and properly create a multiindex
-        return cast(MultiIndex, Index(pd_index))
+        return cast(MultiIndex, Index(pd_index, session=session))
 
     @classmethod
     def from_arrays(
         cls,
         arrays,
         sortorder: int | None = None,
         names=None,
+        *,
+        session: Optional[bigframes.session.Session] = None,
     ) -> MultiIndex:
         pd_index = pandas.MultiIndex.from_arrays(arrays, sortorder, names)
         # Index.__new__ should detect multiple levels and properly create a multiindex
-        return cast(MultiIndex, Index(pd_index))
+        return cast(MultiIndex, Index(pd_index, session=session))
 
     def __eq__(self, other) -> Index:  # type: ignore
         import bigframes.operations as ops
@@ -71,3 +78,38 @@ def __eq__(self, other) -> Index:  # type: ignore
                 index_labels=[None],
             )
         )
+
+
+class MultiIndexAccessor:
+    """Proxy to MultiIndex constructors to allow a session to be passed in."""
+
+    def __init__(self, session: bigframes.session.Session):
+        self._session = session
+
+    def __call__(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :class:`bigframes.pandas.MultiIndex`.
+        """
+        return MultiIndex(*args, session=self._session, **kwargs)
+
+    def from_arrays(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :func:`bigframes.pandas.MultiIndex.from_arrays`.
+        """
+        return MultiIndex.from_arrays(*args, session=self._session, **kwargs)
+
+    def from_frame(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :func:`bigframes.pandas.MultiIndex.from_frame`.
+        """
+        return cast(MultiIndex, MultiIndex.from_frame(*args, **kwargs))
+
+    def from_tuples(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :func:`bigframes.pandas.MultiIndex.from_tuples`.
+        """
+        return MultiIndex.from_tuples(*args, session=self._session, **kwargs)
diff --git a/bigframes/core/log_adapter.py b/bigframes/core/log_adapter.py
@@ -155,7 +155,9 @@ def method_logger(method=None, /, *, custom_base_name: Optional[str] = None):
     def outer_wrapper(method):
         @functools.wraps(method)
         def wrapper(*args, **kwargs):
-            api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__)
+            api_method_name = getattr(
+                method, LOG_OVERRIDE_NAME, method.__name__
+            ).lower()
             if custom_base_name is None:
                 qualname_parts = getattr(method, "__qualname__", method.__name__).split(
                     "."
diff --git a/bigframes/core/reshape/tile.py b/bigframes/core/reshape/tile.py
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import typing
+from typing import Optional, TYPE_CHECKING
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.reshape.tile as vendored_pandas_tile
@@ -31,6 +32,9 @@
 import bigframes.operations.aggregations as agg_ops
 import bigframes.series
 
+if TYPE_CHECKING:
+    import bigframes.session
+
 
 def cut(
     x,
@@ -42,6 +46,7 @@ def cut(
     *,
     right: typing.Optional[bool] = True,
     labels: typing.Union[typing.Iterable[str], bool, None] = None,
+    session: Optional[bigframes.session.Session] = None,
 ) -> bigframes.series.Series:
     if (
         labels is not None
@@ -65,7 +70,7 @@ def cut(
         raise ValueError("Cannot cut empty array.")
 
     if not isinstance(x, bigframes.series.Series):
-        x = bigframes.series.Series(x)
+        x = bigframes.series.Series(x, session=session)
 
     if isinstance(bins, int):
         if bins <= 0:
diff --git a/bigframes/core/tools/datetimes.py b/bigframes/core/tools/datetimes.py
@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from collections.abc import Mapping
 from datetime import date, datetime
-from typing import Optional, Union
+from typing import Optional, TYPE_CHECKING, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.tools.datetimes as vendored_pandas_datetimes
@@ -25,6 +27,9 @@
 import bigframes.operations as ops
 import bigframes.series
 
+if TYPE_CHECKING:
+    import bigframes.session
+
 
 def to_datetime(
     arg: Union[
@@ -37,6 +42,7 @@ def to_datetime(
     utc: bool = False,
     format: Optional[str] = None,
     unit: Optional[str] = None,
+    session: Optional[bigframes.session.Session] = None,
 ) -> Union[pd.Timestamp, datetime, bigframes.series.Series]:
     if isinstance(arg, (int, float, str, datetime, date)):
         return pd.to_datetime(
@@ -52,7 +58,7 @@ def to_datetime(
             f"to datetime is not implemented. {constants.FEEDBACK_LINK}"
         )
 
-    arg = bigframes.series.Series(arg)
+    arg = bigframes.series.Series(arg, session=session)
 
     if format and unit and arg.dtype in (bigframes.dtypes.INT_DTYPE, bigframes.dtypes.FLOAT_DTYPE):  # type: ignore
         raise ValueError("cannot specify both format and unit")
diff --git a/bigframes/formatting_helpers.py b/bigframes/formatting_helpers.py
@@ -105,8 +105,14 @@ def progress_callback(
     """Displays a progress bar while the query is running"""
     global current_display, current_display_id, previous_display_html
 
-    import bigframes._config
-    import bigframes.core.events
+    try:
+        import bigframes._config
+        import bigframes.core.events
+    except ImportError:
+        # Since this gets called from __del__, skip if the import fails to avoid
+        # ImportError: sys.meta_path is None, Python is likely shutting down.
+        # This will allow cleanup to continue.
+        return
 
     progress_bar = bigframes._config.options.display.progress_bar
 
diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py
@@ -16,8 +16,8 @@
 
 from __future__ import annotations
 
-from collections import namedtuple
-from datetime import date, datetime
+import collections
+import datetime
 import inspect
 import sys
 import typing
@@ -198,18 +198,18 @@ def to_datetime(
 
 @typing.overload
 def to_datetime(
-    arg: Union[int, float, str, datetime, date],
+    arg: Union[int, float, str, datetime.datetime, datetime.date],
     *,
     utc: bool = False,
     format: Optional[str] = None,
     unit: Optional[str] = None,
-) -> Union[pandas.Timestamp, datetime]:
+) -> Union[pandas.Timestamp, datetime.datetime]:
     ...
 
 
 def to_datetime(
     arg: Union[
-        Union[int, float, str, datetime, date],
+        Union[int, float, str, datetime.datetime, datetime.date],
         vendored_pandas_datetimes.local_iterables,
         bigframes.series.Series,
         bigframes.dataframe.DataFrame,
@@ -218,8 +218,9 @@ def to_datetime(
     utc: bool = False,
     format: Optional[str] = None,
     unit: Optional[str] = None,
-) -> Union[pandas.Timestamp, datetime, bigframes.series.Series]:
-    return bigframes.core.tools.to_datetime(
+) -> Union[pandas.Timestamp, datetime.datetime, bigframes.series.Series]:
+    return global_session.with_default_session(
+        bigframes.session.Session.to_datetime,
         arg,
         utc=utc,
         format=format,
@@ -322,7 +323,7 @@ def clean_up_by_session_id(
 __version__ = bigframes.version.__version__
 
 # Other public pandas attributes
-NamedAgg = namedtuple("NamedAgg", ["column", "aggfunc"])
+NamedAgg = collections.namedtuple("NamedAgg", ["column", "aggfunc"])
 
 options = config.options
 """Global :class:`~bigframes._config.Options` to configure BigQuery DataFrames."""
diff --git a/bigframes/pandas/core/tools/timedeltas.py b/bigframes/pandas/core/tools/timedeltas.py
@@ -35,7 +35,7 @@ def to_timedelta(
         return arg._apply_unary_op(ops.ToTimedeltaOp(canonical_unit))
 
     if pdtypes.is_list_like(arg):
-        return to_timedelta(series.Series(arg), unit, session=session)
+        return to_timedelta(series.Series(arg, session=session), unit, session=session)
 
     return pd.to_timedelta(arg, unit)
 
diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py
diff --git a/tests/system/small/test_session_as_bpd.py b/tests/system/small/test_session_as_bpd.py
diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py
