How can I pass a Scala UserDefinedFunction where output is a complex type (using StructType and StructField) to be used from Pyspark

user_2081152 0
Viewed 356

So, I would like to create a scala UDF that can be used in Pyspark. What I want is to accept a list of strings as x and a list of strings as y and get all the string combinations

so if I have x = ["a","b] and y=["A","B"] I expect the output to be out = [[a,A],[a,B],[b,A],[b,B]]

The Scala code that I have managed to whip up that works is simple

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

I have created a scala UDF that does that. It works on Scala Spark.

My problem I have is trying to make this callable from pyspark.

In order to do that I have done this :

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()
  }
}

I have created a jar including this code together with other functions (that work with no problem) This code compiles with no issue.

The output column type when i use this in scala spark is Array(ArrayType(StructType(StructField(_1,StringType,true), StructField(_2,StringType,true)),true))

My issue is that I cannot get this to work with Pyspark. I cannot define a correct return type when I register this function.

here is how I try to register the UDF

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

Return type is optional but if I omit it then the result is [] (an empty list)

So... how can I actually use this scala UDF in pyspark?

I realise that there could be many things going wrong hence tried to describe the whole setup as well as I could.

pyspark apache-spark user-defined-functions interop scala
1 Answers

The declaration of DualArrayExplode is

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

This means that a udf is declared that takes two string sequences as input and returns an udf. This should be changed into

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)
  }
}

The return type of the udf has been changed to a sequence of string tuples.

This udf can now be registered in Pyspark with

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)
user_21298010
Related