Snowpark中使用TableFunction如何实现partition by窗口分区？

功能支持说明 #

Snowpark 1.6.0及以上版本已原生支持UDTF和OVER分区逻辑的搭配使用，更低版本可以通过直接执行SQL的方式实现需求。

正确Scala实现代码 #

import com.snowflake.snowpark._
import com.snowflake.snowpark.functions._

val session = Session.builder.configs(configs).create
val df = session.table("CUSTOMER")
// 定义按name分区的窗口规则
val windowSpec = Window.partitionBy(col("name"))
// 加载自定义UDTF
val mapCountUdtf = tableFunction("map_count")
// 调用UDTF并绑定分区窗口，取返回结果
val result = df.select(
  callTableFunction(mapCountUdtf, col("name")).over(windowSpec).as("mcount")
).select(col("mcount")("RESULT").as("result"))

result.show()

上述代码逻辑和你提供的原生SQL完全对齐：通过over方法将分区规则传入UDTF调用逻辑，最后按UDTF定义的输出列名取对应结果即可。

低版本兼容方案 #

如果使用的Snowpark版本低于1.6.0，可以直接执行原生SQL实现相同逻辑：

val result = session.sql("""
select mcount.result from CUSTOMER, table(map_count(name) over (partition by name)) mcount
""")
result.show()

注意事项 #

• 当前Snowpark UDTF的OVER窗口仅支持partition by配置，暂不支持order by和窗口帧定义
• 代码中取值用的RESULT需要和你编写的JavaScript UDTF定义的输出列名保持一致，若你的UDTF输出列名不同需要对应修改

内容的提问来源于stack exchange，提问作者Sella