Snowpark UDFs In Scala

1.介绍

当你使用 Snowpark API 创建一个 UDF 时，Snowpark 库会将你的 UDF 代码序列化并上传到一个Internal Stage。当你调用 UDF 时，Snowpark 库会在数据所在的服务器上执行你的函数。因此，数据不需要传输到客户端，函数就能处理数据。

在你的自定义代码中，你还可以调用打包在 JAR 文件中的代码（例如，用于第三方库的 Java 类）。

有以下另种方法创建UDF：

• 你可以创建一个匿名的 UDF，并将该函数赋值给一个变量。只要该变量在作用域内，你就可以使用该变量来调用该 UDF。

// Create and register an anonymous UDF (doubleUdf).
val doubleUdf = udf((x: Int) => x + x)
// Call the anonymous UDF.
val dfWithDoubleNum = df.withColumn("doubleNum", doubleUdf(col("num")))

• 你可以创建一个带有名称的 UDF，并通过名称来调用该 UDF。例如，如果你需要通过名称调用一个 UDF，或者在后续的会话中使用该 UDF，你可以使用这种方式。

// Create and register a permanent named UDF ("doubleUdf").
session.udf.registerPermanent("doubleUdf", (x: Int) => x + x, "mystage")
// Call the named UDF.
val dfWithDoubleNum = df.withColumn("doubleNum", callUDF("doubleUdf", col("num")))

2.UDF中支持的数据类型

image.png

3.（注意）在具有App Trait的对象中创建UDF

Scala提供了一个App trait，你可以扩展它以将你的Scala对象转化为可执行程序。App trait提供了一个main方法，它会自动执行对象定义体中的所有代码。（对象定义体中的代码实际上成为了主方法。）

扩展App trait的一个影响是，直到调用main方法之前，对象中的字段不会被初始化。如果你的对象扩展了App，并且你定义了一个使用之前初始化的对象字段的UDF，上传到服务器的UDF定义将不包括对象字段的初始化值。

例如，假设你在对象中定义并初始化了一个名为myConst的字段，并在一个UDF中使用了该字段：

object Main extends App {
  ...
  // Initialize a field.
  val myConst = "Prefix "
  // Use the field in a UDF.
  // Because the App trait delays the initialization of the object fields,
  // myConst in the UDF definition resolves to null.
  val myUdf = udf((s : String) =>  myConst + s )
  ...
}

当Snowpark将UDF定义序列化并上传到Snowflake时，myConst没有被初始化，因此解析为null。因此，调用UDF会返回myConst的null值。

为了解决这个问题，你需要修改你的对象，不再扩展App trait，并为你的代码实现一个独立的main方法：

object Main {
  ...
  def main(args: Array[String]): Unit = {
    ... // Your code ...
  }
  ...
}

4.为UDF指定依赖

为了通过Snowpark API定义一个UDF，你必须调用Session.addDependency()来添加包含UDF所依赖的类和资源的文件（例如JAR文件、资源文件等）。Snowpark库将这些文件上传到一个Internal Stage，并在执行UDF时将这些文件添加到类路径中。
这样做的目的是确保在执行UDF时能够访问到所依赖的类和资源。

注意：如果你不希望在每次运行应用程序时都上传文件，可以将文件上传到一个阶段（stage）。在调用 addDependency 时，传递文件在阶段中的路径。这样做的好处是，可以避免重复上传文件，提高应用程序的执行效率。通过指定阶段中文件的路径，Snowpark库将能够在执行UDF时访问该文件。

• add jar file in a stage

// Add a JAR file that you uploaded to a stage.
session.addDependency("@my_stage/<path>/my-library.jar")

5.创建匿名UDF

要创建一个匿名的UDF，你可以采取以下两种方法之一：

1. 调用 com.snowflake.snowpark.functions 对象中的 udf 函数，传入匿名函数的定义。
2. 调用 UDFRegistration 类中的 registerTemporary 方法，传入匿名函数的定义。由于你注册的是一个匿名UDF，所以必须使用不带有名称参数的方法签名。

注意：当编写多线程代码（例如使用并行集合）时，建议使用 registerTemporary 方法注册UDF，而不是使用 udf 函数。这可以避免出现找不到默认的Snowflake Session对象的错误。

