feat: Add Index.__eq__ for consts, aligned objects (#2141)

Author: TrevorBergeron

bigframes/core/indexes/base.py

@@ -754,6 +754,58 @@ def item(self):
         # Docstring is in third_party/bigframes_vendored/pandas/core/indexes/base.py
         return self.to_series().peek(2).item()
 
+    def __eq__(self, other) -> Index:  # type: ignore
+        return self._apply_binop(other, ops.eq_op)
+
+    def _apply_binop(self, other, op: ops.BinaryOp) -> Index:
+        # TODO: Handle local objects, or objects not implicitly alignable? Gets ambiguous with partial ordering though
+        if isinstance(other, (bigframes.series.Series, Index)):
+            other = Index(other)
+            if other.nlevels != self.nlevels:
+                raise ValueError("Dimensions do not match")
+
+            lexpr = self._block.expr
+            rexpr = other._block.expr
+            join_result = lexpr.try_row_join(rexpr)
+            if join_result is None:
+                raise ValueError("Cannot align objects")
+
+            expr, (lmap, rmap) = join_result
+
+            expr, res_ids = expr.compute_values(
+                [
+                    op.as_expr(lmap[lid], rmap[rid])
+                    for lid, rid in zip(lexpr.column_ids, rexpr.column_ids)
+                ]
+            )
+            return Index(
+                blocks.Block(
+                    expr.select_columns(res_ids),
+                    index_columns=res_ids,
+                    column_labels=[],
+                    index_labels=[None] * len(res_ids),
+                )
+            )
+        elif (
+            isinstance(other, bigframes.dtypes.LOCAL_SCALAR_TYPES) and self.nlevels == 1
+        ):
+            block, id = self._block.project_expr(
+                op.as_expr(self._block.index_columns[0], ex.const(other))
+            )
+            return Index(block.select_column(id))
+        elif isinstance(other, tuple) and len(other) == self.nlevels:
+            block = self._block.project_exprs(
+                [
+                    op.as_expr(self._block.index_columns[i], ex.const(other[i]))
+                    for i in range(self.nlevels)
+                ],
+                labels=[None] * self.nlevels,
+                drop=True,
+            )
+            return Index(block.set_index(block.value_columns))
+        else:
+            return NotImplemented
+
 
 def _should_create_datetime_index(block: blocks.Block) -> bool:
     if len(block.index.dtypes) != 1:

bigframes/core/indexes/multi.py

@@ -19,6 +19,8 @@
 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
 
 
@@ -46,3 +48,26 @@ def from_arrays(
         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))
+
+    def __eq__(self, other) -> Index:  # type: ignore
+        import bigframes.operations as ops
+        import bigframes.operations.aggregations as agg_ops
+
+        eq_result = self._apply_binop(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],
+            )
+        )

tests/system/small/test_index.py

@@ -668,3 +668,20 @@ def test_custom_index_setitem_error():
 
     with pytest.raises(TypeError, match="Index does not support mutable operations"):
         custom_index[2] = 999
+
+
+def test_index_eq_const(scalars_df_index, scalars_pandas_df_index):
+    bf_result = (scalars_df_index.index == 3).to_pandas()
+    pd_result = scalars_pandas_df_index.index == 3
+    assert bf_result == pd.Index(pd_result)
+
+
+def test_index_eq_aligned_index(scalars_df_index, scalars_pandas_df_index):
+    bf_result = (
+        bpd.Index(scalars_df_index.int64_col)
+        == bpd.Index(scalars_df_index.int64_col.abs())
+    ).to_pandas()
+    pd_result = pd.Index(scalars_pandas_df_index.int64_col) == pd.Index(
+        scalars_pandas_df_index.int64_col.abs()
+    )
+    assert bf_result == pd.Index(pd_result)

tests/system/small/test_multiindex.py

@@ -1474,3 +1474,13 @@ def test_multi_index_contains(scalars_df_index, scalars_pandas_df_index, key):
     pd_result = key in scalars_pandas_df_index.set_index(col_name).index
 
     assert bf_result == pd_result
+
+
+def test_multiindex_eq_const(scalars_df_index, scalars_pandas_df_index):
+    col_name = ["int64_col", "bool_col"]
+    bf_result = scalars_df_index.set_index(col_name).index == (2, False)
+    pd_result = scalars_pandas_df_index.set_index(col_name).index == (2, False)
+
+    pandas.testing.assert_index_equal(
+        pandas.Index(pd_result, dtype="boolean"), bf_result.to_pandas()
+    )