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

connector/connect/src/main/scala/org/apache/spark/sql/connect/dsl/package.scala

@@ -64,6 +64,36 @@ package object dsl {
           .build()
     }
 
+    /**
+     * Create an unresolved function from name parts.
+     *
+     * @param nameParts
+     * @param args
+     * @return Expression wrapping the unresolved function.
+     */
+    def fun(nameParts: Seq[String], args: Seq[proto.Expression]): proto.Expression = {
+      proto.Expression.newBuilder().setUnresolvedFunction(
+        proto.Expression.UnresolvedFunction.newBuilder()
+          .addAllParts(nameParts.asJava)
+          .addAllArguments(args.asJava)
+      ).build()
+    }
+
+    /**
+     * Creates an UnresolvedFunction from a single identifier.
+     *
+     * @param name
+     * @param args
+     * @return Expression wrapping the unresolved function.
+     */
+    def callFunction(name: String, args: Seq[proto.Expression]): proto.Expression = {
+      proto.Expression.newBuilder().setUnresolvedFunction(
+        proto.Expression.UnresolvedFunction.newBuilder()
+          .addParts(name)
+          .addAllArguments(args.asJava)
+      ).build()
+    }
+
     implicit def intToLiteral(i: Int): proto.Expression =
       proto.Expression
         .newBuilder()

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.
    *
@@ -192,21 +188,14 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
    * @return
    */
   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)
+    if (fun.getPartsCount == 1 && fun.getParts(0).contains(".")) {
+      throw new IllegalArgumentException(
+        "Function identifier must be passed as sequence of name parts.")
     }
+    UnresolvedFunction(
+      fun.getPartsList.asScala.toSeq,
+      fun.getArgumentsList.asScala.map(transformExpression).toSeq,
+      isDistinct = false)
   }
 
   private def transformAlias(alias: proto.Expression.Alias): Expression = {

connector/connect/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala

@@ -193,7 +193,7 @@ class SparkConnectPlannerSuite extends SparkFunSuite with SparkConnectPlanTest {
 
     val joinCondition = proto.Expression.newBuilder.setUnresolvedFunction(
       proto.Expression.UnresolvedFunction.newBuilder
-        .addAllParts(Seq("eq").asJava)
+        .addAllParts(Seq("==").asJava)
         .addArguments(unresolvedAttribute)
         .addArguments(unresolvedAttribute)
         .build())

connector/connect/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectProtoSuite.scala

@@ -49,6 +49,33 @@ class SparkConnectProtoSuite extends PlanTest with SparkConnectPlanTest {
     comparePlans(connectPlan.analyze, sparkPlan.analyze, false)
   }
 
+  test("UnresolvedFunction resolution.") {
+    {
+      import org.apache.spark.sql.connect.dsl.expressions._
+      import org.apache.spark.sql.connect.dsl.plans._
+      assertThrows[IllegalArgumentException] {
+        transform(connectTestRelation.select(callFunction("default.hex", Seq("id".protoAttr))))
+      }
+    }
+
+    val connectPlan = {
+      import org.apache.spark.sql.connect.dsl.expressions._
+      import org.apache.spark.sql.connect.dsl.plans._
+      transform(connectTestRelation.select(fun(Seq("default", "hex"), Seq("id".protoAttr))))
+    }
+
+    assertThrows[UnsupportedOperationException] {
+      connectPlan.analyze
+    }
+
+    val validPlan = {
+      import org.apache.spark.sql.connect.dsl.expressions._
+      import org.apache.spark.sql.connect.dsl.plans._
+      transform(connectTestRelation.select(fun(Seq("hex"), Seq("id".protoAttr))))
+    }
+    assert(validPlan.analyze != null)
+  }
+
   test("Basic filter") {
     val connectPlan = {
       import org.apache.spark.sql.connect.dsl.expressions._

python/pyspark/sql/connect/column.py

@@ -26,11 +26,51 @@
     import pyspark.sql.connect.proto as proto
 
 
+def _bin_op(
+    name: str, doc: str = "binary function", reverse: bool = False
+) -> Callable[["ColumnRef", Any], "Expression"]:
+    def _(self: "ColumnRef", other: Any) -> "Expression":
+        if isinstance(other, get_args(PrimitiveType)):
+            other = LiteralExpression(other)
+        if not reverse:
+            return ScalarFunctionExpression(name, self, other)
+        else:
+            return ScalarFunctionExpression(name, other, self)
+
+    return _
+
+
 class Expression(object):
     """
     Expression base class.
     """
 
+    __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(">=")
+    __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("==", self, other)
+
     def __init__(self) -> None:
         pass
 
@@ -73,20 +113,6 @@ def __str__(self) -> str:
         return f"Literal({self._value})"
 
 
-def _bin_op(
-    name: str, doc: str = "binary function", reverse: bool = False
-) -> Callable[["ColumnRef", Any], Expression]:
-    def _(self: "ColumnRef", other: Any) -> Expression:
-        if isinstance(other, get_args(PrimitiveType)):
-            other = LiteralExpression(other)
-        if not reverse:
-            return ScalarFunctionExpression(name, self, other)
-        else:
-            return ScalarFunctionExpression(name, other, self)
-
-    return _
-
-
 class ColumnRef(Expression):
     """Represents a column reference. There is no guarantee that this column
     actually exists. In the context of this project, we refer by its name and
@@ -105,32 +131,6 @@ 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)
-    __pow__ = _bin_op("pow")
-    __rpow__ = _bin_op("pow", reverse=True)
-    __ge__ = _bin_op("greterEquals")
-    __le__ = _bin_op("lessEquals")
-
-    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)
-
     def to_plan(self, session: Optional["RemoteSparkSession"]) -> proto.Expression:
         """Returns the Proto representation of the expression."""
         expr = proto.Expression()

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