Snowflake Python UDTF中sys._xoptions获取导入目录的正确方式

问题描述 #

开发Snowflake Python用户定义表函数（UDTF）时，需要从文件读取机器学习（ML）模型。参考相关示例代码后产生两点疑问：

1. 示例假设类级别代码会在Snowflake服务器执行，通过伪造sys._xoptions读取上传文件，但我认为类级别代码是本地解释器处理类定义时执行的，此时df_mapping会在本地创建后序列化发送到Snowflake，本地的sys._xoptions无法获取远程导入路径。这个观点是否正确？
2. Snowflake文档提到UDTF类的__init__方法会在处理每个分区前调用，是否应该在此处查询sys._xoptions并加载文件？

解答 #

1. 类级别代码执行时机的判断 #

你的观点部分正确，需要补充关键细节：

• 本地定义UDTF类时，类级别代码确实会在本地执行，示例里伪造sys._xoptions只是为了避免本地定义类时因找不到snowflake_import_directory键而报错。
• 但UDTF注册到Snowflake后，远程执行时Snowflake会重新初始化这个类，此时类级别代码会再次在Snowflake服务器上执行，这时候sys._xoptions会被Snowflake自动设置为正确的远程导入路径，所以示例中的逻辑在远程是能正常读取文件的。

不过这种写法存在明显缺陷：

• 本地必须存在对应文件才能完成类定义，否则会抛出文件不存在的错误
• 若本地文件与上传到Snowflake的文件内容不一致，本地测试和远程执行结果会出现差异，增加调试难度

2. 正确的文件加载时机：__init__方法 #

是的，应该将文件加载逻辑放到__init__方法中，理由如下：

• __init__方法会在Snowflake服务器上每个分区处理前被调用，此时sys._xoptions已经被正确设置，能直接获取到远程的导入目录
• 本地定义类时不会执行__init__内的代码，无需伪造sys._xoptions，避免了本地依赖文件的问题
• 分区级别的初始化也能保证每个分区加载一次文件，平衡资源占用与执行效率

修改后的示例代码 #

##################################################################
## Define the class for the UDTF

# Import the required modules 
import pandas
import sys

# 无需再伪造sys._xoptions

# Define handler class
class leverage_external_mapping_file :
    IMPORT_DIRECTORY_NAME = "snowflake_import_directory"

    def __init__(self):
        # 在__init__中获取导入目录并加载文件
        self.import_dir = sys._xoptions[self.IMPORT_DIRECTORY_NAME]
        self.df_mapping = pandas.read_excel(self.import_dir + 'Dummy Mapping File.xlsx', skiprows=5, usecols="C:D")
  
    ## Define process method that acts
    ## on each individual input row
    def process(
          self
        , input_item: str
      ) :

        ### Apply the mapping to retrieve the mapped value
        df_mapped_group = self.df_mapping[self.df_mapping['Item']==input_item]
    
        mapped_group = 'No matching group found'
        
        if len(df_mapped_group.index) > 0 :
          mapped_group = df_mapped_group.iloc[0]['Group']

        yield(mapped_group,)

##################################################################
## Register UDTF in Snowflake

### Add packages and data types
from snowflake.snowpark.types import StructType, StructField
from snowflake.snowpark.types import StringType
snowpark_session.add_packages('pandas', 'openpyxl') # openpyxl required for pandas to read xlsx
snowpark_session.add_import('Supporting Files/Dummy Mapping File.xlsx')

### Define output schema
output_schema = StructType([
      StructField("MAPPED_ITEM", StringType())
  ])

### Upload UDTF to Snowflake
snowpark_session.udtf.register(
    handler = leverage_external_mapping_file
  , output_schema = output_schema
  , input_types = [StringType()]
  , is_permanent = True
  , name = 'SNOWPARK_LEVERAGE_EXTERNAL_MAPPING_FILE'
  , replace = True
  , stage_location = '@UDTF_STAGE'
)

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