当输出是要从Pyspark使用的复杂类型(使用StructType和StructField)时，如何传递Scala UserDefinedFunction

Question

因此，我想创建一个可以在Pyspark中使用的scala UDF。我想要的是接受一个字符串列表作为x，一个字符串列表作为y，并获得所有的字符串组合

因此，如果我有x= "a"，"B“和y="A"，"b”，我期望输出= [a，A，a，B，b，A，b，B]

我编写的Scala代码非常简单

(x: Seq[String], y: Seq[String]) => {for (a <- x; b <-y) yield (a,b)}

我已经创建了一个scala UDF来做这件事。它适用于Scala Spark。

我的问题是试图让这个在pyspark中可调用。

为了做到这一点，我这样做了：

import org.apache.spark.sql.types._
import org.apache.spark.sql.expressions.UserDefinedFunction

import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.api.java.UDF2
import org.apache.spark.sql.api.java.UDF1

class DualArrayExplode extends UDF2[Seq[String], Seq[String], UserDefinedFunction] {
  override  def call(x: Seq[String], y: Seq[String]):UserDefinedFunction = {
    // (worker node stuff)
    
  val DualArrayExplode =  (x: Seq[String], y: Seq[String]) => {for (a <- x; b <-y) yield (a,b)}
  val DualArrayExplodeUDF = (udf(DualArrayExplode))

  return DualArrayExplodeUDF

  }
}

object DualArrayExplode {
  def apply(): DualArrayExplode = {
    new DualArrayExplode()
  }
}

我已经创建了一个jar，其中包含此代码以及其他函数(工作正常)此代码编译时没有任何问题。

当我在scala spark中使用它时，输出列类型是Array(ArrayType(StructType(StructField(_1,StringType,true), StructField(_2,StringType,true)),true))

我的问题是我不能让它在Pyspark上工作。注册此函数时，我无法定义正确的返回类型。

下面是我尝试注册UDF的方法

spark.udf.registerJavaFunction('DualArrayExplode', 
                               'blah.blah.blah.blah.blah.DualArrayExplode', <WHAT_TYPE_HERE???>)

返回类型是可选的，但是如果我省略了它，那么结果是[] (一个空列表)

所以..。我如何在pyspark中实际使用这个scala UDF？

我意识到可能有很多地方出了问题，因此我尽可能好地描述了整个设置。

Answer 1

DualArrayExplode的声明是

class DualArrayExplode extends UDF2[Seq[String], Seq[String], UserDefinedFunction]

这意味着声明了一个udf，它接受两个字符串序列作为输入并返回一个udf。应将其更改为

class DualArrayExplode extends UDF2[Seq[String], Seq[String], Seq[(String,String)]] {
  override  def call(x: Seq[String], y: Seq[String]): Seq[(String,String)]= {
    // (worker node stuff)
    for (a <- x; b <-y) yield (a,b)
  }
}

udf的返回类型已更改为字符串元组序列。

该udf现在可以在Pyspark中注册为

from pyspark.sql import types as T
rt = T.ArrayType(T.StructType([T.StructField("_1",T.StringType()), 
                               T.StructField("_2",T.StringType())]))
spark.udf.registerJavaFunction(name='DualArrayExplode', 
            javaClassName='blah.blah.DualArrayExplode', returnType=rt)