dataframe.py

# Copyright 2023 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.

"""DataFrame is a two dimensional data structure."""

from __future__ import annotations

import datetime
import inspect
import itertools
import re
import sys
import textwrap
import typing
from typing import (
    Any,
    Callable,
    cast,
    Dict,
    Hashable,
    Iterable,
    List,
    Literal,
    Mapping,
    Optional,
    overload,
    Sequence,
    Tuple,
    Union,
)
import warnings

import bigframes_vendored.constants as constants
import bigframes_vendored.pandas.core.frame as vendored_pandas_frame
import bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing
import google.api_core.exceptions
import google.cloud.bigquery as bigquery
import numpy
import pandas
from pandas.api import extensions as pd_ext
import pandas.io.formats.format
import pyarrow
import tabulate

import bigframes.constants
import bigframes.core
from bigframes.core import agg_expressions
import bigframes.core.block_transforms as block_ops
import bigframes.core.blocks as blocks
import bigframes.core.convert
import bigframes.core.explode
import bigframes.core.expression as ex
import bigframes.core.groupby as groupby
import bigframes.core.guid
import bigframes.core.indexers as indexers
import bigframes.core.indexes as indexes
import bigframes.core.interchange
from bigframes.core.logging import log_adapter
import bigframes.core.ordering as order
import bigframes.core.utils as utils
import bigframes.core.validations as validations
import bigframes.core.window
from bigframes.core.window import rolling
import bigframes.core.window_spec as windows
import bigframes.dtypes
import bigframes.exceptions as bfe
import bigframes.formatting_helpers as formatter
import bigframes.functions
from bigframes.functions import function_typing
import bigframes.operations as ops
import bigframes.operations.aggregations as agg_ops
import bigframes.operations.ai
import bigframes.operations.plotting as plotting
import bigframes.operations.semantics
import bigframes.operations.structs
import bigframes.series
import bigframes.session._io.bigquery
import bigframes.session.execution_spec as ex_spec

if typing.TYPE_CHECKING:
    from _typeshed import SupportsRichComparison

    import bigframes.session

    SingleItemValue = Union[
        bigframes.series.Series, int, float, str, pandas.Timedelta, Callable
    ]
    MultiItemValue = Union[
        "DataFrame", Sequence[int | float | str | pandas.Timedelta | Callable]
    ]

LevelType = typing.Hashable
LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]

ERROR_IO_ONLY_GS_PATHS = f"Only Google Cloud Storage (gs://...) paths are supported. {constants.FEEDBACK_LINK}"
ERROR_IO_REQUIRES_WILDCARD = (
    "Google Cloud Storage path must contain a wildcard '*' character. See: "
    "https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#export_data_statement"
    f"{constants.FEEDBACK_LINK}"
)


