feat: support left_index and right_index for merge

sycai · sycai · commit 94cace336782 · 2025-10-31T23:30:56.000Z
diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py
@@ -2303,6 +2303,8 @@ def merge(
         right_join_ids: typing.Sequence[str],
         sort: bool,
         suffixes: tuple[str, str] = ("_x", "_y"),
+        left_index: bool = False,
+        right_index: bool = False,
     ) -> Block:
         conditions = tuple(
             (lid, rid) for lid, rid in zip(left_join_ids, right_join_ids)
@@ -2324,9 +2326,8 @@ def merge(
             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[matching_right_id]
+                if self.col_id_to_label[col_id] == other.col_id_to_label.get(
+                    matching_right_id, None
                 ):
                     matching_join_labels.append(self.col_id_to_label[col_id])
                     result_columns.append(coalesced_ids[key_part])
@@ -2371,13 +2372,15 @@ def merge(
             or other.index.is_null
             or self.session._default_index_type == bigframes.enums.DefaultIndexKind.NULL
         ):
-            expr = joined_expr
-            index_columns = []
+            return Block(joined_expr, index_columns=[], column_labels=labels)
+        elif left_index:
+            return Block(joined_expr, index_columns=[left_post_join_ids], column_labels=labels)
+        elif right_index:
+            return Block(joined_expr, index_columns=[right_post_join_ids], column_labels=labels)
         else:
             expr, offset_index_id = joined_expr.promote_offsets()
             index_columns = [offset_index_id]
-
-        return Block(expr, index_columns=index_columns, column_labels=labels)
+            return Block(expr, index_columns=index_columns, column_labels=labels)
 
     def _align_both_axes(
         self, other: Block, how: str
diff --git a/bigframes/core/reshape/merge.py b/bigframes/core/reshape/merge.py
@@ -40,6 +40,8 @@ def merge(
     *,
     left_on: blocks.Label | Sequence[blocks.Label] | None = None,
     right_on: blocks.Label | Sequence[blocks.Label] | None = None,
+    left_index: bool = False,
+    right_index: bool = False,
     sort: bool = False,
     suffixes: tuple[str, str] = ("_x", "_y"),
 ) -> dataframe.DataFrame:
@@ -59,42 +61,25 @@ def merge(
         )
         return dataframe.DataFrame(result_block)
 
-    left_on, right_on = _validate_left_right_on(
-        left, right, on, left_on=left_on, right_on=right_on
+    left_join_ids, right_join_ids = _validate_left_right_on(
+        left,
+        right,
+        on,
+        left_on=left_on,
+        right_on=right_on,
+        left_index=left_index,
+        right_index=right_index,
     )
 
-    if utils.is_list_like(left_on):
-        left_on = list(left_on)  # type: ignore
-    else:
-        left_on = [left_on]
-
-    if utils.is_list_like(right_on):
-        right_on = list(right_on)  # type: ignore
-    else:
-        right_on = [right_on]
-
-    left_join_ids = []
-    for label in left_on:  # type: ignore
-        left_col_id = left._resolve_label_exact(label)
-        # 0 elements already throws an exception
-        if not left_col_id:
-            raise ValueError(f"No column {label} found in self.")
-        left_join_ids.append(left_col_id)
-
-    right_join_ids = []
-    for label in right_on:  # type: ignore
-        right_col_id = right._resolve_label_exact(label)
-        if not right_col_id:
-            raise ValueError(f"No column {label} found in other.")
-        right_join_ids.append(right_col_id)
-
     block = left._block.merge(
         right._block,
         how,
         left_join_ids,
         right_join_ids,
         sort=sort,
         suffixes=suffixes,
+        left_index=left_index,
+        right_index=right_index
     )
     return dataframe.DataFrame(block)
 
@@ -127,30 +112,97 @@ def _validate_left_right_on(
     *,
     left_on: blocks.Label | Sequence[blocks.Label] | None = None,
     right_on: blocks.Label | Sequence[blocks.Label] | None = None,
-):
-    if on is not None:
+    left_index: bool = False,
+    right_index: bool = False,
+) -> tuple[list[str], list[str]]:
+    # Turn left_on and right_on to lists
+    if left_on is not None and not isinstance(left_on, (tuple, list)):
+        left_on = [left_on]
+    if right_on is not None and not isinstance(right_on, (tuple, list)):
+        right_on = [right_on]
+
+    # The following checks are copied from Pandas.
+    if on is None and left_on is None and right_on is None:
+        if left_index and right_index:
+            return list(left._block.index_columns), (right._block.index_columns)
+        elif left_index:
+            raise ValueError("Must pass right_on or right_index=True")
+        elif right_index:
+            raise ValueError("Must pass left_on or left_index=True")
+        else:
+            # use the common columns
+            common_cols = left.columns.intersection(right.columns)
+            if len(common_cols) == 0:
+                raise ValueError(
+                    "No common columns to perform merge on. "
+                    f"Merge options: left_on={left_on}, "
+                    f"right_on={right_on}, "
+                    f"left_index={left_index}, "
+                    f"right_index={right_index}"
+                )
+            if (
+                not left.columns.join(common_cols, how="inner").is_unique
+                or not right.columns.join(common_cols, how="inner").is_unique
+            ):
+                raise ValueError(f"Data columns not unique: {repr(common_cols)}")
+            return _to_col_ids(left, common_cols), _to_col_ids(right, common_cols)
+
+    elif on is not None:
         if left_on is not None or right_on is not None:
             raise ValueError(
-                "Can not pass both `on` and `left_on` + `right_on` params."
+                'Can only pass argument "on" OR "left_on" '
+                'and "right_on", not a combination of both.'
             )
-        return on, on
-
-    if left_on is not None and right_on is not None:
-        return left_on, right_on
-
-    left_cols = left.columns
-    right_cols = right.columns
-    common_cols = left_cols.intersection(right_cols)
-    if len(common_cols) == 0:
-        raise ValueError(
-            "No common columns to perform merge on."
-            f"Merge options: left_on={left_on}, "
-            f"right_on={right_on}, "
-        )
-    if (
-        not left_cols.join(common_cols, how="inner").is_unique
-        or not right_cols.join(common_cols, how="inner").is_unique
-    ):
-        raise ValueError(f"Data columns not unique: {repr(common_cols)}")
+        if left_index or right_index:
+            raise ValueError(
+                'Can only pass argument "on" OR "left_index" '
+                'and "right_index", not a combination of both.'
+            )
+        return _to_col_ids(left, on), _to_col_ids(right, on)
 
-    return common_cols, common_cols
+    elif left_on is not None:
+        if left_index:
+            raise ValueError(
+                'Can only pass argument "left_on" OR "left_index" not both.'
+            )
+        if not right_index and right_on is None:
+            raise ValueError('Must pass "right_on" OR "right_index".')
+        n = len(left_on)
+        if right_index:
+            if len(left_on) != right.index.nlevels:
+                raise ValueError(
+                    "len(left_on) must equal the number "
+                    'of levels in the index of "right"'
+                )
+            return _to_col_ids(left, left_on), list(right._block.index_columns)
+
+    elif right_on is not None:
+        if right_index:
+            raise ValueError(
+                'Can only pass argument "right_on" OR "right_index" not both.'
+            )
+        if not left_index and left_on is None:
+            raise ValueError('Must pass "left_on" OR "left_index".')
+        n = len(right_on)
+        if left_index:
+            if len(right_on) != left.index.nlevels:
+                raise ValueError(
+                    "len(right_on) must equal the number "
+                    'of levels in the index of "left"'
+                )
+            return list(left._block.index_columns), _to_col_ids(right, right_on)
+
+    # The user correctly specified left_on and right_on
+    if len(right_on) != len(left_on):
+        raise ValueError("len(right_on) must equal len(left_on)")
+
+    return _to_col_ids(left, left_on), _to_col_ids(right, right_on)
+
+
+def _to_col_ids(
+    df: dataframe.DataFrame, join_cols: blocks.Label | Sequence[blocks.Label]
+) -> list[str]:
+    if utils.is_list_like(join_cols):
+        return [df._block.resolve_label_exact_or_error(col) for col in join_cols]
+
+    return [df._block.resolve_label_exact_or_error(join_cols)]
diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py
@@ -3650,6 +3650,8 @@ def merge(
         *,
         left_on: Union[blocks.Label, Sequence[blocks.Label], None] = None,
         right_on: Union[blocks.Label, Sequence[blocks.Label], None] = None,
+        left_index: bool = False,
+        right_index: bool = False,
         sort: bool = False,
         suffixes: tuple[str, str] = ("_x", "_y"),
     ) -> DataFrame:
@@ -3662,6 +3664,8 @@ def merge(
             on,
             left_on=left_on,
             right_on=right_on,
+            left_index=left_index,
+            right_index=right_index,
             sort=sort,
             suffixes=suffixes,
         )
diff --git a/tests/system/small/core/test_reshape.py b/tests/system/small/core/test_reshape.py
@@ -0,0 +1,96 @@
+# Copyright 2025 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.
+
+import pandas as pd
+import pandas.testing
+import pytest
+
+from bigframes import session
+from bigframes.core.reshape import merge
+
+
+@pytest.mark.parametrize(
+    ("left_on", "right_on", "left_index", "right_index"),
+    [
+        ("col_a", None, False, True),
+        (None, "col_c", True, False),
+        (None, None, True, True),
+    ],
+)
+def test_join_with_index(
+    session: session.Session, left_on, right_on, left_index, right_index
+):
+    df1 = pd.DataFrame({"col_a": [1, 2, 3], "col_b": [2, 3, 4]})
+    bf1 = session.read_pandas(df1)
+    df2 = pd.DataFrame({"col_c": [1, 2, 3], "col_d": [2, 3, 4]})
+    bf2 = session.read_pandas(df2)
+
+    bf_result = merge.merge(
+        bf1,
+        bf2,
+        left_on=left_on,
+        right_on=right_on,
+        left_index=left_index,
+        right_index=right_index,
+    ).to_pandas()
+    pd_result = pd.merge(
+        df1,
+        df2,
+        left_on=left_on,
+        right_on=right_on,
+        left_index=left_index,
+        right_index=right_index,
+    )
+
+    pandas.testing.assert_frame_equal(
+        bf_result, pd_result, check_dtype=False, check_index_type=False
+    )
+
+@pytest.mark.parametrize(
+    ("left_on", "right_on", "left_index", "right_index"),
+    [
+        (["col_a", "col_b"], None, False, True),
+        (None, ["col_c", "col_d"], True, False),
+        (None, None, True, True),
+    ],
+)
+def test_join_with_multiindex(
+    session: session.Session, left_on, right_on, left_index, right_index
+):
+    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)
+    bf1 = session.read_pandas(df1)
+    df2 = pd.DataFrame({"col_c": [1, 2, 3], "col_d": [2, 3, 4]}, index=multi_idx)
+    bf2 = session.read_pandas(df2)
+
+    bf_result = merge.merge(
+        bf1,
+        bf2,
+        left_on=left_on,
+        right_on=right_on,
+        left_index=left_index,
+        right_index=right_index,
+    ).to_pandas()
+    pd_result = pd.merge(
+        df1,
+        df2,
+        left_on=left_on,
+        right_on=right_on,
+        left_index=left_index,
+        right_index=right_index,
+    )
+
+    pandas.testing.assert_frame_equal(
+        bf_result, pd_result, check_dtype=False, check_index_type=False
+    )
