SQLMesh Python 模型入门（二）：依赖管理、引擎实战与内存优化

本系列基于 SQLMesh 官方文档（https://sqlmesh.readthedocs.io/en/stable/concepts/models/python_models/）整理，共 3 篇，面向初学者。本篇是第二篇，解决"数据怎么取、算在哪、内存爆了怎么办"三个实战问题。

• 第（一）篇：基础语法与核心概念
• 第（三）篇：前后置语句、蓝图批量建模、避坑清单

---

1. 依赖管理：resolve_table 与 depends_on

想读取上游模型的数据，必须先解析出它在当前环境下的真实表名 ，这就是 resolve_table 的职责：

table = context.resolve_table("docs_example.upstream_model")
df = context.fetchdf(f"SELECT * FROM {table}")

resolve_table 有两重作用：

1. 返回当前运行环境（如 dev 环境的 schema__dev 前缀）下正确的表名；
2. 自动把被引用的模型登记为当前模型的依赖。

另一种声明依赖的方式是在 @model 装饰器里显式写 depends_on。规则是：装饰器里显式声明的依赖优先于函数体内的动态引用。看这个官方例子：

@model(
    "my_model.with_explicit_dependencies",
    depends_on=["docs_example.upstream_dependency"],  # ✅ 会被捕获
)
def execute(context, start, end, execution_time, **kwargs):
    # ❌ 由于装饰器里已声明依赖，这里的引用会被忽略
    context.resolve_table("docs_example.another_dependency")
    ...

此外，用户自定义的全局变量或蓝图变量也可以出现在 resolve_table 的调用中：

@model(
    "@schema_name.test_model2",
    kind="FULL",
    columns={"id": "INT"},
)
def execute(context, **kwargs):
    table = context.resolve_table(f"{context.var('schema_name')}.test_model1")
    select_query = exp.select("*").from_(table)
    return context.fetchdf(select_query)

---

2. 实战示例三连：Basic → SQL+Pandas → PySpark

2.1 查询上游模型 + Pandas 处理

最常用的一种模式：SQL 负责取数和粗筛，pandas 负责灵活加工：

import typing as t
from datetime import datetime

import pandas as pd
from sqlmesh import ExecutionContext, model


@model(
    "docs_example.sql_pandas",
    columns={
        "id": "int",
        "name": "text",
    },
)
def execute(
    context: ExecutionContext,
    start: datetime,
    end: datetime,
    execution_time: datetime,
    **kwargs: t.Any,
) -> pd.DataFrame:
    # 获取上游模型的表名，并自动登记为依赖
    table = context.resolve_table("upstream_model")

    # 把数据取回来；如果引擎是 Spark，这里返回的就是 Spark DataFrame
    df = context.fetchdf(f"SELECT id, name FROM {table}")

    # 做一些 pandas 擅长的事
    df["id"] += 1
    return df

2.2 PySpark 示例：分布式计算的正确姿势

如果你使用 Spark 引擎，推荐直接用 Spark DataFrame API 而不是 Pandas------数据全程在集群分布式计算，不会拉到本地。

import typing as t
from datetime import datetime

from pyspark.sql import DataFrame, functions
from sqlmesh import ExecutionContext, model


@model(
    "docs_example.pyspark",
    columns={
        "id": "int",
        "name": "text",
        "country": "text",
    },
)
def execute(
    context: ExecutionContext,
    start: datetime,
    end: datetime,
    execution_time: datetime,
    **kwargs: t.Any,
) -> DataFrame:
    # 获取上游模型表名并登记依赖
    table = context.resolve_table("upstream_model")

    # 用 Spark DataFrame API 增加一列 country
    df = context.spark.table(table).withColumn("country", functions.lit("USA"))

    # 直接返回 PySpark DataFrame，本地不计算任何数据
    return df

三个关键学习点：

1. 通过 context.spark 拿到 SparkSession，再用 spark.table(表名) 载入数据，比先 fetchdf 转成 Pandas 高效得多；
2. 返回类型注解写 pyspark.sql.DataFrame 而不是 pd.DataFrame；
3. 只要最终返回的是 Spark DataFrame，计算就发生在集群上，没有本地内存瓶颈。

---

3. 同样思路的另外两种引擎：Snowpark 与 Bigframe

3.1 Snowpark（Snowflake 引擎）

用 context.snowpark 操作 DataFrame，计算下推到 Snowflake：

import typing as t
from datetime import datetime

from snowflake.snowpark.dataframe import DataFrame
from sqlmesh import ExecutionContext, model


