wrap up support for single index

Author: sycai

bigframes/core/blocks.py

@@ -2351,7 +2351,7 @@ def merge(
                 joined_expr, left_post_join_ids, right_post_join_ids
             )
         else:
-            joined_expr, resolved_join_ids = resolve_col_join_ids(
+            joined_expr, resolved_join_ids = resolve_col_join_ids( # type: ignore
                 joined_expr,
                 left_post_join_ids,
                 right_post_join_ids,

bigframes/core/reshape/merge.py

@@ -20,6 +20,7 @@
 
 from typing import Literal, Sequence
 
+from bigframes_vendored import constants
 import bigframes_vendored.pandas.core.reshape.merge as vendored_pandas_merge
 
 from bigframes import dataframe, series
@@ -116,15 +117,24 @@ def _validate_left_right_on(
     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)):
+    if left_on is not None and not isinstance(left_on, Sequence):
         left_on = [left_on]
-    if right_on is not None and not isinstance(right_on, (tuple, list)):
+    if right_on is not None and not isinstance(right_on, Sequence):
         right_on = [right_on]
 
+    if left_index and left.index.nlevels > 1:
+        raise ValueError(
+            f"Joining with multi-level index is not supported. {constants.FEEDBACK_LINK}"
+        )
+    if right_index and right.index.nlevels > 1:
+        raise ValueError(
+            f"Joining with multi-level index is not supported. {constants.FEEDBACK_LINK}"
+        )
+
     # 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)
+            return list(left._block.index_columns), list(right._block.index_columns)
         elif left_index:
             raise ValueError("Must pass right_on or right_index=True")
         elif right_index:
@@ -145,7 +155,7 @@ def _validate_left_right_on(
                 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)
+            return _to_col_ids(left, common_cols.to_list()), _to_col_ids(right, common_cols.to_list())
 
     elif on is not None:
         if left_on is not None or right_on is not None:
@@ -167,7 +177,6 @@ def _validate_left_right_on(
             )
         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(
@@ -183,7 +192,6 @@ def _validate_left_right_on(
             )
         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(
@@ -193,7 +201,7 @@ def _validate_left_right_on(
             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):
+    if len(right_on) != len(left_on): # type: ignore
         raise ValueError("len(right_on) must equal len(left_on)")
 
     return _to_col_ids(left, left_on), _to_col_ids(right, right_on)

tests/system/small/core/test_reshape.py

@@ -44,7 +44,7 @@ def test_join_with_index(
         right_on=right_on,
         left_index=left_index,
         right_index=right_index,
-        how=how
+        how=how,
     ).to_pandas()
     pd_result = pd.merge(
         df1,
@@ -53,14 +53,42 @@ def test_join_with_index(
         right_on=right_on,
         left_index=left_index,
         right_index=right_index,
-        how=how
+        how=how,
     )
 
     pandas.testing.assert_frame_equal(
         bf_result, pd_result, check_dtype=False, check_index_type=False
     )
 
 
+@pytest.mark.parametrize(
+    ("on", "left_on", "right_on", "left_index", "right_index"),
+    [
+        (None, "col_a", None, True, False),
+        (None, None, "col_c", None, True),
+        ("col_a", None, None, True, True),
+    ],
+)
+def test_join_with_index_invalid_index_arg_raise_error(
+    session: session.Session, on, left_on, right_on, left_index, right_index
+):
+    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]}, index=[2, 3, 4])
+    bf2 = session.read_pandas(df2)
+
+    with pytest.raises(ValueError):
+        merge.merge(
+            bf1,
+            bf2,
+            on=on,
+            left_on=left_on,
+            right_on=right_on,
+            left_index=left_index,
+            right_index=right_index,
+        ).to_pandas()
+
+
 @pytest.mark.parametrize(
     ("left_on", "right_on", "left_index", "right_index"),
     [
@@ -70,7 +98,7 @@ def test_join_with_index(
     ],
 )
 @pytest.mark.parametrize("how", ["inner", "left", "right", "outer"])
-def test_join_with_multiindex(
+def test_join_with_multiindex_raises_error(
     session: session.Session, left_on, right_on, left_index, right_index, how
 ):
     multi_idx1 = pd.MultiIndex.from_tuples([(1, 2), (2, 3), (3, 5)])
@@ -80,25 +108,13 @@ def test_join_with_multiindex(
     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(
-        bf1,
-        bf2,
-        left_on=left_on,
-        right_on=right_on,
-        left_index=left_index,
-        right_index=right_index,
-        how=how
-    ).to_pandas()
-    pd_result = pd.merge(
-        df1,
-        df2,
-        left_on=left_on,
-        right_on=right_on,
-        left_index=left_index,
-        right_index=right_index,
-        how=how
-    )
-
-    pandas.testing.assert_frame_equal(
-        bf_result.sort_index(), pd_result.sort_index(), check_dtype=False, check_index_type=False,
-    )
+    with pytest.raises(ValueError):
+        merge.merge(
+            bf1,
+            bf2,
+            left_on=left_on,
+            right_on=right_on,
+            left_index=left_index,
+            right_index=right_index,
+            how=how,
+        )

third_party/bigframes_vendored/pandas/core/frame.py

@@ -4593,6 +4593,8 @@ def merge(
         *,
         left_on: Optional[str] = None,
         right_on: Optional[str] = None,
+        left_index: bool = False,
+        right_index: bool = False,
         sort: bool = False,
         suffixes: tuple[str, str] = ("_x", "_y"),
     ) -> DataFrame:
@@ -4705,6 +4707,10 @@ def merge(
             right_on (label or list of labels):
                 Columns to join on in the right DataFrame. Either on or left_on + right_on
                 must be passed in.
+            left_index (bool, default False):
+                Use the index from the left DataFrame as the join key.
+            right_index (bool, default False):
+                Use the index from the right DataFrame as the join key.
             sort:
                 Default False. Sort the join keys lexicographically in the
                 result DataFrame. If False, the order of the join keys depends

third_party/bigframes_vendored/pandas/core/reshape/merge.py

@@ -13,6 +13,8 @@ def merge(
     *,
     left_on=None,
     right_on=None,
+    left_index: bool = False,
+    right_index: bool = False,
     sort=False,
     suffixes=("_x", "_y"),
 ):
@@ -61,6 +63,10 @@ def merge(
         right_on (label or list of labels):
             Columns to join on in the right DataFrame. Either on or left_on + right_on
             must be passed in.
+        left_index (bool, default False):
+            Use the index from the left DataFrame as the join key.
+        right_index (bool, default False):
+            Use the index from the right DataFrame as the join key.
         sort:
             Default False. Sort the join keys lexicographically in the
             result DataFrame. If False, the order of the join keys depends