checkpoint: single-index dev complete. still facing errors when dealing with multi-index

sycai · sycai · commit ed8047bb7fec · 2025-11-05T19:13:54.000Z
diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py
@@ -2345,30 +2345,48 @@ def merge(
         left_post_join_ids = tuple(get_column_left[id] for id in left_join_ids)
         right_post_join_ids = tuple(get_column_right[id] for id in right_join_ids)
 
-        joined_expr, coalesced_ids = coalesce_columns(
-            joined_expr, left_post_join_ids, right_post_join_ids, how=how, drop=False
-        )
+        if left_index or right_index:
+            # For some reason pandas coalesces two joining columns if one side is an index.
+            joined_expr, resolved_join_ids = coalesce_columns(
+                joined_expr, left_post_join_ids, right_post_join_ids
+            )
+        else:
+            joined_expr, resolved_join_ids = resolve_col_join_ids(
+                joined_expr,
+                left_post_join_ids,
+                right_post_join_ids,
+                how=how,
+                drop=False,
+            )
 
         result_columns = []
         matching_join_labels = []
+
         # Select left value columns
         for col_id in self.value_columns:
             if col_id in left_join_ids:
                 key_part = left_join_ids.index(col_id)
                 matching_right_id = right_join_ids[key_part]
-                if self.col_id_to_label[col_id] == other.col_id_to_label.get(
-                    matching_right_id, None
+                if (
+                    right_index
+                    or self.col_id_to_label[col_id]
+                    == other.col_id_to_label[matching_right_id]
                 ):
                     matching_join_labels.append(self.col_id_to_label[col_id])
-                    result_columns.append(coalesced_ids[key_part])
+                    result_columns.append(resolved_join_ids[key_part])
                 else:
                     result_columns.append(get_column_left[col_id])
             else:
                 result_columns.append(get_column_left[col_id])
+
+        # Select right value columns
         for col_id in other.value_columns:
             if col_id in right_join_ids:
                 if other.col_id_to_label[col_id] in matching_join_labels:
                     pass
+                elif left_index:
+                    key_part = right_join_ids.index(col_id)
+                    result_columns.append(resolved_join_ids[key_part])
                 else:
                     result_columns.append(get_column_right[col_id])
             else:
@@ -2379,14 +2397,19 @@ def merge(
             joined_expr = joined_expr.order_by(
                 [
                     ordering.OrderingExpression(ex.deref(col_id))
-                    for col_id in coalesced_ids
+                    for col_id in resolved_join_ids
                 ],
             )
 
         left_idx_id_post_join = [get_column_left[id] for id in self.index_columns]
         right_idx_id_post_join = [get_column_right[id] for id in other.index_columns]
         index_cols = _resolve_index_col(
-            left_idx_id_post_join, right_idx_id_post_join, left_index, right_index, how
+            left_idx_id_post_join,
+            right_idx_id_post_join,
+            resolved_join_ids,
+            left_index,
+            right_index,
+            how,
         )
 
         joined_expr = joined_expr.select_columns(result_columns + index_cols)
@@ -3123,7 +3146,7 @@ def join_mono_indexed(
     left_index = get_column_left[left.index_columns[0]]
     right_index = get_column_right[right.index_columns[0]]
     # Drop original indices from each side. and used the coalesced combination generated by the join.
-    combined_expr, coalesced_join_cols = coalesce_columns(
+    combined_expr, coalesced_join_cols = resolve_col_join_ids(
         combined_expr, [left_index], [right_index], how=how
     )
     if sort:
@@ -3188,7 +3211,7 @@ def join_multi_indexed(
     left_ids_post_join = [get_column_left[id] for id in left_join_ids]
     right_ids_post_join = [get_column_right[id] for id in right_join_ids]
     # Drop original indices from each side. and used the coalesced combination generated by the join.
-    combined_expr, coalesced_join_cols = coalesce_columns(
+    combined_expr, coalesced_join_cols = resolve_col_join_ids(
         combined_expr, left_ids_post_join, right_ids_post_join, how=how
     )
     if sort:
@@ -3231,13 +3254,17 @@ def resolve_label_id(label: Label) -> str:
 
 
 # TODO: Rewrite just to return expressions
-def coalesce_columns(
+def resolve_col_join_ids(
     expr: core.ArrayValue,
     left_ids: typing.Sequence[str],
     right_ids: typing.Sequence[str],
     how: str,
     drop: bool = True,
 ) -> Tuple[core.ArrayValue, Sequence[str]]:
+    """
+    Collapses and selects the joining column IDs, with the assumption that
+    the ids are all belong to value columns.
+    """
     result_ids = []
     for left_id, right_id in zip(left_ids, right_ids):
         if how == "left" or how == "inner" or how == "cross":
@@ -3249,7 +3276,6 @@ def coalesce_columns(
             if drop:
                 expr = expr.drop_columns([left_id])
         elif how == "outer":
-            coalesced_id = guid.generate_guid()
             expr, coalesced_id = expr.project_to_id(
                 ops.coalesce_op.as_expr(left_id, right_id)
             )
@@ -3261,6 +3287,21 @@ def coalesce_columns(
     return expr, result_ids
 
 
+def coalesce_columns(
+    expr: core.ArrayValue,
+    left_ids: typing.Sequence[str],
+    right_ids: typing.Sequence[str],
+) -> tuple[core.ArrayValue, list[str]]:
+    result_ids = []
+    for left_id, right_id in zip(left_ids, right_ids):
+        expr, coalesced_id = expr.project_to_id(
+            ops.coalesce_op.as_expr(left_id, right_id)
+        )
+        result_ids.append(coalesced_id)
+
+    return expr, result_ids
+
+
 def _cast_index(block: Block, dtypes: typing.Sequence[bigframes.dtypes.Dtype]):
     original_block = block
     result_ids = []
@@ -3481,6 +3522,7 @@ def _pd_index_to_array_value(
 def _resolve_index_col(
     left_index_cols: list[str],
     right_index_cols: list[str],
+    resolved_join_ids: list[str],
     left_index: bool,
     right_index: bool,
     how: typing.Literal[
@@ -3497,12 +3539,13 @@ def _resolve_index_col(
         if how == "right":
             return right_index_cols
         if how == "outer":
-            return []
+            return resolved_join_ids
         else:
             return []
     elif left_index and not right_index:
         return right_index_cols
     elif right_index and not left_index:
         return left_index_cols
     else:
+        # Joining with value columns only. Existing indices will be discarded.
         return []
diff --git a/tests/system/small/core/test_reshape.py b/tests/system/small/core/test_reshape.py
@@ -24,16 +24,17 @@
     ("left_on", "right_on", "left_index", "right_index"),
     [
         ("col_a", None, False, True),
-        (None, "col_c", True, False),
+        (None, "col_d", True, False),
         (None, None, True, True),
     ],
 )
+@pytest.mark.parametrize("how", ["inner", "left", "right", "outer"])
 def test_join_with_index(
-    session: session.Session, left_on, right_on, left_index, right_index
+    session: session.Session, left_on, right_on, left_index, right_index, how
 ):
-    df1 = pd.DataFrame({"col_a": [1, 2, 3], "col_b": [2, 3, 4]})
+    df1 = pd.DataFrame({"col_a": [1, 2, 3], "col_b": [2, 3, 4]}, index=[1, 2, 3])
     bf1 = session.read_pandas(df1)
-    df2 = pd.DataFrame({"col_c": [1, 2, 3], "col_d": [2, 3, 4]})
+    df2 = pd.DataFrame({"col_c": [1, 2, 3], "col_d": [2, 3, 4]}, index=[2, 3, 4])
     bf2 = session.read_pandas(df2)
 
     bf_result = merge.merge(
@@ -43,6 +44,7 @@ def test_join_with_index(
         right_on=right_on,
         left_index=left_index,
         right_index=right_index,
+        how=how
     ).to_pandas()
     pd_result = pd.merge(
         df1,
@@ -51,6 +53,7 @@ def test_join_with_index(
         right_on=right_on,
         left_index=left_index,
         right_index=right_index,
+        how=how
     )
 
     pandas.testing.assert_frame_equal(
@@ -66,13 +69,15 @@ def test_join_with_index(
         (None, None, True, True),
     ],
 )
+@pytest.mark.parametrize("how", ["inner", "left", "right", "outer"])
 def test_join_with_multiindex(
-    session: session.Session, left_on, right_on, left_index, right_index
+    session: session.Session, left_on, right_on, left_index, right_index, how
 ):
-    multi_idx = pd.MultiIndex.from_tuples([(1, 2), (2, 3), (3, 4)])
-    df1 = pd.DataFrame({"col_a": [1, 2, 3], "col_b": [2, 3, 4]}, index=multi_idx)
+    multi_idx1 = pd.MultiIndex.from_tuples([(1, 2), (2, 3), (3, 5)])
+    df1 = pd.DataFrame({"col_a": [1, 2, 3], "col_b": [2, 3, 4]}, index=multi_idx1)
     bf1 = session.read_pandas(df1)
-    df2 = pd.DataFrame({"col_c": [1, 2, 3], "col_d": [2, 3, 4]}, index=multi_idx)
+    multi_idx2 = pd.MultiIndex.from_tuples([(1, 2), (2, 3), (3, 2)])
+    df2 = pd.DataFrame({"col_c": [1, 2, 3], "col_d": [2, 3, 4]}, index=multi_idx2)
     bf2 = session.read_pandas(df2)
 
     bf_result = merge.merge(
@@ -82,6 +87,7 @@ def test_join_with_multiindex(
         right_on=right_on,
         left_index=left_index,
         right_index=right_index,
+        how=how
     ).to_pandas()
     pd_result = pd.merge(
         df1,
@@ -90,8 +96,9 @@ def test_join_with_multiindex(
         right_on=right_on,
         left_index=left_index,
         right_index=right_index,
+        how=how
     )
 
     pandas.testing.assert_frame_equal(
-        bf_result, pd_result, check_dtype=False, check_index_type=False
+        bf_result.sort_index(), pd_result.sort_index(), check_dtype=False, check_index_type=False,
     )
