plumb through session for from_tuples and from_arrays

Author: tswast

bigframes/core/indexes/multi.py

@@ -14,13 +14,16 @@
 
 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
 
 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__
@@ -31,18 +34,50 @@ 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))
+
+
+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_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)

bigframes/session/__init__.py

@@ -68,6 +68,7 @@
 import bigframes.core
 from bigframes.core import blocks, log_adapter, utils
 import bigframes.core.indexes
+import bigframes.core.indexes.multi
 import bigframes.core.pyformat
 
 # Even though the ibis.backends.bigquery import is unused, it's needed
@@ -2316,20 +2317,17 @@ def DataFrame(self, *args, **kwargs):
 
         return bigframes.dataframe.DataFrame(*args, session=self, **kwargs)
 
-    def MultiIndex(self, *args, **kwargs):
+    @property
+    def MultiIndex(self) -> bigframes.core.indexes.multi.MultiIndexAccessor:
         """Constructs a MultiIndex.
 
         Included for compatibility between bpd and Session.
 
         See :class:`bigframes.pandas.MulitIndex` for full documentation.
         """
-        import bigframes.core.indexes
-
-        return bigframes.core.indexes.MultiIndex(*args, session=self, **kwargs)
+        import bigframes.core.indexes.multi
 
-    MultiIndex.from_tuples = bigframes.core.indexes.MultiIndex.from_tuples  # type: ignore
-    MultiIndex.from_frame = bigframes.core.indexes.MultiIndex.from_frame  # type: ignore
-    MultiIndex.from_arrays = bigframes.core.indexes.MultiIndex.from_arrays  # type: ignore
+        return bigframes.core.indexes.multi.MultiIndexAccessor(self)
 
     def Index(self, *args, **kwargs):
         """Constructs a Index.