# Inherits from pandas DataFrame so that we can use the same docstrings.
@log_adapter.class_logger
class DataFrame(vendored_pandas_frame.DataFrame):
    __doc__ = vendored_pandas_frame.DataFrame.__doc__
    # internal flag to disable cache at all
    _disable_cache_override: bool = False
    # Must be above 5000 for pandas to delegate to bigframes for binops
    __pandas_priority__ = 15000

    def __init__(
        self,
        data=None,
        index: vendored_pandas_typing.Axes | None = None,
        columns: vendored_pandas_typing.Axes | None = None,
        dtype: typing.Optional[
            bigframes.dtypes.DtypeString | bigframes.dtypes.Dtype
        ] = None,
        copy: typing.Optional[bool] = None,
        *,
        session: typing.Optional[bigframes.session.Session] = None,
    ):
        global bigframes

        self._query_job: Optional[bigquery.QueryJob] = None

        if copy is not None and not copy:
            raise ValueError(
                f"DataFrame constructor only supports copy=True. {constants.FEEDBACK_LINK}"
            )
        # Ignore object dtype if provided, as it provides no additional
        # information about what BigQuery type to use.
        if dtype is not None and bigframes.dtypes.is_object_like(dtype):
            dtype = None

        # Check to see if constructing from BigQuery-backed objects before
        # falling back to pandas constructor
        block = None
        if isinstance(data, blocks.Block):
            block = data

        elif isinstance(data, DataFrame):
            block = data._get_block()

        # Dict of Series
        elif (
            utils.is_dict_like(data)
            and len(data) >= 1
            and any(
                isinstance(data[key], bigframes.series.Series) for key in data.keys()
            )
        ):
            if not all(
                isinstance(data[key], bigframes.series.Series) for key in data.keys()
            ):
                # TODO(tbergeron): Support local list/series data by converting to memtable.
                raise NotImplementedError(
                    f"Cannot mix Series with other types. {constants.FEEDBACK_LINK}"
                )
            keys = list(data.keys())
            first_label, first_series = keys[0], data[keys[0]]
            block = (
                typing.cast(bigframes.series.Series, first_series)
                ._get_block()
                .with_column_labels([first_label])
            )

            for key in keys[1:]:
                other = typing.cast(bigframes.series.Series, data[key])
                other_block = other._block.with_column_labels([key])
                # Pandas will keep original sorting if all indices are aligned.
                # We cannot detect this easily however, and so always sort on index
                block, _ = block.join(  # type:ignore
                    other_block, how="outer", sort=True
                )

        if block:
            if index is not None:
                bf_index = indexes.Index(index)
                idx_block = bf_index._block
                idx_cols = idx_block.index_columns
                block, (_, r_mapping) = block.reset_index().join(
                    bf_index._block.reset_index(), how="inner"
                )
                block = block.set_index([r_mapping[idx_col] for idx_col in idx_cols])
            if columns:
                column_ids = [
                    block.resolve_label_exact_or_error(label) for label in list(columns)
                ]
                block = block.select_columns(column_ids)  # type:ignore
            if dtype:
                bf_dtype = bigframes.dtypes.bigframes_type(dtype)
                block = block.multi_apply_unary_op(ops.AsTypeOp(to_type=bf_dtype))

        else:
            if isinstance(dtype, str) and dtype.lower() == "json":
                dtype = bigframes.dtypes.JSON_DTYPE

            import bigframes.pandas

            pd_dataframe = pandas.DataFrame(
                data=data,
                index=index,  # type:ignore
                columns=columns,  # type:ignore
                dtype=dtype,  # type:ignore
            )
            if session:
                block = session.read_pandas(pd_dataframe)._get_block()
            else:
                block = bigframes.pandas.read_pandas(pd_dataframe)._get_block()

        # We use _block as an indicator in __getattr__ and __setattr__ to see
        # if the object is fully initialized, so make sure we set the _block
        # attribute last.
        self._block = block
        self._block.session._register_object(self)

    def __dir__(self):
        return dir(type(self)) + [
            label
            for label in self._block.column_labels
            if label and isinstance(label, str)
        ]

    def _ipython_key_completions_(self) -> List[str]:
        return list(
            [
                label
                for label in self._block.column_labels
                if label and isinstance(label, str)
            ]
        )

    def _find_indices(
        self,
        columns: Union[blocks.Label, Sequence[blocks.Label]],
        tolerance: bool = False,
    ) -> Sequence[int]:
        """Find corresponding indices in df._block.column_labels for column name(s).
        Order is kept the same as input names order.

        Args:
            columns: column name(s)
            tolerance: True to pass through columns not found. False to raise
                ValueError.
        """
        col_ids = self._sql_names(columns, tolerance)
        return [self._block.value_columns.index(col_id) for col_id in col_ids]

    def _resolve_label_exact(self, label) -> Optional[str]:
        return self._block.resolve_label_exact(label)

    def _sql_names(
        self,
        columns: Union[blocks.Label, Sequence[blocks.Label], pandas.Index],
        tolerance: bool = False,
    ) -> Sequence[str]:
        """Retrieve sql name (column name in BQ schema) of column(s)."""
        labels = (
            columns
            if utils.is_list_like(columns) and not isinstance(columns, tuple)
            else [columns]
        )  # type:ignore
        results: Sequence[str] = []
        for label in labels:
            col_ids = self._block.label_to_col_id.get(label, [])
            if not tolerance and len(col_ids) == 0:
                raise ValueError(f"Column name {label} doesn't exist")
            results = (*results, *col_ids)
        return results

    @property
    @validations.requires_index
    def index(
        self,
    ) -> indexes.Index:
        return indexes.Index.from_frame(self)

    @index.setter
    def index(self, value):
        # TODO: Handle assigning MultiIndex
        result = self._assign_single_item("_new_bf_index", value).set_index(
            "_new_bf_index"
        )
        self._set_block(result._get_block())
        self.index.name = value.name if hasattr(value, "name") else None

    @property
    @validations.requires_index
    def loc(self) -> indexers.LocDataFrameIndexer:
        return indexers.LocDataFrameIndexer(self)

    @property
    @validations.requires_ordering()
    def iloc(self) -> indexers.ILocDataFrameIndexer:
        return indexers.ILocDataFrameIndexer(self)

    @property
    @validations.requires_ordering()
    def iat(self) -> indexers.IatDataFrameIndexer:
        return indexers.IatDataFrameIndexer(self)

    @property
    @validations.requires_index
    def at(self) -> indexers.AtDataFrameIndexer:
        return indexers.AtDataFrameIndexer(self)

    @property
    def dtypes(self) -> pandas.Series:
        dtypes = self._block.dtypes
        bigframes.dtypes.warn_on_db_dtypes_json_dtype(dtypes)
        return pandas.Series(data=dtypes, index=self._block.column_labels)

    @property
    def columns(self) -> pandas.Index:
        return self.dtypes.index

    @columns.setter
    def columns(self, labels: pandas.Index):
        new_block = self._block.with_column_labels(labels)
        self._set_block(new_block)

    @property
    def shape(self) -> Tuple[int, int]:
        return self._block.shape

    @property
    def size(self) -> int:
        rows, cols = self.shape
        return rows * cols

    @property
    def ndim(self) -> int:
        return 2

    @property
    def empty(self) -> bool:
        return self.size == 0

    @property
    def values(self) -> numpy.ndarray:
        return self.to_numpy()

    @property
    def bqclient(self) -> bigframes.Session:
        """BigQuery REST API Client the DataFrame uses for operations."""
        return self._session.bqclient

    @property
    def _session(self) -> bigframes.Session:
        return self._get_block().expr.session

    @property
    def _has_index(self) -> bool:
        return len(self._block.index_columns) > 0

    @property
    @validations.requires_ordering()
    def T(self) -> DataFrame:
        return DataFrame(self._get_block().transpose())

    @validations.requires_index
    @validations.requires_ordering()
    def transpose(self) -> DataFrame:
        return self.T

    def __len__(self):
        rows, _ = self.shape
        return rows

    __len__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__len__)

    def __iter__(self):
        return iter(self.columns)

    def __contains__(self, key) -> bool:
        return key in self.columns

    def astype(
        self,
        dtype: Union[
            bigframes.dtypes.DtypeString,
            bigframes.dtypes.Dtype,
            type,
            dict[str, Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype]],
        ],
        *,
        errors: Literal["raise", "null"] = "raise",
    ) -> DataFrame:
        if errors not in ["raise", "null"]:
            raise ValueError("Arg 'error' must be one of 'raise' or 'null'")

        safe_cast = errors == "null"

        if isinstance(dtype, dict):
            result = self.copy()
            for col, to_type in dtype.items():
                result[col] = result[col].astype(to_type)
            return result

        dtype = bigframes.dtypes.bigframes_type(dtype)

        return self._apply_unary_op(ops.AsTypeOp(dtype, safe_cast))

    def _should_sql_have_index(self) -> bool:
        """Should the SQL we pass to BQML and other I/O include the index?"""

        return self._has_index and (
            self.index.name is not None or len(self.index.names) > 1
        )

    def _to_placeholder_table(self, dry_run: bool = False) -> bigquery.TableReference:
        """Compiles this DataFrame's expression tree to SQL and saves it to a
        (temporary) view or table (in the case of a dry run).
        """
        return self._block.to_placeholder_table(
            include_index=self._should_sql_have_index(), dry_run=dry_run
        )

    def _to_sql_query(
        self, include_index: bool, enable_cache: bool = True
    ) -> Tuple[str, list[str], list[blocks.Label]]:
        """Compiles this DataFrame's expression tree to SQL, optionally
        including index columns.

        Args:
            include_index (bool):
                whether to include index columns.

        Returns:
            Tuple[sql_string, index_column_id_list, index_column_label_list]:
                If include_index is set to False, index_column_id_list and index_column_label_list
                return empty lists.
        """
        return self._block.to_sql_query(include_index, enable_cache=enable_cache)

    @property
    def sql(self) -> str:
        """Compiles this DataFrame's expression tree to SQL.

        Returns:
            str:
                string representing the compiled SQL.
        """
        try:
            include_index = self._should_sql_have_index()
            sql, _, _ = self._to_sql_query(include_index=include_index)
            return sql
        except AttributeError as e:
            # Workaround for a development-mode debugging issue:
            # An `AttributeError` originating *inside* this @property getter (e.g., due to
            # a typo or referencing a non-existent attribute) can be mistakenly intercepted
            # by the class's __getattr__ method if one is defined.
            # We catch the AttributeError and raise SyntaxError instead to make it clear
            # the error originates *here* in the property implementation.
            # See: https://stackoverflow.com/questions/50542177/correct-handling-of-attributeerror-in-getattr-when-using-property
            raise SyntaxError(
                "AttributeError encountered. Please check the implementation for incorrect attribute access."
            ) from e

    @property
    def query_job(self) -> Optional[bigquery.QueryJob]:
        """BigQuery job metadata for the most recent query.

        Returns:
            None or google.cloud.bigquery.QueryJob:
                The most recent `QueryJob
                <https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob>`_.
        """
        if self._query_job is None:
            self._set_internal_query_job(self._compute_dry_run())
        return self._query_job

    def memory_usage(self, index: bool = True):
        n_rows, _ = self.shape
        # like pandas, treat all variable-size objects as just 8-byte pointers, ignoring actual object
        column_sizes = self.dtypes.map(
            lambda dtype: bigframes.dtypes.DTYPE_BYTE_SIZES.get(dtype, 8) * n_rows
        )
        if index and self._has_index:
            index_size = pandas.Series([self.index._memory_usage()], index=["Index"])
            column_sizes = pandas.concat([index_size, column_sizes])
        return column_sizes

    def info(
        self,
        verbose: Optional[bool] = None,
        buf=None,
        max_cols: Optional[int] = None,
        memory_usage: Optional[bool] = None,
        show_counts: Optional[bool] = None,
    ):
        obuf = buf or sys.stdout

        n_rows, n_columns = self.shape

        max_cols = (
            max_cols
            if max_cols is not None
            else bigframes.options.display.max_info_columns
        )

        show_all_columns = verbose if verbose is not None else (n_columns < max_cols)

        obuf.write(f"{type(self)}\n")

        if self._block.has_index:
            index_type = "MultiIndex" if self.index.nlevels > 1 else "Index"

            index_stats = f"{n_rows} entries"
            if n_rows > 0:
                # These accessses are kind of expensive, maybe should try to skip?
                first_indice = self.index[0]
                last_indice = self.index[-1]
                index_stats += f", {first_indice} to {last_indice}"
            obuf.write(f"{index_type}: {index_stats}\n")
        else:
            obuf.write("NullIndex\n")

        if n_columns == 0:
            # We don't display any more information if the dataframe has no columns
            obuf.write("Empty DataFrame\n")
            return

        dtype_strings = self.dtypes.astype("string")
        if show_all_columns:
            obuf.write(f"Data columns (total {n_columns} columns):\n")
            column_info = self.columns.to_frame(name="Column")

            max_rows = bigframes.options.display.max_info_rows
            too_many_rows = n_rows > max_rows if max_rows is not None else False

            if show_counts if show_counts is not None else (not too_many_rows):
                non_null_counts = self.count().to_pandas()
                column_info["Non-Null Count"] = non_null_counts.map(
                    lambda x: f"{int(x)} non-null"
                )

            column_info["Dtype"] = dtype_strings

            column_info = column_info.reset_index(drop=True)
            column_info.index.name = "#"

            column_info_formatted = tabulate.tabulate(column_info, headers="keys")  # type: ignore
            obuf.write(column_info_formatted)
            obuf.write("\n")

        else:  # Just number of columns and first, last
            obuf.write(
                f"Columns: {n_columns} entries, {self.columns[0]} to {self.columns[-1]}\n"
            )
        dtype_counts = dtype_strings.value_counts().sort_index(ascending=True).items()
        dtype_counts_formatted = ", ".join(
            f"{dtype}({count})" for dtype, count in dtype_counts
        )
        obuf.write(f"dtypes: {dtype_counts_formatted}\n")

        show_memory = (
            memory_usage
            if memory_usage is not None
            else bigframes.options.display.memory_usage
        )
        if show_memory:
            # TODO: Convert to different units (kb, mb, etc.)
            obuf.write(f"memory usage: {self.memory_usage().sum()} bytes\n")

    def select_dtypes(self, include=None, exclude=None) -> DataFrame:
        # Create empty pandas dataframe with same schema and then leverage actual pandas implementation
        as_pandas = pandas.DataFrame(
            {
                col_id: pandas.Series([], dtype=dtype)
                for col_id, dtype in zip(self._block.value_columns, self._block.dtypes)
            }
        )
        selected_columns = tuple(
            as_pandas.select_dtypes(include=include, exclude=exclude).columns
        )
        return DataFrame(self._block.select_columns(selected_columns))

    def _set_internal_query_job(self, query_job: Optional[bigquery.QueryJob]):
        self._query_job = query_job

    @overload
    def __getitem__(
        self,
        key: bigframes.series.Series,
    ) -> DataFrame:
        ...

    @overload
    def __getitem__(
        self,
        key: slice,
    ) -> DataFrame:
        ...

    @overload
    def __getitem__(
        self,
        key: List[str],
    ) -> DataFrame:
        ...

    @overload
    def __getitem__(
        self,
        key: List[blocks.Label],
    ) -> DataFrame:
        ...

    @overload
    def __getitem__(self, key: pandas.Index) -> DataFrame:
        ...

    @overload
    def __getitem__(
        self,
        key: blocks.Label,
    ) -> bigframes.series.Series:
        ...

    def __getitem__(
        self,
        key: Union[
            blocks.Label,
            List[str],
            List[blocks.Label],
            # Index of column labels can be treated the same as a sequence of column labels.
            pandas.Index,
            bigframes.series.Series,
            slice,
        ],
    ):  # No return type annotations (like pandas) as type cannot always be determined statically
        # NOTE: This implements the operations described in
        # https://pandas.pydata.org/docs/getting_started/intro_tutorials/03_subset_data.html

        if isinstance(key, bigframes.series.Series):
            return self._getitem_bool_series(key)

        if isinstance(key, slice):
            return self.iloc[key]

        # TODO(tswast): Fix this pylance warning: Class overlaps "Hashable"
        # unsafely and could produce a match at runtime
        if isinstance(key, blocks.Label):
            return self._getitem_label(key)

        if utils.is_list_like(key):
            return self._getitem_columns(key)
        else:
            # TODO(tswast): What case is this supposed to be handling?
            return self._getitem_columns([cast(Hashable, key)])

    __getitem__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__getitem__)

    def _getitem_columns(self, key: Sequence[blocks.Label]) -> DataFrame:
        selected_ids: Tuple[str, ...] = ()
        for label in key:
            col_ids = self._block.label_to_col_id[label]
            selected_ids = (*selected_ids, *col_ids)

        return DataFrame(self._block.select_columns(selected_ids))

    def _getitem_label(self, key: blocks.Label):
        col_ids = self._block.cols_matching_label(key)
        if len(col_ids) == 0:
            raise KeyError(
                f"{key} not found in DataFrame columns: {self._block.column_labels}"
            )
        block = self._block.select_columns(col_ids)
        if isinstance(self.columns, pandas.MultiIndex):
            # Multiindex should drop-level if not selecting entire
            key_levels = len(key) if isinstance(key, tuple) else 1
            index_levels = self.columns.nlevels
            if key_levels < index_levels:
                block = block.with_column_labels(
                    block.column_labels.droplevel(list(range(key_levels)))
                )
                # Force return DataFrame in this case, even if only single column
                return DataFrame(block)

        if len(col_ids) == 1:
            return bigframes.series.Series(block)
        return DataFrame(block)

    # Bool Series selects rows
    def _getitem_bool_series(self, key: bigframes.series.Series) -> DataFrame:
        if not key.dtype == pandas.BooleanDtype():
            raise NotImplementedError(
                f"Only boolean series currently supported for indexing. {constants.FEEDBACK_LINK}"
            )
            # TODO: enforce stricter alignment
        combined_index, (
            get_column_left,
            get_column_right,
        ) = self._block.join(key._block, how="left")
        block = combined_index
        filter_col_id = get_column_right[key._value_column]
        block = block.filter_by_id(filter_col_id)
        block = block.drop_columns([filter_col_id])
        return DataFrame(block)

    def __getattr__(self, key: str):
        # To allow subclasses to set private attributes before the class is
        # fully initialized, protect against recursion errors with
        # uninitialized DataFrame objects. Note: this comes at the downside
        # that columns with a leading `_` won't be treated as columns.
        #
        # See:
        # https://github.com/googleapis/python-bigquery-dataframes/issues/728
        # and
        # https://nedbatchelder.com/blog/201010/surprising_getattr_recursion.html
        if key == "_block":
            raise AttributeError(key)

        if key in self._block.column_labels:
            return self.__getitem__(key)

        if hasattr(pandas.DataFrame, key):
            log_adapter.submit_pandas_labels(
                self._block.expr.session.bqclient, self.__class__.__name__, key
            )
            raise AttributeError(
                textwrap.dedent(
                    f"""
                    BigQuery DataFrames has not yet implemented an equivalent to
                    'pandas.DataFrame.{key}'. {constants.FEEDBACK_LINK}
                    """
                )
            )
        raise AttributeError(key)

    def __setattr__(self, key: str, value):
        if key == "_block":
            object.__setattr__(self, key, value)
            return

        # To allow subclasses to set private attributes before the class is
        # fully initialized, assume anything set before `_block` is initialized
        # is a regular attribute.
        if not hasattr(self, "_block"):
            object.__setattr__(self, key, value)
            return

        # If someone has a column named the same as a normal attribute
        # (e.g. index), we want to set the normal attribute, not the column.
        # To do that, check if there is a normal attribute by using
        # __getattribute__ (not __getattr__, because that includes columns).
        # If that returns a value without raising, then we know this is a
        # normal attribute and we should prefer that.
        try:
            object.__getattribute__(self, key)
            return object.__setattr__(self, key, value)
        except AttributeError:
            pass

        # If we made it here, then we know that it's not a regular attribute
        # already, so it might be a column to update. Note: we don't allow
        # adding new columns using __setattr__, only __setitem__, that way we
        # can still add regular new attributes.
        if key in self._block.column_labels:
            self[key] = value
        else:
            object.__setattr__(self, key, value)

    def __repr__(self) -> str:
        """Converts a DataFrame to a string. Calls to_pandas.

        Only represents the first `bigframes.options.display.max_rows`.
        """
        # Protect against errors with uninitialized DataFrame. See:
        # https://github.com/googleapis/python-bigquery-dataframes/issues/728
        if not hasattr(self, "_block"):
            return object.__repr__(self)

        opts = bigframes.options.display
        max_results = opts.max_rows
        if opts.repr_mode == "deferred":
            return formatter.repr_query_job(self._compute_dry_run())

        # TODO(swast): pass max_columns and get the true column count back. Maybe
        # get 1 more column than we have requested so that pandas can add the
        # ... for us?
        pandas_df, row_count, query_job = self._block.retrieve_repr_request_results(
            max_results
        )

        self._set_internal_query_job(query_job)
        from bigframes.display import plaintext

        return plaintext.create_text_representation(
            pandas_df,
            row_count,
            is_series=False,
            has_index=self._has_index,
            column_count=len(self.columns),
        )

    def _get_display_df_and_blob_cols(self) -> tuple[DataFrame, list[str]]:
        """Process blob columns for display."""
        df = self
        blob_cols = []
        if bigframes.options.display.blob_display:
            blob_cols = [
                series_name
                for series_name, series in self.items()
                if series.dtype == bigframes.dtypes.OBJ_REF_DTYPE
            ]
            if blob_cols:
                df = self.copy()
                for col in blob_cols:
                    # TODO(garrettwu): Not necessary to get access urls for all the rows. Update when having a to get URLs from local data.
                    df[col] = df[col].blob._get_runtime(mode="R", with_metadata=True)
        return df, blob_cols

    def _repr_mimebundle_(self, include=None, exclude=None):
        """
        Custom display method for IPython/Jupyter environments.
        This is called by IPython's display system when the object is displayed.
        """
        # TODO(b/467647693): Anywidget integration has been tested in Jupyter, VS Code, and
        # BQ Studio, but there is a known compatibility issue with Marimo that needs to be addressed.
        from bigframes.display import html

        return html.repr_mimebundle(self, include=include, exclude=exclude)

    def __delitem__(self, key: str):
        df = self.drop(columns=[key])
        self._set_block(df._get_block())

    def __setitem__(
        self,
        key: str | list[str] | pandas.Index,
        value: SingleItemValue | MultiItemValue,
    ):
        if isinstance(key, (list, pandas.Index)):
            df = self._assign_multi_items(key, value)
        else:
            df = self._assign_single_item(key, value)
        self._set_block(df._get_block())

    __setitem__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__setitem__)

    def _apply_binop(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        op,
        axis: str | int = "columns",
        how: str = "outer",
        reverse: bool = False,
    ):
        if isinstance(other, bigframes.dtypes.LOCAL_SCALAR_TYPES):
            return self._apply_scalar_binop(other, op, reverse=reverse)
        elif isinstance(other, DataFrame):
            return self._apply_dataframe_binop(other, op, how=how, reverse=reverse)
        elif isinstance(other, pandas.DataFrame):
            return self._apply_dataframe_binop(
                DataFrame(other), op, how=how, reverse=reverse
            )
        elif utils.get_axis_number(axis) == 0:
            return self._apply_series_binop_axis_0(other, op, how, reverse)
        elif utils.get_axis_number(axis) == 1:
            return self._apply_series_binop_axis_1(other, op, how, reverse)
        raise NotImplementedError(
            f"binary operation is not implemented on the second operand of type {type(other).__name__}."
            f"{constants.FEEDBACK_LINK}"
        )

    def _apply_scalar_binop(
        self,
        other: bigframes.dtypes.LOCAL_SCALAR_TYPE,
        op: ops.BinaryOp,
        reverse: bool = False,
    ) -> DataFrame:
        if reverse:
            expr = op.as_expr(
                left_input=ex.const(other),
                right_input=ex.free_var("var1"),
            )
        else:
            expr = op.as_expr(
                left_input=ex.free_var("var1"),
                right_input=ex.const(other),
            )
        return DataFrame(self._block.multi_apply_unary_op(expr))

    def _apply_series_binop_axis_0(
        self,
        other,
        op: ops.BinaryOp,
        how: str = "outer",
        reverse: bool = False,
    ) -> DataFrame:
        bf_series = bigframes.core.convert.to_bf_series(
            other, self.index if self._has_index else None, self._session
        )
        aligned_block, columns, expr_pairs = self._block._align_axis_0(
            bf_series._block, how=how
        )
        result = aligned_block._apply_binop(
            op, inputs=expr_pairs, labels=columns, reverse=reverse
        )
        return DataFrame(result)

    def _apply_series_binop_axis_1(
        self,
        other,
        op: ops.BinaryOp,
        how: str = "outer",
        reverse: bool = False,
    ) -> DataFrame:
        """Align dataframe with pandas series by inlining series values as literals."""
        # If we already know the transposed schema (from the transpose cache), we don't need to materialize rows from other
        # Instead, can fully defer execution (as a cross-join)
        if (
            isinstance(other, bigframes.series.Series)
            and other._block._transpose_cache is not None
        ):
            aligned_block, columns, expr_pairs = self._block._align_series_block_axis_1(
                other._block, how=how
            )
        else:
            # Fallback path, materialize `other` locally
            pd_series = bigframes.core.convert.to_pd_series(other, self.columns)
            aligned_block, columns, expr_pairs = self._block._align_pd_series_axis_1(
                pd_series, how=how
            )
        result = aligned_block._apply_binop(
            op, inputs=expr_pairs, labels=columns, reverse=reverse
        )
        return DataFrame(result)

    def _apply_dataframe_binop(
        self,
        other: DataFrame,
        op: ops.BinaryOp,
        how: str = "outer",
        reverse: bool = False,
    ) -> DataFrame:
        aligned_block, columns, expr_pairs = self._block._align_both_axes(
            other._block, how=how
        )
        result = aligned_block._apply_binop(
            op, inputs=expr_pairs, labels=columns, reverse=reverse
        )
        return DataFrame(result)

    def eq(self, other: typing.Any, axis: str | int = "columns") -> DataFrame:
        return self._apply_binop(other, ops.eq_op, axis=axis)

    def __eq__(self, other) -> DataFrame:  # type: ignore
        return self.eq(other)

    __eq__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__eq__)

    def ne(self, other: typing.Any, axis: str | int = "columns") -> DataFrame:
        return self._apply_binop(other, ops.ne_op, axis=axis)

    def __ne__(self, other) -> DataFrame:  # type: ignore
        return self.ne(other)

    __ne__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__ne__)

    def __invert__(self) -> DataFrame:
        return self._apply_unary_op(ops.invert_op)

    __invert__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__invert__)

    def le(self, other: typing.Any, axis: str | int = "columns") -> DataFrame:
        return self._apply_binop(other, ops.le_op, axis=axis)

    def __le__(self, other) -> DataFrame:
        return self.le(other)

    __le__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__le__)

    def lt(self, other: typing.Any, axis: str | int = "columns") -> DataFrame:
        return self._apply_binop(other, ops.lt_op, axis=axis)

    def __lt__(self, other) -> DataFrame:
        return self.lt(other)

    __lt__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__lt__)

    def ge(self, other: typing.Any, axis: str | int = "columns") -> DataFrame:
        return self._apply_binop(other, ops.ge_op, axis=axis)

    def __ge__(self, other) -> DataFrame: