[SPARK-40915][CONNECT] Improve on in Join in Python client

### What changes were proposed in this pull request?

1. Fix Join's `on` from ANY to concrete types (e.g. str, list of str, column, list of columns)
2. When `on` is str or list of str, it should generate a proto plan with `using_columns`
3. When `on` is column or list of column, it should generate a proto plan with `join_condition`.

### Why are the changes needed?

Improve API coverage

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

UT

Closes #38393 from amaliujia/python_join_on_improvement.

Authored-by: Rui Wang <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>

Author: amaliujia

python/pyspark/sql/connect/dataframe.py

@@ -217,8 +217,13 @@ def groupBy(self, *cols: "ColumnOrString") -> GroupingFrame:
     def head(self, n: int) -> Optional["pandas.DataFrame"]:
         return self.limit(n).toPandas()
 
-    # TODO(martin.grund) fix mypu
-    def join(self, other: "DataFrame", on: Any, how: Optional[str] = None) -> "DataFrame":
+    # TODO: extend `on` to also be type List[ColumnRef].
+    def join(
+        self,
+        other: "DataFrame",
+        on: Optional[Union[str, List[str], ColumnRef]] = None,
+        how: Optional[str] = None,
+    ) -> "DataFrame":
         if self._plan is None:
             raise Exception("Cannot join when self._plan is empty.")
         if other._plan is None:

python/pyspark/sql/connect/plan.py

@@ -537,7 +537,7 @@ def __init__(
         self,
         left: Optional["LogicalPlan"],
         right: "LogicalPlan",
-        on: "ColumnOrString",
+        on: Optional[Union[str, List[str], ColumnRef]],
         how: Optional[str],
     ) -> None:
         super().__init__(left)
@@ -575,7 +575,14 @@ def plan(self, session: Optional["RemoteSparkSession"]) -> proto.Relation:
         rel = proto.Relation()
         rel.join.left.CopyFrom(self.left.plan(session))
         rel.join.right.CopyFrom(self.right.plan(session))
-        rel.join.join_condition.CopyFrom(self.to_attr_or_expression(self.on, session))
+        if self.on is not None:
+            if not isinstance(self.on, list):
+                if isinstance(self.on, str):
+                    rel.join.using_columns.append(self.on)
+                else:
+                    rel.join.join_condition.CopyFrom(self.to_attr_or_expression(self.on, session))
+            else:
+                rel.join.using_columns.extend(self.on)
         rel.join.join_type = self.how
         return rel
 

python/pyspark/sql/tests/connect/test_connect_plan_only.py

@@ -37,6 +37,23 @@ def test_simple_project(self):
         self.assertIsNotNone(plan.root, "Root relation must be set")
         self.assertIsNotNone(plan.root.read)
 
+    def test_join_using_columns(self):
+        left_input = self.connect.readTable(table_name=self.tbl_name)
+        right_input = self.connect.readTable(table_name=self.tbl_name)
+        plan = left_input.join(other=right_input, on="join_column")._plan.to_proto(self.connect)
+        self.assertEqual(len(plan.root.join.using_columns), 1)
+
+        plan2 = left_input.join(other=right_input, on=["col1", "col2"])._plan.to_proto(self.connect)
+        self.assertEqual(len(plan2.root.join.using_columns), 2)
+
+    def test_join_condition(self):
+        left_input = self.connect.readTable(table_name=self.tbl_name)
+        right_input = self.connect.readTable(table_name=self.tbl_name)
+        plan = left_input.join(
+            other=right_input, on=left_input.name == right_input.name
+        )._plan.to_proto(self.connect)
+        self.assertIsNotNone(plan.root.join.join_condition)
+
     def test_filter(self):
         df = self.connect.readTable(table_name=self.tbl_name)
         plan = df.filter(df.col_name > 3)._plan.to_proto(self.connect)