@model(
    "docs_example.snowpark",
    columns={
        "id": "int",
        "name": "text",
        "country": "text",
    },
)
def execute(
    context: ExecutionContext,
    start: datetime,
    end: datetime,
    execution_time: datetime,
    **kwargs: t.Any,
) -> DataFrame:
    # 直接返回 snowpark DataFrame，本地不计算任何数据
    df = context.snowpark.create_dataframe(
        [[1, "a", "usa"], [2, "b", "cad"]], schema=["id", "name", "country"]
    )
    df = df.filter(df.id > 1)
    return df

3.2 Bigframe（BigQuery 引擎）

用 context.bigframe，所有计算都在 BigQuery 完成。它甚至支持把本地 Python 函数注册为 remote function 在集群上执行：

import typing as t
from datetime import datetime

from bigframes.pandas import DataFrame
from sqlmesh import ExecutionContext, model


def get_bucket(num: int):
    if not num:
        return "NA"
    boundary = 10
    return "at_or_above_10" if num >= boundary else "below_10"


@model(
    "mart.wiki",
    columns={
        "title": "text",
        "views": "int",
        "bucket": "text",
    },
)
def execute(context, start, end, execution_time, **kwargs) -> DataFrame:
    # 把本地 Python 函数包装成 BigQuery remote function
    remote_get_bucket = context.bigframe.remote_function([int], str)(get_bucket)

    # 只返回 Bigframe 句柄，数据不落到本地
    df = context.bigframe.read_gbq("bigquery-samples.wikipedia_pageviews.200809h")
    df = (
        df[df.title.str.contains(r"[Gg]oogle")]
        .groupby(["title"], as_index=False)["views"]
        .sum(numeric_only=True)
        .sort_values("views", ascending=False)
    )
    return df.assign(bucket=df["views"].apply(remote_get_bucket))

四种返回类型选择速查：

场景	返回类型	入口	计算位置
通用 / 小数据量	Pandas DataFrame	context.fetchdf	本地内存
Spark 引擎	PySpark DataFrame	context.spark	集群分布式
Snowflake 引擎	Snowpark DataFrame	context.snowpark	Snowflake 内
BigQuery 引擎	Bigframe DataFrame	context.bigframe	BigQuery 内

---

4. 大输出分块：用生成器 yield 降低内存占用

Pandas 是单机内存型框架，数据量太大时内存会爆，又用不了 Spark 时，SQLMesh 允许用 Python 生成器 yield 把输出拆成多批，每次只把一小块数据加载进内存：

@model(
    "docs_example.batching",
    columns={
        "id": "int",
    },
)
def execute(
    context: ExecutionContext,
    start: datetime,
    end: datetime,
    execution_time: datetime,
    **kwargs: t.Any,
) -> pd.DataFrame:
    table = context.resolve_table("upstream_model")

    for i in range(3):
        # 分 3 次查询，每次只取一块数据，避免内存耗尽
        df = context.fetchdf(f"SELECT id from {table} WHERE id = {i}")
        yield df

---

5. 空表禁忌：绝对不能 return 空 DataFrame

Python 模型不允许返回空的 DataFrame 。如果你的代码有可能产出空结果，必须改成条件 yield：

@model(
    "my_model.empty_df"
)
def execute(
    context: ExecutionContext,
) -> pd.DataFrame:
    # ... 生成 df 的代码 ...

    if df.empty:
        yield from ()        # 空的话什么都不产出
    else:
        yield df

记住口诀：有数据 → yield df；可能没数据 → 永远不要 return 空表。

---

6. 序列化：代码到底在哪里运行？

SQLMesh 通过自研的序列化框架，在运行 SQLMesh 的机器本地执行 Python 代码。这意味着：

• 你的 Python 环境（依赖包版本）需要就绪；
• 如果引擎是 Spark/Snowflake/BigQuery 且你返回的是对应 DataFrame，重计算会下推到集群，本地只做编排。

---

小结

1. 读上游模型必先 resolve_table------硬编码表名会在 dev/prod 切换时拿错数据；
2. depends_on 显式声明优先于函数体内的动态引用；
3. 引擎有原生 DataFrame API（spark/snowpark/bigframe）时，尽量让计算留在集群，别拉回本地；
4. 输出太大用生成器分批 yield，可能为空的结果绝不用 return 空表。

下一篇讲工程化细节：前后置语句、宏变量在属性中的坑、以及用蓝图一次生成多个模型。

---

参考资料：SQLMesh 官方文档 --- Python models（https://sqlmesh.readthedocs.io/en/stable/concepts/models/python_models/）