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
©著作权归作者所有,转载或内容合作请联系作者
  • 人面猴
    序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
    沈念sama阅读 203,271评论 5赞 476
  • 死咒
    序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
    沈念sama阅读 85,275评论 2赞 380
  • 救了他两次的神仙让他今天三更去死
    文/潘晓璐 我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
    开封第一讲书人阅读 150,151评论 0赞 336
  • 道士缉凶录：失踪的卖姜人
    文/不坏的土叔 我叫张陵，是天一观的道长。 经常有香客问我，道长，这世上最难降的妖魔是什么？ 我笑而不...
    开封第一讲书人阅读 54,550评论 1赞 273
  • 港岛之恋（遗憾婚礼）
    正文 为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
    茶点故事阅读 63,553评论 5赞 365
  • 恶毒庶女顶嫁案：这布局不是一般人想出来的
    文/花漫 我一把揭开白布。 她就那样静静地躺着，像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上，一...
    开封第一讲书人阅读 48,559评论 1赞 281
  • 城市分裂传说
    那天，我揣着相机与录音，去河边找鬼。 笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。 我是一名探鬼主播，决...
    沈念sama阅读 37,924评论 3赞 395
  • 双鸳鸯连环套：你想象不到人心有多黑
    文/苍兰香墨 我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
    开封第一讲书人阅读 36,580评论 0赞 257
  • 万荣杀人案实录
    序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
    沈念sama阅读 40,826评论 1赞 297
  • 护林员之死
    正文 独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,578评论 2赞 320
  • 白月光启示录
    正文 我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,661评论 1赞 329
  • 活死人
    序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
    沈念sama阅读 33,363评论 4赞 318
  • 日本核电站爆炸内幕
    正文 年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
    茶点故事阅读 38,940评论 3赞 307
  • 男人毒药：我在死后第九天来索命
    文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,926评论 0赞 19
  • 一桩弑父案，背后竟有这般阴谋
    文/苍兰香墨 我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,156评论 1赞 259
  • 情欲美人皮
    我被黑心中介骗来泰国打工， 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。 一个月前我还...
    沈念sama阅读 42,872评论 2赞 349
  • 代替公主和亲
    正文 我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
    茶点故事阅读 42,391评论 2赞 342

推荐阅读更多精彩内容

  • 面试题1
    一.Hadoop 1.hdfs写流程 2.hdfs读流程 3.hdfs的体系结构 4.一个datanode 宕机,...
    qydong阅读 2,304评论 0赞 0
  • hive积累大全
    此篇内容：hive自定义函数UDF、UDTF，压缩存储方式，hive优化、hive实际编程基本涵盖了hive基础及...
    大数据阶梯之路阅读 802评论 0赞 4
  • 2021-02-08b
    30.域的含义 域英文叫DOMAIN——域(Domain)是Windows网络中独立运行的单位，域之间相互访问则需...
    saluch阅读 104评论 0赞 0
  • Flink源码阅读（13）--- Flink 不同版本新特性整理
    一. Flink 1.9 新特性 总的变更：https://issues.apache.org/jira/secu...
    sj_91d7阅读 2,738评论 0赞 2
  • Hive基础
    一、Hive基本概念 Hive 是基于 Hadoop 的一个数据仓库工具，可以将结构化的数据文件映射为一张表，并提...
    CJ21阅读 1,729评论 0赞 13