Skip to content

[SPARK-30127][SQL] Support case class parameter for typed Scala UDF #27937

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Ngone51 wants to merge 11 commits into from
Closed

[SPARK-30127][SQL] Support case class parameter for typed Scala UDF #27937

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
2421c72
init
Ngone51 Mar 16, 2020
dd90298
update
Ngone51 Mar 17, 2020
2b186bd
support inner class
Ngone51 Mar 17, 2020
84e7a7f
fix non-static inner class
Ngone51 Mar 18, 2020
d600cf7
use expressionencoder
Ngone51 Mar 18, 2020
867ad06
fix Any and untyped Scala UDF
Ngone51 Mar 18, 2020
23ca098
address comment
Ngone51 Mar 19, 2020
842d6fa
use encoder only
Ngone51 Mar 19, 2020
174e017
no need resolvedEnc
Ngone51 Mar 23, 2020
8e82f3f
remove applyOption
Ngone51 Mar 23, 2020
b0b298e
update comment
Ngone51 Mar 23, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2707,7 +2707,7 @@ class Analyzer(

case p => p transformExpressionsUp {

case udf @ ScalaUDF(_, _, inputs, inputPrimitives, _, _, _, _)
case udf @ ScalaUDF(_, _, inputs, inputPrimitives, _, _, _, _, _)
if inputPrimitives.contains(true) =>
// Otherwise, add special handling of null for fields that can't accept null.
// The result of operations like this, when passed null, is generally to return null.
Expand Down
527 changes: 270 additions & 257 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
View file
Open in desktop

Large diffs are not rendered by default.

69 changes: 46 additions & 23 deletions sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
View file
Open in desktop

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions sql/core/src/main/scala/org/apache/spark/sql/expressions/UserDefinedFunction.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ private[spark] case class SparkUserDefinedFunction(
f: AnyRef,
dataType: DataType,
inputSchemas: Seq[Option[ScalaReflection.Schema]],

@koertkuipers koertkuipers May 6, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

quite a lot of outside libraries use this constructor (or to be precise the older equivalent version with DataType instead of Schema that used to be in UserDefinedFunction). see for example:
https://github.com/salesforce/TransmogrifAI/blob/master/features/src/main/scala/com/salesforce/op/features/FeatureSparkTypes.scala#L269

can you suggest a workaround? i see no easy way to go from DataType to ExpressionEncoder...

@Ngone51 Ngone51 May 7, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @koertkuipers , the example seems use UserDefinedFunction instead of SparkUserDefinedFunction. And since Spark v3.0, UserDefinedFunction has been changed from case class to abstract class, see SPARK-26216.

cc @cloud-fan

@koertkuipers koertkuipers May 7, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hi @Ngone51 yes changing it to use SparkUserDefinedFunction instead of UserDefinedFunction was easy. but the change from inputSchemas to inputEncoders poses a bit of a challenge.

@cloud-fan cloud-fan May 7, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not go from DataType to ExpressionEncoder, which is impossible. It's from T <: FeatureType : TypeTag to ExpressionEncoder[T]. I think you need to adjust FeatureSparkTypes and support creating ExpressionEncoder. This is the cost of relying on internal APIs.

@koertkuipers koertkuipers May 10, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

while its nice to see case classes supported, i was surprised to find that another feature of Encoder is not supported: Option to indicate a udf argument could be null. for example one could define a udf for (Option[String], Int) => String to explicitly handle the case where the first input argument is null (or should i say its now None).
a better example is (String, Option[Int]) => String to work around the udf behavior that if primitivess are null the udf is not called and the output is also null.

especially when designing generic systems that use udfs this becomes really important. say you write something that does transform (X, Y) => X for generic types X and Y (with typetags). now say Y could be null (the udf could be called after a left join for example where Y could be joined in). the behavior would now change based on the concrete type of Y... for Strings nulls would get passed in to the udf while for Ints the udf would be skipped. i dont think anyone would want to deal with that. so instead you wants to write a udf for (X, Option[Y]) => X in this case.

@koertkuipers koertkuipers May 10, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in my testing i also found the api now to be somewhat inconsistent/confusing. basically sometimes CatalystTypeConverters.createToScalaConverter is used and sometimes ExpressionEncoder.fromRow, depending solely on if the the argument is a top level struct or not. but CatalystTypeConverters.createToScalaConverter and ExpressionEncoder.fromRow behave very differently, leading to inconsistent application.

for example this (contrived) usage works:

case class Person(name: String, age: Option[Int])
Seq((1, Person("john", Some(55))), (2, Person("mary", None))).toDF("id", "person").withColumn("age", udf{ p: Person1 => p.age }.apply(col("person")))

but this does not:

Seq((1, Seq(Person("john", Some(55)))), (2, Seq(Person("mary", None)))).toDF("id", "persons").withColumn("ages", udf{ s: Seq[Person1] => s.map(_.age) }.apply(col("persons")))

and while Option works nicely in Person case class (and also in tuples) Option does not work in a simple Seq:

Seq(Seq(Some(1), None)).toDF.withColumn("value", udf{ s: Seq[Option[Int]] => s.map(_.map(_ + 1)) }.apply(col("value")) )

and Option also does not work for a function argument:

Seq(None, Some(1), None).toDF.withColumn("value", udf{ o: Option[Int] => o.map(_ + 1) }.apply(col("value")))

this inconsistency will be hard to understand. and this inconsistency is not limited to Options. it also applies to many other things. for example tuples inside maps will not work (still have to use Row there) but tuples inside maps will work if its inside a case class. that is hard to explain to a user...

finally let me give some background why i am a little nervous about this change...
spark udfs have been somewhat limited for a long time. no support for case class, tuples, options. so many libraries have worked around that by defining their own udfs on top on SparkUserDefinedFunction. we do this inhouse too. it is easy to do this with type classes thanks to the composability of inputSchemas.
so now you replaced inputSchemas with inputEncoders. but ExpressionEncoder and TypeTags are not composable. i do not see a way for us to build on top of this for our own inhouse udfs. so then the alternative for us is to abandon our inhouse udfs and start using spark's udfs again, which now support case classes and tuples, which is nice! but the inconsistency of the api and lack of support for option makes that currently not viable to me. i realize this is a spark internal api and this is entirely my own problem. but i thought it was worth pointing out because i suspect i am not the only one that has done this. i think this is one of the more typical workarounds people have done using spark (and i am aware of multiple implementations of this workaround).

sorry for the long posts(s)

@cloud-fan cloud-fan May 11, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

inputSchemas is composable but it contains too little information, that's why the Spark UDF was so limited before.

Our goal is to make Spark UDF powerful enough so that people don't need to use internal APIs to build inhouse UDFs. But you are right that the support is not completed. @Ngone51 Can you take a closer look and see how to make it completed?

BTW, if you do need to keep your inhouse UDFs for a while, there is a way to create ExpressionEncoder from Seq[DataType], by calling RowEncoder.apply. It only supports standard Spark external types, i.e. Row, not case class, which is the same as older versions of Spark.

@Ngone51 Ngone51 May 11, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Ngone51 Can you take a closer look and see how to make it completed?

Yea, I'm looking into it.

@koertkuipers koertkuipers May 11, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@cloud-fan thanks i will take a look at RowEncoder

inputEncoders: Seq[ExpressionEncoder[_]] = Nil,

@cloud-fan cloud-fan Mar 18, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we have a simple inputInfos: Seq[Option[(ScalaReflection.Schema, ExpressionEncoder)]]?

@cloud-fan cloud-fan Mar 18, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

or is it possible to get the schema and nullability from the encoder?

name: Option[String] = None,
nullable: Boolean = true,
deterministic: Boolean = true) extends UserDefinedFunction {
Expand All @@ -115,6 +116,7 @@ private[spark] case class SparkUserDefinedFunction(
dataType,
exprs,
inputsPrimitive,
inputEncoders,
inputTypes,
udfName = name,
nullable = nullable,
Expand Down
33 changes: 22 additions & 11 deletions sql/core/src/main/scala/org/apache/spark/sql/functions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4388,7 +4388,8 @@ object functions {
def udf[RT: TypeTag](f: Function0[RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4404,7 +4405,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag](f: Function1[A1, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4420,7 +4422,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag](f: Function2[A1, A2, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: Nil

@cloud-fan cloud-fan Mar 18, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the type is Any, we may fail to create encoder and we should catch it, like what we've done for ScalaReflection.schemaFor above.

val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4436,7 +4439,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag](f: Function3[A1, A2, A3, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4452,7 +4456,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag](f: Function4[A1, A2, A3, A4, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A4])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: ExpressionEncoder[A4]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4468,7 +4473,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag](f: Function5[A1, A2, A3, A4, A5, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A4])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A5])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: ExpressionEncoder[A4]() :: ExpressionEncoder[A5]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4484,7 +4490,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag](f: Function6[A1, A2, A3, A4, A5, A6, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A4])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A5])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A6])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: ExpressionEncoder[A4]() :: ExpressionEncoder[A5]() :: ExpressionEncoder[A6]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4500,7 +4507,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag](f: Function7[A1, A2, A3, A4, A5, A6, A7, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A4])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A5])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A6])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A7])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: ExpressionEncoder[A4]() :: ExpressionEncoder[A5]() :: ExpressionEncoder[A6]() :: ExpressionEncoder[A7]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4516,7 +4524,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag](f: Function8[A1, A2, A3, A4, A5, A6, A7, A8, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A4])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A5])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A6])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A7])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A8])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: ExpressionEncoder[A4]() :: ExpressionEncoder[A5]() :: ExpressionEncoder[A6]() :: ExpressionEncoder[A7]() :: ExpressionEncoder[A8]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4532,7 +4541,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag](f: Function9[A1, A2, A3, A4, A5, A6, A7, A8, A9, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A4])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A5])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A6])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A7])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A8])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A9])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: ExpressionEncoder[A4]() :: ExpressionEncoder[A5]() :: ExpressionEncoder[A6]() :: ExpressionEncoder[A7]() :: ExpressionEncoder[A8]() :: ExpressionEncoder[A9]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand All @@ -4548,7 +4558,8 @@ object functions {
def udf[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag](f: Function10[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, RT]): UserDefinedFunction = {
val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
val inputSchemas = Try(ScalaReflection.schemaFor(typeTag[A1])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A2])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A3])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A4])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A5])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A6])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A7])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A8])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A9])).toOption :: Try(ScalaReflection.schemaFor(typeTag[A10])).toOption :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas)
val inputEncoders: Seq[ExpressionEncoder[_]] = ExpressionEncoder[A1]() :: ExpressionEncoder[A2]() :: ExpressionEncoder[A3]() :: ExpressionEncoder[A4]() :: ExpressionEncoder[A5]() :: ExpressionEncoder[A6]() :: ExpressionEncoder[A7]() :: ExpressionEncoder[A8]() :: ExpressionEncoder[A9]() :: ExpressionEncoder[A10]() :: Nil
val udf = SparkUserDefinedFunction(f, dataType, inputSchemas, inputEncoders)
if (nullable) udf else udf.asNonNullable()
}

Expand Down
29 changes: 28 additions & 1 deletion sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@ import org.apache.spark.sql.test.SQLTestData._
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.QueryExecutionListener


private case class FunctionResult(f1: String, f2: String)

class UDFSuite extends QueryTest with SharedSparkSession {
Expand Down Expand Up @@ -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 ") {

@cloud-fan cloud-fan Mar 18, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we test nested case calss as well?

val f = (d1: TestData) => TestData(d1.key * 2, "copy")
val myUdf = udf(f)
val df = Seq(("data", TestData(50, "d2"))).toDF("col1", "col2")
checkAnswer(df.select(myUdf(Column("col2"))), Row(Row(100, "copy")) :: Nil)
}
}