[SPARK-40538] [CONNECT] Improve built-in function support for Python client.

connector/connect/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -179,10 +179,6 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
       limitExpr = expressions.Literal(limit.getLimit, IntegerType))
   }
 
-  private def lookupFunction(name: String, args: Seq[Expression]): Expression = {
-    UnresolvedFunction(Seq(name), args, isDistinct = false)
-  }
-
   /**
    * Translates a scalar function from proto to the Catalyst expression.
    *
@@ -193,20 +189,10 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
    */
   private def transformScalarFunction(fun: proto.Expression.UnresolvedFunction): Expression = {
     val funName = fun.getPartsList.asScala.mkString(".")
-    funName match {
-      case "gt" =>
-        assert(fun.getArgumentsCount == 2, "`gt` function must have two arguments.")
-        expressions.GreaterThan(
-          transformExpression(fun.getArguments(0)),
-          transformExpression(fun.getArguments(1)))
-      case "eq" =>
-        assert(fun.getArgumentsCount == 2, "`eq` function must have two arguments.")
-        expressions.EqualTo(
-          transformExpression(fun.getArguments(0)),
-          transformExpression(fun.getArguments(1)))
-      case _ =>
-        lookupFunction(funName, fun.getArgumentsList.asScala.map(transformExpression).toSeq)
-    }
+    UnresolvedFunction(
+      Seq(funName),
+      fun.getArgumentsList.asScala.map(transformExpression).toSeq,
+      isDistinct = false)
   }
 
   private def transformAlias(alias: proto.Expression.Alias): Expression = {

python/pyspark/sql/connect/column.py

@@ -105,31 +105,31 @@ def name(self) -> str:
         """Returns the qualified name of the column reference."""
         return ".".join(self._parts)
 
-    __gt__ = _bin_op("gt")
-    __lt__ = _bin_op("lt")
-    __add__ = _bin_op("plus")
-    __sub__ = _bin_op("minus")
-    __mul__ = _bin_op("multiply")
-    __div__ = _bin_op("divide")
-    __truediv__ = _bin_op("divide")
-    __mod__ = _bin_op("modulo")
-    __radd__ = _bin_op("plus", reverse=True)
-    __rsub__ = _bin_op("minus", reverse=True)
-    __rmul__ = _bin_op("multiply", reverse=True)
-    __rdiv__ = _bin_op("divide", reverse=True)
-    __rtruediv__ = _bin_op("divide", reverse=True)
+    __gt__ = _bin_op(">")
+    __lt__ = _bin_op(">")
+    __add__ = _bin_op("+")
+    __sub__ = _bin_op("-")
+    __mul__ = _bin_op("*")
+    __div__ = _bin_op("/")
+    __truediv__ = _bin_op("/")
+    __mod__ = _bin_op("%")
+    __radd__ = _bin_op("+", reverse=True)
+    __rsub__ = _bin_op("-", reverse=True)
+    __rmul__ = _bin_op("*", reverse=True)
+    __rdiv__ = _bin_op("/", reverse=True)
+    __rtruediv__ = _bin_op("/", reverse=True)
     __pow__ = _bin_op("pow")
     __rpow__ = _bin_op("pow", reverse=True)
-    __ge__ = _bin_op("greterEquals")
-    __le__ = _bin_op("lessEquals")
+    __ge__ = _bin_op(">=")
+    __le__ = _bin_op("<=")
 
     def __eq__(self, other: Any) -> Expression:  # type: ignore[override]
         """Returns a binary expression with the current column as the left
         side and the other expression as the right side.
         """
         if isinstance(other, get_args(PrimitiveType)):
             other = LiteralExpression(other)
-        return ScalarFunctionExpression("eq", self, other)
+        return ScalarFunctionExpression("==", self, other)
 
     def to_plan(self, session: Optional["RemoteSparkSession"]) -> proto.Expression:
         """Returns the Proto representation of the expression."""
@@ -161,7 +161,7 @@ def to_plan(self, session: Optional["RemoteSparkSession"]) -> proto.Expression:
         return self.ref.to_plan(session)
 
 
-class ScalarFunctionExpression(Expression):
+class ScalarFunctionExpression(ColumnRef):
     def __init__(
         self,
         op: str,

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

@@ -19,9 +19,12 @@
 import unittest
 import tempfile
 
+import pandas
+
 from pyspark.sql import SparkSession, Row
 from pyspark.sql.connect.client import RemoteSparkSession
 from pyspark.sql.connect.function_builder import udf
+from pyspark.sql.connect.functions import lit
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
 from pyspark.testing.utils import ReusedPySparkTestCase
 
@@ -79,6 +82,15 @@ def test_simple_explain_string(self):
         result = df.explain()
         self.assertGreater(len(result), 0)
 
+    def test_simple_binary_expressions(self):
+        """Test complex expression"""
+        df = self.connect.read.table(self.tbl_name)
+        pd = df.select(df.id).where(df.id % lit(30) == lit(0)).sort(df.id.asc()).toPandas()
+        self.assertEqual(len(pd.index), 4)
+
+        res = pandas.DataFrame(data={"id": [0, 30, 60, 90]})
+        self.assert_(pd.equals(res), f"{pd.to_string()} != {res.to_string()}")
+
 
 if __name__ == "__main__":
     from pyspark.sql.tests.connect.test_connect_basic import *  # noqa: F401

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

@@ -16,6 +16,7 @@
 #
 
 from pyspark.testing.connectutils import PlanOnlyTestFixture
+from pyspark.sql.connect.proto import Expression as ProtoExpression
 import pyspark.sql.connect as c
 import pyspark.sql.connect.plan as p
 import pyspark.sql.connect.column as col
@@ -51,6 +52,34 @@ def test_column_literals(self):
         plan = fun.lit(10).to_plan(None)
         self.assertIs(plan.literal.i32, 10)
 
+    def test_column_expressions(self):
+        """Test a more complex combination of expressions and their translation into
+        the protobuf structure."""
+        df = c.DataFrame.withPlan(p.Read("table"))
+
+        expr = df.id % fun.lit(10) == fun.lit(10)
+        expr_plan = expr.to_plan(None)
+        self.assertIsNotNone(expr_plan.unresolved_function)
+        self.assertEqual(expr_plan.unresolved_function.parts[0], "==")
+
+        lit_fun = expr_plan.unresolved_function.arguments[1]
+        self.assertIsInstance(lit_fun, ProtoExpression)
+        self.assertIsInstance(lit_fun.literal, ProtoExpression.Literal)
+        self.assertEqual(lit_fun.literal.i32, 10)
+
+        mod_fun = expr_plan.unresolved_function.arguments[0]
+        self.assertIsInstance(mod_fun, ProtoExpression)
+        self.assertIsInstance(mod_fun.unresolved_function, ProtoExpression.UnresolvedFunction)
+        self.assertEqual(len(mod_fun.unresolved_function.arguments), 2)
+        self.assertIsInstance(mod_fun.unresolved_function.arguments[0], ProtoExpression)
+        self.assertIsInstance(
+            mod_fun.unresolved_function.arguments[0].unresolved_attribute,
+            ProtoExpression.UnresolvedAttribute,
+        )
+        self.assertEqual(
+            mod_fun.unresolved_function.arguments[0].unresolved_attribute.parts, ["id"]
+        )
+
 
 if __name__ == "__main__":
     import unittest