feat: add equals methods to series/dataframe (#76)

* feat: add equals methods to series/dataframe

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -25,6 +25,39 @@
 import bigframes.operations.aggregations as agg_ops
 
 
+def equals(block1: blocks.Block, block2: blocks.Block) -> bool:
+    if not block1.column_labels.equals(block2.column_labels):
+        return False
+    if block1.dtypes != block2.dtypes:
+        return False
+    # TODO: More advanced expression tree traversals to short circuit actually querying data
+
+    block1 = block1.reset_index(drop=False)
+    block2 = block2.reset_index(drop=False)
+
+    joined, (lmap, rmap) = block1.index.join(block2.index, how="outer")
+    joined_block = joined._block
+
+    equality_ids = []
+    for lcol, rcol in zip(block1.value_columns, block2.value_columns):
+        lcolmapped = lmap(lcol)
+        rcolmapped = rmap(rcol)
+        joined_block, result_id = joined_block.apply_binary_op(
+            lcolmapped, rcolmapped, ops.eq_nulls_match_op
+        )
+        joined_block, result_id = joined_block.apply_unary_op(
+            result_id, ops.partial_right(ops.fillna_op, False)
+        )
+        equality_ids.append(result_id)
+
+    joined_block = joined_block.select_columns(equality_ids).with_column_labels(
+        list(range(len(equality_ids)))
+    )
+    stacked_block = joined_block.stack(dropna=False, sort=False)
+    result = stacked_block.get_stat(stacked_block.value_columns[0], agg_ops.all_op)
+    return typing.cast(bool, result)
+
+
 def indicate_duplicates(
     block: blocks.Block, columns: typing.Sequence[str], keep: str = "first"
 ) -> typing.Tuple[blocks.Block, str]:

bigframes/dataframe.py

@@ -1066,6 +1066,12 @@ def rename_axis(
             labels = [mapper]
         return DataFrame(self._block.with_index_labels(labels))
 
+    def equals(self, other: typing.Union[bigframes.series.Series, DataFrame]) -> bool:
+        # Must be same object type, same column dtypes, and same label values
+        if not isinstance(other, DataFrame):
+            return False
+        return block_ops.equals(self._block, other._block)
+
     def assign(self, **kwargs) -> DataFrame:
         # TODO(garrettwu) Support list-like values. Requires ordering.
         # TODO(garrettwu) Support callable values.

bigframes/operations/__init__.py

@@ -705,6 +705,16 @@ def eq_op(
     return x == y
 
 
+def eq_nulls_match_op(
+    x: ibis_types.Value,
+    y: ibis_types.Value,
+):
+    """Variant of eq_op where nulls match each other. Only use where dtypes are known to be same."""
+    left = x.cast(ibis_dtypes.str).fillna(ibis_types.literal("$NULL_SENTINEL$"))
+    right = y.cast(ibis_dtypes.str).fillna(ibis_types.literal("$NULL_SENTINEL$"))
+    return left == right
+
+
 def ne_op(
     x: ibis_types.Value,
     y: ibis_types.Value,

bigframes/series.py

@@ -209,6 +209,14 @@ def rename_axis(
             labels = [mapper]
         return Series(self._block.with_index_labels(labels))
 
+    def equals(
+        self, other: typing.Union[Series, bigframes.dataframe.DataFrame]
+    ) -> bool:
+        # Must be same object type, same column dtypes, and same label values
+        if not isinstance(other, Series):
+            return False
+        return block_ops.equals(self._block, other._block)
+
     def reset_index(
         self,
         *,

tests/system/small/test_dataframe.py

@@ -2551,6 +2551,74 @@ def test_df_reindex_columns(scalars_df_index, scalars_pandas_df_index):
     )
 
 
+def test_df_equals_identical(scalars_df_index, scalars_pandas_df_index):
+    unsupported = [
+        "geography_col",
+    ]
+    scalars_df_index = scalars_df_index.drop(columns=unsupported)
+    scalars_pandas_df_index = scalars_pandas_df_index.drop(columns=unsupported)
+
+    bf_result = scalars_df_index.equals(scalars_df_index)
+    pd_result = scalars_pandas_df_index.equals(scalars_pandas_df_index)
+
+    assert pd_result == bf_result
+
+
+def test_df_equals_series(scalars_df_index, scalars_pandas_df_index):
+    bf_result = scalars_df_index[["int64_col"]].equals(scalars_df_index["int64_col"])
+    pd_result = scalars_pandas_df_index[["int64_col"]].equals(
+        scalars_pandas_df_index["int64_col"]
+    )
+
+    assert pd_result == bf_result
+
+
+def test_df_equals_different_dtype(scalars_df_index, scalars_pandas_df_index):
+    columns = ["int64_col", "int64_too"]
+    scalars_df_index = scalars_df_index[columns]
+    scalars_pandas_df_index = scalars_pandas_df_index[columns]
+
+    bf_modified = scalars_df_index.copy()
+    bf_modified = bf_modified.astype("Float64")
+
+    pd_modified = scalars_pandas_df_index.copy()
+    pd_modified = pd_modified.astype("Float64")
+
+    bf_result = scalars_df_index.equals(bf_modified)
+    pd_result = scalars_pandas_df_index.equals(pd_modified)
+
+    assert pd_result == bf_result
+
+
+def test_df_equals_different_values(scalars_df_index, scalars_pandas_df_index):
+    columns = ["int64_col", "int64_too"]
+    scalars_df_index = scalars_df_index[columns]
+    scalars_pandas_df_index = scalars_pandas_df_index[columns]
+
+    bf_modified = scalars_df_index.copy()
+    bf_modified["int64_col"] = bf_modified.int64_col + 1
+
+    pd_modified = scalars_pandas_df_index.copy()
+    pd_modified["int64_col"] = pd_modified.int64_col + 1
+
+    bf_result = scalars_df_index.equals(bf_modified)
+    pd_result = scalars_pandas_df_index.equals(pd_modified)
+
+    assert pd_result == bf_result
+
+
+def test_df_equals_extra_column(scalars_df_index, scalars_pandas_df_index):
+    columns = ["int64_col", "int64_too"]
+    more_columns = ["int64_col", "int64_too", "float64_col"]
+
+    bf_result = scalars_df_index[columns].equals(scalars_df_index[more_columns])
+    pd_result = scalars_pandas_df_index[columns].equals(
+        scalars_pandas_df_index[more_columns]
+    )
+
+    assert pd_result == bf_result
+
+
 def test_df_reindex_like(scalars_df_index, scalars_pandas_df_index):
     reindex_target_bf = scalars_df_index.reindex(
         columns=["not_a_col", "int64_col", "int64_too"], index=[5, 1, 3, 99, 1]

tests/system/small/test_series.py

@@ -112,6 +112,44 @@ def test_series_get_column_default(scalars_dfs):
     assert result == "default_val"
 
 
+def test_series_equals_identical(scalars_df_index, scalars_pandas_df_index):
+    bf_result = scalars_df_index.int64_col.equals(scalars_df_index.int64_col)
+    pd_result = scalars_pandas_df_index.int64_col.equals(
+        scalars_pandas_df_index.int64_col
+    )
+
+    assert pd_result == bf_result
+
+
+def test_series_equals_df(scalars_df_index, scalars_pandas_df_index):
+    bf_result = scalars_df_index["int64_col"].equals(scalars_df_index[["int64_col"]])
+    pd_result = scalars_pandas_df_index["int64_col"].equals(
+        scalars_pandas_df_index[["int64_col"]]
+    )
+
+    assert pd_result == bf_result
+
+
+def test_series_equals_different_dtype(scalars_df_index, scalars_pandas_df_index):
+    bf_series = scalars_df_index["int64_col"]
+    pd_series = scalars_pandas_df_index["int64_col"]
+
+    bf_result = bf_series.equals(bf_series.astype("Float64"))
+    pd_result = pd_series.equals(pd_series.astype("Float64"))
+
+    assert pd_result == bf_result
+
+
+def test_series_equals_different_values(scalars_df_index, scalars_pandas_df_index):
+    bf_series = scalars_df_index["int64_col"]
+    pd_series = scalars_pandas_df_index["int64_col"]
+
+    bf_result = bf_series.equals(bf_series + 1)
+    pd_result = pd_series.equals(pd_series + 1)
+
+    assert pd_result == bf_result
+
+
 def test_series_get_with_default_index(scalars_dfs):
     col_name = "float64_col"
     key = 2

third_party/bigframes_vendored/pandas/core/frame.py

@@ -399,6 +399,28 @@ def to_orc(self, path=None, **kwargs) -> bytes | None:
     # ----------------------------------------------------------------------
     # Unsorted
 
+    def equals(self, other) -> bool:
+        """
+        Test whether two objects contain the same elements.
+
+        This function allows two Series or DataFrames to be compared against
+        each other to see if they have the same shape and elements. NaNs in
+        the same location are considered equal.
+
+        The row/column index do not need to have the same type, as long
+        as the values are considered equal. Corresponding columns must be of
+        the same dtype.
+
+        Args:
+            other (Series or DataFrame):
+                The other Series or DataFrame to be compared with the first.
+
+        Returns:
+            bool: True if all elements are the same in both objects, False
+            otherwise.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     def assign(self, **kwargs) -> DataFrame:
         r"""
         Assign new columns to a DataFrame.