📚 课程目录第4章 Snowpark for Python核心机制

4. Snowpark for Python核心机制

Snowpark for Python 的价值不在于“能在 Snowflake 里写 Python”，而在于它把 Python 代码的执行位置从本地客户端推进到 Snowflake 的计算引擎内部。理解这一点，才能理解 DataFrame API、UDF、UDTF、存储过程这四类机制各自解决什么问题、代码在哪里运行、数据如何流动。本节围绕“代码执行位置”这一主线展开，先建立整体执行模型，再逐层拆解四类机制，最后给出在 Snowflake 内执行 Python 机器学习代码的完整实现路径。

一、整体执行模型：客户端只负责“描述”，引擎负责“执行”

Snowpark 采用惰性求值（lazy evaluation）模型。当你在本地 Python 环境中写下 session.table("T").filter(...).select(...) 时，这些操作并不会立即执行，而是在客户端构建一棵逻辑计划树。只有遇到 collect()、show()、count()、to_pandas() 等动作算子（action）时，Snowpark 才把这棵计划树编译成 SQL，提交给 Snowflake 执行，再把结果拉回客户端。

这意味着 DataFrame API 的“计算”发生在 Snowflake 的虚拟仓库里，客户端只承担计划构建与结果接收。而 UDF、UDTF、存储过程则进一步把用户自定义的 Python 逻辑也推入引擎侧执行。四类机制的执行位置关系如下：

二、DataFrame API：把 Python 表达式编译成 SQL

DataFrame API 是 Snowpark 的基础层。它提供与 Pandas 相似的链式写法，但底层生成的是 Snowflake SQL。例如对一张客户表做分组统计：

from snowflake.snowpark import Session
from snowflake.snowpark.functions import col, avg, count

session = Session.builder.configs(conn_params).create()

df = session.table("ML_DB.PUBLIC.CUSTOMER_FEATURES")

result = (
    df.filter(col("REGION") == "EAST")
      .group_by("SEGMENT")
      .agg(
          count("CUSTOMER_ID").alias("CUST_CNT"),
          avg("AVG_ORDER_AMT").alias("AVG_AMT")
      )
      .sort(col("AVG_AMT").desc())
)

result.show(10)

这段代码在客户端只构建计划，show(10) 触发后编译为一条带 WHERE、GROUP BY、ORDER BY 的 SQL，在虚拟仓库中并行执行。关键判断依据是：DataFrame API 适合能用 SQL 表达的集合运算，它不启动 Python 进程，因此吞吐高、开销低。当逻辑无法用 SQL 表达（如调用 scikit-learn 模型、复杂字符串解析）时，才需要下沉到 UDF 或 UDTF。

三、UDF：逐行标量变换

UDF（User Defined Function）接收一行的一个或多个列值，返回一个标量。它适合特征工程中的单行计算，例如把原始金额做对数变换，或用已训练模型对单行打分。注册方式有两种：匿名临时 UDF 与持久化 UDF。持久化 UDF 需要指定 packages、imports 与 stage_location。

from snowflake.snowpark.functions import udf
from snowflake.snowpark.types import FloatType
import numpy as np

@udf(name="LOG_AMT_UDF",
     return_type=FloatType(),
     packages=["numpy"],
     stage_location="@ML_DB.PUBLIC.ML_STAGE")
def log_amt(x: float) -> float:
    if x is None or x <= 0:
        return 0.0
    return float(np.log1p(x))

df2 = df.with_column("LOG_AMT", log_amt(col("AVG_ORDER_AMT")))
df2.show(5)

执行时，Snowflake 在每个并行工作节点上启动 Python 沙箱，按批次把行数据喂给 log_amt，返回结果再拼回结果集。需要注意：UDF 是逐行调用，Python 解释器开销会随行数放大，因此应尽量把逻辑向量化，或改用 UDTF 批量处理。判断 UDF 是否合适，可以看三点：输入是否单行、输出是否单值、逻辑是否无法用 SQL 内置函数表达。

四、UDTF：单行输入、多行输出

UDTF（User Defined Table Function）用于把一行展开成多行，典型场景是批量预测：输入一行客户特征，输出该客户在多个候选商品上的预测分数。UDTF 通过 process() 方法逐行 yield 结果，并用 end_partition() 做分区收尾。

from snowflake.snowpark.functions import udtf
from snowflake.snowpark.types import StructType, StructField, StringType, FloatType

