cloud-fan commented on a change in pull request #27937: [WIP][SPARK-30127][SQL]
Support case class parameter for typed Scala UDF
URL: https://github.com/apache/spark/pull/27937#discussion_r394353700
##########
File path: sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
##########
@@ -551,4 +550,32 @@ class UDFSuite extends QueryTest with SharedSparkSession {
}
assert(e.getMessage.contains("Invalid arguments for function cast"))
}
+
+ test("only one case class parameter") {
+ val f = (d: TestData) => d.key * d.value.toInt
+ val myUdf = udf(f)
+ val df = Seq(("data", TestData(50, "2"))).toDF("col1", "col2")
+ checkAnswer(df.select(myUdf(Column("col2"))), Row(100) :: Nil)
+ }
+
+ test("one case class with primitive parameter") {
+ val f = (i: Int, p: TestData) => p.key * i
+ val myUdf = udf(f)
+ val df = Seq((2, TestData(50, "data"))).toDF("col1", "col2")
+ checkAnswer(df.select(myUdf(Column("col1"), Column("col2"))), Row(100) ::
Nil)
+ }
+
+ test("multiple case class parameters") {
+ val f = (d1: TestData, d2: TestData) => d1.key * d2.key
+ val myUdf = udf(f)
+ val df = Seq((TestData(10, "d1"), TestData(50, "d2"))).toDF("col1", "col2")
+ checkAnswer(df.select(myUdf(Column("col1"), Column("col2"))), Row(500) ::
Nil)
+ }
+
+ test("input case class parameter and return case class ") {
Review comment:
can we test nested case calss as well?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
