Snowpark+DBT中外部函数在apply_in_pandas报错ModuleNotFoundError

场景与现有代码 #

使用Snowpark和dbt在Snowflake中进行数据转换并创建表/视图，当前模型代码结构如下：

import pandas as pd
from snowflake.snowpark import Session
from snowflake.snowpark.functions import col
from snowflake.snowpark.types import StructType, StructField, StringType, BooleanType, FloatType

def model(dbt, session: Session):
    dbt.config(materialized="table")
    df = dbt.ref("<SOME_TABLE>")
    df_grouped = df.group_by(col('ID1_COL'), col('ID2_COL'))

    def transform_function(df: pd.DataFrame) -> pd.DataFrame:
        """
        用于分组数据转换的函数
        """
        [...]
        return df_transformed

    final_df = df_grouped.apply_in_pandas(
        transform_function,
        output_schema=StructType([
            StructField("COL1", StringType()),
            StructField("COL2", BooleanType()),
            [...]
            StructField("COLN", FloatType()),
        ])
    )

    return final_df

问题现象 #

希望将transform_function定义在model函数外部，但无论采用哪种方式，都会触发如下错误：

ModuleNotFoundError: No module named 'main_module'

已尝试的解决方法 #

• 使用@udf装饰器及多种函数注册方式，均未解决问题
• 尝试pandas_udf，但执行时出现长时间挂起
• 该函数在不用于apply_in_pandas时可正常运行
• 暂未尝试将函数写入单独库文件，且希望避免这种实现方式

内容的提问来源于stack exchange，提问作者Juan Saiz Lomas