# Copyright 2024 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.
from __future__ import annotations
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 import blocks
from bigframes.core import expression as ex
from bigframes.core.indexes.base import Index
if TYPE_CHECKING:
import bigframes.session
[docs]
class MultiIndex(Index, vendored_pandas_multindex.MultiIndex):
__doc__ = vendored_pandas_multindex.MultiIndex.__doc__
[docs]
@classmethod
def from_tuples(
cls,
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, session=session))
[docs]
@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, session=session))
def __eq__(self, other) -> Index: # type: ignore
import bigframes.operations as ops
import bigframes.operations.aggregations as agg_ops
eq_result = self._apply_binary_op(other, ops.eq_op)._block.expr
as_array = ops.ToArrayOp().as_expr(
*(
ops.fillna_op.as_expr(col, ex.const(False))
for col in eq_result.column_ids
)
)
reduced = ops.ArrayReduceOp(agg_ops.all_op).as_expr(as_array)
result_expr, result_ids = eq_result.compute_values([reduced])
return Index(
blocks.Block(
result_expr.select_columns(result_ids),
index_columns=result_ids,
column_labels=(),
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)