perf: Simplify merge join key coalescing (#1361)

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -2077,14 +2077,12 @@ def merge(
         result_columns = []
         matching_join_labels = []
 
-        coalesced_ids = []
-        for left_id, right_id in zip(left_join_ids, right_join_ids):
-            joined_expr, coalesced_id = joined_expr.project_to_id(
-                ops.coalesce_op.as_expr(
-                    get_column_left[left_id], get_column_right[right_id]
-                ),
-            )
-            coalesced_ids.append(coalesced_id)
+        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
+        )
 
         for col_id in self.value_columns:
             if col_id in left_join_ids:
@@ -2102,7 +2100,6 @@ def merge(
                 result_columns.append(get_column_left[col_id])
         for col_id in other.value_columns:
             if col_id in right_join_ids:
-                key_part = right_join_ids.index(col_id)
                 if other.col_id_to_label[matching_right_id] in matching_join_labels:
                     pass
                 else:
@@ -2928,26 +2925,31 @@ def resolve_label_id(label: Label) -> str:
     )
 
 
+# TODO: Rewrite just to return expressions
 def coalesce_columns(
     expr: core.ArrayValue,
     left_ids: typing.Sequence[str],
     right_ids: typing.Sequence[str],
     how: str,
+    drop: bool = True,
 ) -> Tuple[core.ArrayValue, Sequence[str]]:
     result_ids = []
     for left_id, right_id in zip(left_ids, right_ids):
         if how == "left" or how == "inner" or how == "cross":
             result_ids.append(left_id)
-            expr = expr.drop_columns([right_id])
+            if drop:
+                expr = expr.drop_columns([right_id])
         elif how == "right":
             result_ids.append(right_id)
-            expr = expr.drop_columns([left_id])
+            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)
             )
-            expr = expr.drop_columns([left_id, right_id])
+            if drop:
+                expr = expr.drop_columns([left_id, right_id])
             result_ids.append(coalesced_id)
         else:
             raise ValueError(f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}")