如何在Snowpark DataFrame中调用Python UDF并打印结果？

问题根源 #

1. 你写的是标量UDF，但要实现分组聚合逻辑，必须定义成聚合UDF——标量UDF仅能处理单行数据，无法处理分组后的批量数据集合。
2. Snowpark DataFrame的groupBy后没有function方法，这是Pandas的专属用法，Snowpark中需用agg()执行聚合操作。
3. UDF参数类型错误：聚合UDF接收的是分组后的整列数据集合，不是单个float值。

修正后的代码 #

from snowflake.snowpark.functions import agg_udf
import numpy as np

# 定义聚合UDF，处理分组后的列数据集合
@agg_udf(input_types=[float], return_type=float, packages=["numpy"])
def quantile_udf(x):
    return np.quantile(x, 0.95)

@agg_udf(input_types=[float], return_type=float, packages=["numpy"])
def mean_udf(x):
    return np.mean(x)

# 分组后用agg()调用聚合UDF，指定别名方便查看结果
tf = df_operation.groupBy('STORE_ID').agg(mean_udf("REG_SLS_U").alias("AVG_REG_SLS_U"))

# 打印数据：用show()格式化输出到控制台，或collect()转为Python列表
tf.show()
# 也可以用以下方式获取数据列表
# print(tf.collect())

关键说明 #

• 用agg_udf装饰器替代普通udf，明确这是聚合类型UDF，Snowpark会自动将分组后的整列数据传入UDF。
• groupBy().agg()是Snowpark执行分组聚合的标准写法，支持同时传入多个聚合UDF，例如：

  tf = df_operation.groupBy('STORE_ID').agg(
      mean_udf("REG_SLS_U").alias("AVG_REG_SLS_U"),
      quantile_udf("REG_SLS_U").alias("95PCT_REG_SLS_U")
  )
  

• 查看数据时，show()默认输出前10行格式化结果，collect()会将全量数据转为Python列表，适合后续业务处理。

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