// Create and register an anonymous UDF.
val doubleUdf = udf((x: Int) => x + x)
// Call the anonymous UDF, passing in the "num" column.
// The example uses withColumn to return a DataFrame containing
// the UDF result in a new column named "doubleNum".
val dfWithDoubleNum = df.withColumn("doubleNum", doubleUdf(col("num")))

以下示例创建了一个使用自定义类（LanguageDetector，用于检测文本中使用的语言）的匿名UDF。该示例调用匿名UDF来检测DataFrame中的text_data列中的语言，并创建一个新的DataFrame，其中包含一个额外的lang列，表示所使用的语言。

// Import the udf function from the functions object.
import com.snowflake.snowpark.functions._

// Import the package for your custom code.
// The custom code in this example detects the language of textual data.
import com.mycompany.LanguageDetector

// If the custom code is packaged in a JAR file, add that JAR file as
// a dependency.
session.addDependency("$HOME/language-detector.jar")

// Create a detector
val detector = new LanguageDetector()

// Create an anonymous UDF that takes a string of text and returns the language used in that string.
// Note that this captures the detector object created above.
// Assign the UDF to the langUdf variable, which will be used to call the UDF.
val langUdf = udf((s: String) =>
     Option(detector.detect(s)).getOrElse("UNKNOWN"))

// Create a new DataFrame that contains an additional "lang" column that contains the language
// detected by the UDF.
val dfEmailsWithLangCol =
    dfEmails.withColumn("lang", langUdf(col("text_data")))

6.创建具名UDF

如果你想通过名称调用一个UDF（例如使用functions对象中的callUDF函数），或者如果你需要在后续会话中使用一个UDF，你可以创建和注册一个命名的UDF。为此，可以使用UDFRegistration类中的以下方法之一：

• registerTemporary：如果你只计划在当前会话中使用该UDF。
• registerPermanent：如果你计划在后续会话中使用该UDF。

要访问UDFRegistration类的对象，调用Session类的udf方法。
registerTemporary方法创建了一个临时的UDF，你可以在当前会话中使用它。

// Create and register a temporary named UDF.
session.udf.registerTemporary("doubleUdf", (x: Int) => x + x)
// Call the named UDF, passing in the "num" column.
// The example uses withColumn to return a DataFrame containing
// the UDF result in a new column named "doubleNum".
val dfWithDoubleNum = df.withColumn("doubleNum", callUDF("doubleUdf", col("num")))

7.使用不可序列化的对象

当你为lambda或函数创建UDF时，Snowpark库会对lambda闭包进行序列化并将其发送到服务器执行。
如果被lambda闭包捕获的对象不可序列化，Snowpark库将抛出java.io.NotSerializableException异常。

如果出现这种情况，你可以采取以下两种方法之一：

1. 使对象可序列化。
2. 将对象声明为lazy val或使用@transient注解，以避免对对象进行序列化。
   例如：

// Declare the detector object as lazy.
lazy val detector = new LanguageDetector("en")
// The detector object is not serialized but is instead reconstructed on the server.
val langUdf = udf((s: String) =>
     Option(detector.detect(s)).getOrElse("UNKNOWN"))

8.编写UDF的初始化代码

如果你的UDF需要初始化代码或上下文，你可以通过UDF闭包中捕获的值来提供这些内容。

下面的示例使用一个单独的类来初始化三个UDF所需的上下文。

• 第一个UDF在lambda内部创建了该类的一个新实例，因此每次调用UDF时都会执行初始化。
• 第二个UDF捕获了在客户端程序中生成的该类的实例。在客户端生成的上下文被序列化，并被UDF使用。请注意，为了使该方法有效，上下文类必须是可序列化的。
• 第三个UDF捕获了一个lazy val，因此上下文在第一次UDF调用时被惰性实例化，并在后续调用中被重用。即使上下文不可序列化，这种方法也有效。但是，并不能保证数据帧中的所有UDF调用都使用同一个惰性生成的上下文。

import com.snowflake.snowpark._
import com.snowflake.snowpark.functions._
import scala.util.Random

// Context needed for a UDF.
class Context {
  val randomInt = Random.nextInt
}

// Serializable context needed for the UDF.
class SerContext extends Serializable {
  val randomInt = Random.nextInt
}

