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 os
import re
import sys
import textwrap
import typing
from typing import (
    Callable,
    Dict,
    Iterable,
    List,
    Literal,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Union,
)

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
import tabulate

import bigframes
import bigframes._config.display_options as display_options
import bigframes.constants as constants
import bigframes.core
from bigframes.core import log_adapter
import bigframes.core.block_transforms as block_ops
import bigframes.core.blocks as blocks
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.ordering as order
import bigframes.core.utils as utils
import bigframes.core.window
import bigframes.dtypes
import bigframes.formatting_helpers as formatter
import bigframes.operations as ops
import bigframes.operations.aggregations as agg_ops
import bigframes.operations.plotting as plotting
import bigframes.series
import bigframes.series as bf_series
import bigframes.session._io.bigquery

if typing.TYPE_CHECKING:
    import bigframes.session

    SingleItemValue = Union[bigframes.series.Series, int, float, 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__

    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,
    ):
        if copy is not None and not copy:
            raise ValueError(
                f"DataFrame constructor only supports copy=True. {constants.FEEDBACK_LINK}"
            )

        # 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], bf_series.Series) for key in data.keys())
        ):
            if not all(isinstance(data[key], bf_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(bf_series.Series, first_series)
                ._get_block()
                .with_column_labels([first_label])
            )

            for key in keys[1:]:
                other = typing.cast(bf_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:
                block = block.select_columns(list(columns))  # type:ignore
            if dtype:
                block = block.multi_apply_unary_op(
                    block.value_columns, ops.AsTypeOp(to_type=dtype)
                )
            self._block = block

        else:
            import bigframes.pandas

            pd_dataframe = pandas.DataFrame(
                data=data,
                index=index,  # type:ignore
                columns=columns,  # type:ignore
                dtype=dtype,  # type:ignore
            )
            if session:
                self._block = session.read_pandas(pd_dataframe)._get_block()
            else:
                self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block()
        self._query_job: Optional[bigquery.QueryJob] = None

        # Runs strict validations to ensure internal type predictions and ibis are completely in sync
        # Do not execute these validations outside of testing suite.
        if "PYTEST_CURRENT_TEST" in os.environ:
            self._block.expr.validate_schema()

    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]:
        """Returns the column id matching the label if there is exactly
        one such column. If there are multiple columns with the same name,
        raises an error. If there is no such column, returns None."""
        matches = self._block.label_to_col_id.get(label, [])
        if len(matches) > 1:
            raise ValueError(
                f"Multiple columns matching id {label} were found. {constants.FEEDBACK_LINK}"
            )
        return matches[0] if len(matches) != 0 else None

    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
    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
    def loc(self) -> indexers.LocDataFrameIndexer:
        return indexers.LocDataFrameIndexer(self)

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

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

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

    @property
    def dtypes(self) -> pandas.Series:
        return pandas.Series(data=self._block.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 _session(self) -> bigframes.Session:
        return self._get_block().expr.session

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

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

    def astype(
        self,
        dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype],
    ) -> DataFrame:
        return self._apply_unary_op(ops.AsTypeOp(to_type=dtype))

    def _to_sql_query(
        self, include_index: bool
    ) -> 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:
            a tuple of (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)

    @property
    def sql(self) -> str:
        """Compiles this DataFrame's expression tree to SQL."""
        sql, _, _ = self._to_sql_query(include_index=False)
        return sql

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

        Returns:
            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:
            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")

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

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

        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: bigquery.QueryJob):
        self._query_job = query_job

    def __getitem__(
        self,
        key: Union[
            blocks.Label,
            Sequence[blocks.Label],
            # Index of column labels can be treated the same as a sequence of column labels.
            pandas.Index,
            bigframes.series.Series,
        ],
    ):  # No return type annotations (like pandas) as type cannot always be determined statically
        """Gets the specified column(s) from the DataFrame."""
        # 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, typing.Hashable):
            return self._getitem_label(key)
        # Select a subset of columns or re-order columns.
        # In Ibis after you apply a projection, any column objects from the
        # table before the projection can't be combined with column objects
        # from the table after the projection. This is because the table after
        # a projection is considered a totally separate table expression.
        #
        # This is unexpected behavior for a pandas user, who expects their old
        # Series objects to still work with the new / mutated DataFrame. We
        # avoid applying a projection in Ibis until it's absolutely necessary
        # to provide pandas-like semantics.
        # TODO(swast): Do we need to apply implicit join when doing a
        # projection?

        # Select a number of columns as DF.
        key = key if utils.is_list_like(key) else [key]  # type:ignore

        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(key)
        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):
        if key in self._block.column_labels:
            return self.__getitem__(key)
        elif hasattr(pandas.DataFrame, key):
            raise AttributeError(
                textwrap.dedent(
                    f"""
                    BigQuery DataFrames has not yet implemented an equivalent to
                    'pandas.DataFrame.{key}'. {constants.FEEDBACK_LINK}
                    """
                )
            )
        else:
            raise AttributeError(key)

    def __setattr__(self, key: str, value):
        if key in ["_block", "_query_job"]:
            object.__setattr__(self, key, value)
            return
        # Can this be removed???
        try:
            # boring attributes go through boring old path
            object.__getattribute__(self, key)
            return object.__setattr__(self, key, value)
        except AttributeError:
            pass

        # if this fails, go on to more involved attribute setting
        # (note that this matches __getattr__, above).
        try:
            if key in self.columns:
                self[key] = value
            else:
                object.__setattr__(self, key, value)
        # Can this be removed?
        except (AttributeError, TypeError):
            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`.
        """
        opts = bigframes.options.display
        max_results = opts.max_rows
        if opts.repr_mode == "deferred":
            return formatter.repr_query_job(self.query_job)

        self._cached()
        # 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)

        column_count = len(pandas_df.columns)

        with display_options.pandas_repr(opts):
            repr_string = repr(pandas_df)

        # Modify the end of the string to reflect count.
        lines = repr_string.split("\n")
        pattern = re.compile("\\[[0-9]+ rows x [0-9]+ columns\\]")
        if pattern.match(lines[-1]):
            lines = lines[:-2]

        if row_count > len(lines) - 1:
            lines.append("...")

        lines.append("")
        lines.append(f"[{row_count} rows x {column_count} columns]")
        return "\n".join(lines)

    def _repr_html_(self) -> str:
        """
        Returns an html string primarily for use by notebooks for displaying
        a representation of the DataFrame. Displays 20 rows by default since
        many notebooks are not configured for large tables.
        """
        opts = bigframes.options.display
        max_results = bigframes.options.display.max_rows
        if opts.repr_mode == "deferred":
            return formatter.repr_query_job_html(self.query_job)

        self._cached()
        # 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)

        column_count = len(pandas_df.columns)

        with display_options.pandas_repr(opts):
            # _repr_html_ stub is missing so mypy thinks it's a Series. Ignore mypy.
            html_string = pandas_df._repr_html_()  # type:ignore

        html_string += f"[{row_count} rows x {column_count} columns in total]"
        return html_string

    def __setitem__(self, key: str, value: SingleItemValue):
        """Modify or insert a column into the DataFrame.

        Note: This does **not** modify the original table the DataFrame was
        derived from.
        """
        df = self._assign_single_item(key, value)
        self._set_block(df._get_block())

    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, (float, int)):
            return self._apply_scalar_binop(other, op, reverse=reverse)
        elif isinstance(other, indexes.Index):
            return self._apply_series_binop(
                other.to_series(index=self.index),
                op,
                axis=axis,
                how=how,
                reverse=reverse,
            )
        elif isinstance(other, bigframes.series.Series):
            return self._apply_series_binop(
                other, op, axis=axis, how=how, reverse=reverse
            )
        elif isinstance(other, DataFrame):
            return self._apply_dataframe_binop(other, op, how=how, reverse=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: float | int, op: ops.BinaryOp, reverse: bool = False
    ) -> DataFrame:
        block = self._block
        for column_id, label in zip(
            self._block.value_columns, self._block.column_labels
        ):
            expr = (
                op.as_expr(ex.const(other), column_id)
                if reverse
                else op.as_expr(column_id, ex.const(other))
            )
            block, _ = block.project_expr(expr, label)
            block = block.drop_columns([column_id])
        return DataFrame(block)

    def _apply_series_binop(
        self,
        other: bigframes.series.Series,
        op: ops.BinaryOp,
        axis: str | int = "columns",
        how: str = "outer",
        reverse: bool = False,
    ) -> DataFrame:
        if axis not in ("columns", "index", 0, 1):
            raise ValueError(f"Invalid input: axis {axis}.")

        if axis in ("columns", 1):
            raise NotImplementedError(
                f"Row Series operations haven't been supported. {constants.FEEDBACK_LINK}"
            )

        block, (get_column_left, get_column_right) = self._block.join(
            other._block, how=how
        )

        series_column_id = other._value_column
        series_col = get_column_right[series_column_id]
        for column_id, label in zip(
            self._block.value_columns, self._block.column_labels
        ):
            self_col = get_column_left[column_id]
            expr = (
                op.as_expr(series_col, self_col)
                if reverse
                else op.as_expr(self_col, series_col)
            )
            block, _ = block.project_expr(expr, label)
            block = block.drop_columns([get_column_left[column_id]])

        block = block.drop_columns([series_col])
        block = block.with_index_labels(self.index.names)
        return DataFrame(block)

    def _apply_dataframe_binop(
        self,
        other: DataFrame,
        op: ops.BinaryOp,
        how: str = "outer",
        reverse: bool = False,
    ) -> DataFrame:
        # Join rows
        block, (get_column_left, get_column_right) = self._block.join(
            other._block, how=how
        )
        # join columns schema
        # indexers will be none for exact match
        columns, lcol_indexer, rcol_indexer = self.columns.join(
            other.columns, how=how, return_indexers=True
        )

        binop_result_ids = []

        column_indices = zip(
            lcol_indexer if (lcol_indexer is not None) else range(len(columns)),
            rcol_indexer if (lcol_indexer is not None) else range(len(columns)),
        )

        for left_index, right_index in column_indices:
            if left_index >= 0 and right_index >= 0:  # -1 indices indicate missing
                self_col_id = get_column_left[self._block.value_columns[left_index]]
                other_col_id = get_column_right[other._block.value_columns[right_index]]
                expr = (
                    op.as_expr(other_col_id, self_col_id)
                    if reverse
                    else op.as_expr(self_col_id, other_col_id)
                )
            elif left_index >= 0:
                self_col_id = get_column_left[self._block.value_columns[left_index]]
                expr = (
                    op.as_expr(ex.const(None), self_col_id)
                    if reverse
                    else op.as_expr(self_col_id, ex.const(None))
                )
            elif right_index >= 0:
                other_col_id = get_column_right[other._block.value_columns[right_index]]
                expr = (
                    op.as_expr(other_col_id, ex.const(None))
                    if reverse
                    else op.as_expr(ex.const(None), other_col_id)
                )
            else:
                # Should not be possible
                raise ValueError("No right or left index.")
            block, result_col_id = block.project_expr(expr)
            binop_result_ids.append(result_col_id)

        block = block.select_columns(binop_result_ids).with_column_labels(columns)
        return DataFrame(block)

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

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

    __eq__ = eq  # type: ignore

    __ne__ = ne  # type: ignore

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

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

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

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

    __lt__ = lt

    __le__ = le

    __gt__ = gt

    __ge__ = ge

    def add(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        # TODO(swast): Support fill_value parameter.
        # TODO(swast): Support level parameter with MultiIndex.
        return self._apply_binop(other, ops.add_op, axis=axis)

    __radd__ = __add__ = radd = add

    def sub(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        return self._apply_binop(other, ops.sub_op, axis=axis)

    __sub__ = subtract = sub

    def rsub(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        return self._apply_binop(other, ops.sub_op, axis=axis, reverse=True)

    __rsub__ = rsub

    def mul(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        return self._apply_binop(other, ops.mul_op, axis=axis)

    __rmul__ = __mul__ = rmul = multiply = mul

    def truediv(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        return self._apply_binop(other, ops.div_op, axis=axis)

    div = divide = __truediv__ = truediv

    def rtruediv(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        return self._apply_binop(other, ops.div_op, axis=axis, reverse=True)

    __rtruediv__ = rdiv = rtruediv

    def floordiv(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        return self._apply_binop(other, ops.floordiv_op, axis=axis)

    __floordiv__ = floordiv

    def rfloordiv(
        self,
        other: float | int | bigframes.series.Series | DataFrame,
        axis: str | int = "columns",
    ) -> DataFrame:
        return self._apply_binop(other, ops.floordiv_op, axis=axis, reverse=True)

    __rfloordiv__ = rfloordiv

    def mod(self, other: int | bigframes.series.Series | DataFrame, axis: str | int = "columns") -> DataFrame:  # type: ignore
        return self._apply_binop(other, ops.mod_op, axis=axis)

    def rmod(self, other: int | bigframes.series.Series | DataFrame, axis: str | int = "columns") -> DataFrame:  # type: ignore
        return self._apply_binop(other, ops.mod_op, axis=axis, reverse=True)

    __mod__ = mod

    __rmod__ = rmod

    def pow(
        self, other: int | bigframes.series.Series, axis: str | int = "columns"
    ) -> DataFrame:
        return self._apply_binop(other, ops.pow_op, axis=axis)

    def rpow(
        self, other: int | bigframes.series.Series, axis: str | int = "columns"
    ) -> DataFrame:
        return self._apply_binop(other, ops.pow_op, axis=axis, reverse=True)

    __pow__ = pow

    __rpow__ = rpow

    def align(
        self,
        other: typing.Union[DataFrame, bigframes.series.Series],
        join: str = "outer",
        axis: typing.Union[str, int, None] = None,
    ) -> typing.Tuple[
        typing.Union[DataFrame, bigframes.series.Series],
        typing.Union[DataFrame, bigframes.series.Series],
    ]:
        axis_n = utils.get_axis_number(axis) if axis else None
        if axis_n == 1 and isinstance(other, bigframes.series.Series):
            raise NotImplementedError(
                f"align with series and axis=1 not supported. {constants.FEEDBACK_LINK}"
            )
        left_block, right_block = block_ops.align(
            self._block, other._block, join=join, axis=axis
        )
        return DataFrame(left_block), other.__class__(right_block)

    def update(self, other, join: str = "left", overwrite=True, filter_func=None):
        other = other if isinstance(other, DataFrame) else DataFrame(other)
        if join != "left":
            raise ValueError("Only 'left' join supported for update")

        if filter_func is not None:  # Will always take other if possible

            def update_func(
                left: bigframes.series.Series, right: bigframes.series.Series
            ) -> bigframes.series.Series:
                return left.mask(right.notna() & filter_func(left), right)

        elif overwrite:

            def update_func(
                left: bigframes.series.Series, right: bigframes.series.Series
            ) -> bigframes.series.Series:
                return left.mask(right.notna(), right)

        else:

            def update_func(
                left: bigframes.series.Series, right: bigframes.series.Series
            ) -> bigframes.series.Series:
                return left.mask(left.isna(), right)

        result = self.combine(other, update_func, how=join)

        self._set_block(result._block)

    def combine(
        self,
        other: DataFrame,
        func: typing.Callable[
            [bigframes.series.Series, bigframes.series.Series], bigframes.series.Series
        ],
        fill_value=None,
        overwrite: bool = True,
        *,
        how: str = "outer",
    ) -> DataFrame:
        l_aligned, r_aligned = block_ops.align(self._block, other._block, join=how)

        other_missing_labels = self._block.column_labels.difference(
            other._block.column_labels
        )

        l_frame = DataFrame(l_aligned)
        r_frame = DataFrame(r_aligned)
        results = []
        for (label, lseries), (_, rseries) in zip(l_frame.items(), r_frame.items()):
            if not ((label in other_missing_labels) and not overwrite):
                if fill_value is not None:
                    result = func(
                        lseries.fillna(fill_value), rseries.fillna(fill_value)
                    )
                else:
                    result = func(lseries, rseries)
            else:
                result = (
                    lseries.fillna(fill_value) if fill_value is not None else lseries
                )