fix multiindex eq, mypy

Author: TrevorBergeron

bigframes/core/indexes/base.py

@@ -754,7 +754,7 @@ 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:
+    def __eq__(self, other) -> Index:  # type: ignore
         return self._apply_binop(other, ops.eq_op)
 
     def _apply_binop(self, other, op: ops.BinaryOp) -> Index:
@@ -802,7 +802,7 @@ def _apply_binop(self, other, op: ops.BinaryOp) -> Index:
                 labels=[None] * self.nlevels,
                 drop=True,
             )
-            return Index(block)
+            return Index(block.set_index(block.value_columns))
         else:
             return NotImplemented
 

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
 
 
@@ -47,13 +49,25 @@ def from_arrays(
         # Index.__new__ should detect multiple levels and properly create a multiindex
         return cast(MultiIndex, Index(pd_index))
 
-    def __eg__(self, other) -> 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)
+        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(
-            eq_result._block.aggregate_all_and_stack(
-                agg_ops.all_op, axis=1, dropna=False
+            blocks.Block(
+                result_expr.select_columns(result_ids),
+                index_columns=result_ids,
+                column_labels=(),
+                index_labels=[None],
             )
         )

tests/system/small/test_multiindex.py

@@ -1481,4 +1481,6 @@ def test_multiindex_eq_const(scalars_df_index, scalars_pandas_df_index):
     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)
 
-    assert bf_result == pd_result
+    pandas.testing.assert_index_equal(
+        pandas.Index(pd_result, dtype="boolean"), bf_result.to_pandas()
+    )