object TestUdf {
  def main(args: Array[String]): Unit = {
    // Create the session.
    val session = Session.builder.configFile("/<path>/profile.properties").create
    import session.implicits._
    session.range(1, 10, 2).show()

    // Create a DataFrame with two columns ("c" and "d").
    val dummy = session.createDataFrame(Seq((1, 1), (2, 2), (3, 3))).toDF("c", "d")
    dummy.show()

    // Initialize the context once per invocation.
    val udfRepeatedInit = udf((i: Int) => (new Context).randomInt)
    dummy.select(udfRepeatedInit('c)).show()

    // Initialize the serializable context only once,
    // regardless of the number of times that the UDF is invoked.
    val sC = new SerContext
    val udfOnceInit = udf((i: Int) => sC.randomInt)
    dummy.select(udfOnceInit('c)).show()

    // Initialize the non-serializable context only once,
    // regardless of the number of times that the UDF is invoked.
    lazy val unserC = new Context
    val udfOnceInitU = udf((i: Int) => unserC.randomInt)
    dummy.select(udfOnceInitU('c)).show()
  }
}

9.在UDF中读取文件

正如前面提到的，Snowpark库会将UDF上传并在服务器上执行。如果你的UDF需要从文件中读取数据，你必须确保文件与UDF一起上传。

另外，如果文件的内容在调用UDF之间保持不变，你可以编写代码，在第一次调用时加载文件，并在后续调用中不再加载。这可以提高UDF调用的性能。

example:

1.Add the file to a jar file
# Create a new JAR file containing data/hello.txt.
$ jar cvf <path>/myJar.jar data/hello.txt

2.Specify that the JAR file is a dependency
// Specify that myJar.jar contains files that your UDF depends on.
session.addDependency("<path>/myJar.jar")

3.In the UDF, call Class.getResourceAsStream to find the file in the classpath and read the file.
// Read data/hello.txt from myJar.jar.
val resourceName = "/data/hello.txt"
val inputStream = classOf[com.snowflake.snowpark.DataFrame].getResourceAsStream(resourceName)

注意：如果你不希望文件的内容在UDF调用之间发生变化，可以将文件读取操作放在一个lazy val中。这将导致文件加载代码仅在第一次调用UDF时执行，而不会在后续调用中执行。

下面的示例定义了一个对象（UDFCode），其中包含一个函数作为UDF（readFileFunc）使用。该函数读取文件data/hello.txt，该文件预计包含字符串"hello"。该函数将此字符串添加到作为参数传入的字符串前面。

// Create a function object that reads a file.
object UDFCode extends Serializable {

  // The code in this block reads the file. To prevent this code from executing each time that the UDF is called,
  // the code is used in the definition of a lazy val. The code for a lazy val is executed only once when the variable is
  // first accessed.
  lazy val prefix = {
    import java.io._
    val resourceName = "/data/hello.txt"
    val inputStream = classOf[com.snowflake.snowpark.DataFrame]
      .getResourceAsStream(resourceName)
    if (inputStream == null) {
      throw new Exception("Can't find file " + resourceName)
    }
    scala.io.Source.fromInputStream(inputStream).mkString
  }

  val readFileFunc = (s: String) => prefix + " : " + s
}


// Add the JAR file as a dependency.
session.addDependency("<path>/myJar.jar")

// Create a new DataFrame with one column (NAME)
// that contains the name "Raymond".
val myDf = session.sql("select 'Raymond' NAME")

// Register the function that you defined earlier as an anonymous UDF.
val readFileUdf = udf(UDFCode.readFileFunc)

// Call UDF for the values in the NAME column of the DataFrame.
myDf.withColumn("CONCAT", readFileUdf(col("NAME"))).show()

10.创建UDTF(User-Defined Table Functions)

要在Snowpark中创建和注册UDTF，你需要执行以下步骤：
1.为UDTF定义一个类
2.创建UDTF实例并注册

定义UDTF类

定义一个类，该类继承自com.snowflake.snowpark.udtf包中的UDTFn类之一（例如UDTF0、UDTF1等），其中n表示您的UDTF的输入参数数量。例如，如果您的UDTF传入2个输入参数，请继承UDTF2类。

在您的类中，重写以下方法：

• outputSchema()：返回一个types.StructType对象，描述返回行中字段的名称和类型（即输出的“模式”）。
• process()：对输入分区中的每一行调用一次此方法（请参阅下面的注意事项）。
• endPartition()：在所有行都传递给process()后，对每个分区调用一次此方法。

当调用UDTF时，行在传递给UDTF之前被分组到分区中：

• 如果调用UDTF的语句指定了PARTITION子句（显式分区），则该子句确定如何对行进行分区。
• 如果语句没有指定PARTITION子句（隐式分区），Snowflake将确定最佳的行分区方式。

Table Functions and Partitions

在将行传递给表函数之前，这些行可以被分组到分区中。分区有两个主要优势：

1. 分区允许Snowflake将工作负载分割成多个部分，以提高并行性和性能。
2. 分区允许Snowflake将具有相同特征的所有行作为一个组进行处理。您可以基于组中的所有行返回结果，而不仅仅是针对单个行。

例如，您可以将股票价格数据按照每只股票分组为一个组。可以将同一公司的所有股票价格一起分析，同时可以独立地分析每个公司的股票价格，而不受其他公司的影响。

数据可以显式分区或隐式分区。

显式分区

partitioning into multiple groups

SELECT *
    FROM stocks_table AS st,
         TABLE(my_udtf(st.symbol, st.transaction_date, st.price) OVER (PARTITION BY st.symbol))

partitioning into single group

SELECT *
    FROM stocks_table AS st,
         TABLE(my_udtf(st.symbol, st.transaction_date, st.price) OVER (PARTITION BY 1))

隐式分区

如果表函数没有使用PARTITION BY子句显式地对行进行分区，那么Snowflake通常会隐式地对行进行分区，以利用并行处理来提高性能。

分区的数量通常基于诸如处理函数的仓库大小和输入关系的基数等因素。行通常根据行的物理位置（例如，按micro-partition）分配到特定的分区，因此分区分组没有意义。

重写outputSchema方法

def outputSchema(): StructType

例如，如果您的UDTF返回一个只有一个整数字段的行：

override def outputSchema(): StructType = StructType(StructField("C1", IntegerType))

重写process方法

def process(arg0: A0, ... arg<n> A<n>): Iterable[Row]

此方法对输入分区中的每一行调用一次。

在process()方法中，构建并返回一个包含要由UDTF返回的数据的Row对象的Iterable。行中的字段必须使用您在outputSchema方法中指定的类型。
例如，如果您的UDTF生成行，则为生成的行构建并返回一个Row对象的Iterable：

override def process(start: Int, count: Int): Iterable[Row] =
    (start until (start + count)).map(Row(_))

重写endPartition方法

覆盖endPartition方法并添加在所有输入分区的行都已传递给process方法后执行的代码。endPartition方法对于每个输入分区都会被调用一次。

def endPartition(): Iterable[Row]

如果您在每个分区的末尾不需要返回额外的行，请返回一个空的 Row 对象的可迭代对象。例如：

override def endPartition(): Iterable[Row] = Array.empty[Row]

UDTF example

class MyRangeUdtf extends UDTF2[Int, Int] {
  override def process(start: Int, count: Int): Iterable[Row] =
    (start until (start + count)).map(Row(_))
  override def endPartition(): Iterable[Row] = Array.empty[Row]
  override def outputSchema(): StructType = StructType(StructField("C1", IntegerType))
}

注册UDTF

接下来，创建一个新类的实例，并通过调用 UDTFRegistration 的方法来注册该类。您可以注册一个临时的或永久的 UDTF。

注册Temporary UDTF

// Register the MyRangeUdtf class that was defined in the previous example.
val tableFunction = session.udtf.registerTemporary(new MyRangeUdtf())
// Use the returned TableFunction object to call the UDTF.
session.tableFunction(tableFunction, lit(10), lit(5)).show

// Register the MyRangeUdtf class that was defined in the previous example.
val tableFunction = session.udtf.registerTemporary("myUdtf", new MyRangeUdtf())
// Call the UDTF by name.
session.tableFunction(TableFunction("myUdtf"), lit(10), lit(5)).show

注册Permanent UDTF

// Register the MyRangeUdtf class that was defined in the previous example.
val tableFunction = session.udtf.registerPermanent("myUdtf", new MyRangeUdtf(), "@mystage")
// Call the UDTF by name.
session.tableFunction(TableFunction("myUdtf"), lit(10), lit(5)).show