@udtf(name="PREDICT_SCORES_UDTF",
      output_schema=StructType([
          StructField("ITEM_ID", StringType()),
          StructField("SCORE", FloatType())
      ]),
      packages=["scikit-learn", "joblib", "numpy"],
      imports=["@ML_DB.PUBLIC.ML_STAGE/model.pkl"],
      stage_location="@ML_DB.PUBLIC.ML_STAGE")
class PredictScores:
    def __init__(self):
        import joblib
        self.model = joblib.load("model.pkl")

    def process(self, age: int, amt: float, region: str):
        import numpy as np
        X = np.array([[age, amt, hash(region) % 100]])
        probs = self.model.predict_proba(X)[0]
        for i, p in enumerate(probs):
            yield (f"ITEM_{i:03d}", float(p))

    def end_partition(self):
        pass

调用时用 join_table_function 或 LATERAL 语法，把 UDTF 的输出与原始行关联。UDTF 的优势在于模型只加载一次（__init__ 中完成），随后在分区内批量处理，比逐行 UDF 更高效。判断依据是：当输出行数不确定、需要展开或需要模型常驻内存时，优先选 UDTF。

五、存储过程：过程级 Python 编排

存储过程（Stored Procedure）把一段 Python 程序整体注册到 Snowflake，可包含多条 SQL、循环、条件分支、DDL 与调用 UDF/UDTF。它适合训练编排：从特征表读取数据、训练模型、写回模型文件到 Stage、记录训练元数据。

CREATE OR REPLACE PROCEDURE ML_DB.PUBLIC.TRAIN_MODEL()
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = '3.10'
PACKAGES = ('snowflake-snowpark-python','scikit-learn','joblib','pandas')
HANDLER = 'run'
AS
$$
import joblib
from snowflake.snowpark import Session
from sklearn.ensemble import GradientBoostingClassifier

def run(session: Session) -> str:
    df = session.table("ML_DB.PUBLIC.TRAIN_FEATURES").to_pandas()
    X = df[["AGE","AVG_ORDER_AMT","LOG_AMT"]]
    y = df["CHURNED"]
    model = GradientBoostingClassifier(n_estimators=120, max_depth=3)
    model.fit(X, y)
    joblib.dump(model, "/tmp/model.pkl")
    session.file.put("/tmp/model.pkl",
                     "@ML_DB.PUBLIC.ML_STAGE", auto_compress=False, overwrite=True)
    return f"trained rows={len(df)}, acc={model.score(X,y):.4f}"
$$;

调用方式为 CALL ML_DB.PUBLIC.TRAIN_MODEL();。存储过程运行在 Snowflake 引擎侧，数据无需外传，训练完成后模型文件落在 Stage，供 UDF/UDTF 通过 imports 加载。这样就形成了“存储过程训练 → Stage 存模型 → UDF/UDTF 推理”的闭环。

六、四类机制的选型与协作

机制	输入/输出	执行位置	典型用途	选型信号
DataFrame API	表 → 表	仓库 SQL 引擎	过滤、聚合、连接、窗口	逻辑可用 SQL 表达
UDF	单行 → 单值	Python 沙箱	特征变换、单行打分	逐行标量计算
UDTF	单行 → 多行	Python 沙箱	批量预测、展开	输出行数不定、模型常驻
存储过程	过程 → 结果	Python 沙箱	训练编排、调度、DDL	多语句、循环、写文件

实践建议：先用 DataFrame API 完成所有能用 SQL 表达的特征加工，把数据收敛到“每行一个样本”的宽表；再用 UDF 做轻量标量变换；当需要模型推理且输出多行时改用 UDTF；训练与调度逻辑放进存储过程。这样既避免逐行 Python 开销，又让模型文件与依赖包统一由 Stage 管理。

常见误区：把整个训练循环写进 UDF。UDF 是逐行调用，无法承载模型训练；训练应放在存储过程或客户端，训练产物写入 Stage 后再由 UDF/UDTF 加载。另外，UDF/UDTF 的 packages 与 imports 必须显式声明，否则沙箱内会因缺少依赖或模型文件而报错。

把四类机制串起来看，Snowpark for Python 的核心机制可以概括为：DataFrame API 负责“用 Python 写 SQL”，UDF 负责“逐行标量 Python”，UDTF 负责“逐行展开 Python”，存储过程负责“过程级 Python 编排”。四者共享同一套会话、同一份 Stage 与同一批依赖包，最终让 Python 机器学习代码在 Snowflake 内部完成从特征、训练到推理的完整链路，而数据始终不离开平台。