Snowpark Python UDTF是否支持传入多个DataFrame列作为入参

问题描述 #

基于Snowpark/Python编写接收单个列作为参数的UDTF时运行正常，未查询到该特性相关官方文档，确认是否支持向UDTF传入多个列（即DataFrame子集）作为入参。
编写的测试代码运行失败，抛出异常：
"TypeError: 'TABLE FUNCTION' expected Column or str, got: <class 'snowflake.snowpark.dataframe.DataFrame'>"
排除将多个列拼接为单个列传入UDTF的方案，需要可行的实现方式，测试代码如下：

import uuid
@udtf(output_schema=["c1","c2","x"], 
      input_types =[StringType(), StringType(), IntegerType()],
      name="udft_two_col_test", 
      replace=True, 
      session=ses)
class udft_two_col_test:
    def process(self, c1:str, c2:str, n: int) -> Iterable[Tuple[str, str, str]]: 
        for i in range(n):
            yield (c1, c2, f'{n}-{c1}-{c2}')

            
df = ses.create_dataframe([str(uuid.uuid4()).split('-') for i in range(1,10,1)], schema=['c1','c2','c3','c4','c5'])
df.sort('c1','c2').show()

------------------------------------------------
|"C1"      |"C2"  |"C3"  |"C4"  |"C5"          |
------------------------------------------------
|125a9845  |f7e2  |48dd  |b51c  |42ba82531fe7  |
|136da5dc  |62cb  |47c0  |98f9  |4182421e6d2b  |
|300380e2  |b365  |4d6a  |8d6b  |1092e4c24ec8  |
|3d9d9882  |0fb2  |4209  |bf11  |4341b0336946  |
|43c4147d  |1603  |4548  |ad8e  |4df50cddd682  |
|9e1024ca  |61d5  |404d  |88f8  |79393083eb30  |
|bf25e899  |5697  |4c36  |8533  |e3009c68ce9b  |
|d6dd677f  |035b  |49e7  |9236  |316741579f3c  |
|f4b83587  |26e1  |48cf  |8563  |0586ccb6602e  |
------------------------------------------------

df.join_table_function("udft_two_col_test", df["c1","c2"], lit(3)).sort('c1','c2').show(100)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
...
---> 17 df.join_table_function("udft_two_col_test", df["c1","c2"], lit(3)).sort('c1','c2').show(100)
...
TypeError: 'TABLE FUNCTION' expected Column or str, got: <class 'snowflake.snowpark.dataframe.DataFrame'>

报错原因 #

Snowpark Python中使用df["c1","c2"]的多列切片语法，返回值是仅包含选中列的新DataFrame对象，不是多个独立的Column对象。而join_table_function方法要求UDTF的每个入参必须是Column类型对象或者代表列名的字符串，传入DataFrame对象就会触发该类型错误。
Snowpark Python UDTF原生支持传入多个独立列作为入参，不存在功能层面的限制。

可行解决方案 #

• 直接按UDTF定义的参数顺序，逐个传入对应列即可，无需修改UDTF逻辑，也不需要拼接列，修正后的调用代码如下：

df.join_table_function("udft_two_col_test", df["c1"], df["c2"], lit(3)).sort('c1','c2').show(100)

• 如果需要动态传入不确定数量的列，可以先将待传入的列存入列表，调用时用*解包列表传参，适配动态列场景：

# 支持从配置、其他逻辑动态生成待传入的列列表
input_cols = [df["c1"], df["c2"]]
# 解包传入，和逐个传参效果一致
df.join_table_function("udft_two_col_test", *input_cols, lit(3)).sort('c1','c2').show(100)

传入列时也可以直接使用列名字符串，比如"c1"、"c2"，Snowpark会自动识别为对应列。

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