使用Dask将Snowflake的SQL表导入Jupyter时遇序列化错误求助

一、修复当前代码的错误 #

你的代码出现Could not serialize object of type HighLevelGraph错误，核心原因有三个：

• 重复初始化Dask Client：先创建了client = Client()，后续又用with Client() as client重复启动集群，导致任务图序列化冲突。
• 未正确处理Snowflake连接：fetch_data函数中使用的conn未定义，且Snowflake连接对象无法被序列化传递给Dask Worker，必须在每个任务内部独立创建连接。
• 无效的重复查询：循环10次调用全表查询会生成重复数据，完全没必要。

修正后的基础代码 #

import snowflake.connector
from dask.distributed import Client
from dask import delayed
from dask.dataframe import from_delayed
import dask.dataframe as dd

# Snowflake连接参数
connection_parameters = {
        'user': user,
        'password': password,
        'account': account,
        'warehouse': warehouse,
        'database': database,
        'schema': schema
}

query = """SELECT * FROM my_table"""

# 仅初始化一次Dask Client
client = Client()

@delayed
def fetch_data(query, conn_params):
    # 每个任务内部独立创建/关闭Snowflake连接，避免序列化问题
    conn = snowflake.connector.connect(**conn_params)
    cur = conn.cursor()
    cur.execute(query)
    df = cur.fetch_pandas_all()
    cur.close()
    conn.close()
    return df

# 单任务读取全表
delayed_dfs = [fetch_data(query, connection_parameters)]
ddf = dd.from_delayed(delayed_dfs)
result_df = ddf.compute()

如果需要分块读取10万行的表，可以基于主键做分片查询（假设表有自增主键id）：

@delayed
def fetch_chunk(query, conn_params, start, end):
    chunk_query = f"{query} WHERE id BETWEEN {start} AND {end}"
    conn = snowflake.connector.connect(**conn_params)
    cur = conn.cursor()
    cur.execute(chunk_query)
    df = cur.fetch_pandas_all()
    cur.close()
    conn.close()
    return df

# 分成10个1万行的块
delayed_dfs = [fetch_chunk(query, connection_parameters, i*10000+1, (i+1)*10000) for i in range(10)]
ddf = dd.from_delayed(delayed_dfs)
result_df = ddf.compute()

二、更简便的Dask读取Snowflake方法 #

推荐使用官方维护的dask-snowflake库，它专门优化了Dask与Snowflake的交互，无需手动处理连接和分块逻辑。

步骤1：安装依赖 #

pip install dask-snowflake

步骤2：使用示例代码 #

from dask_snowflake import read_snowflake
from dask.distributed import Client

client = Client()

connection_parameters = {
        'user': user,
        'password': password,
        'account': account,
        'warehouse': warehouse,
        'database': database,
        'schema': schema
}

# 自动分块读取Snowflake表
ddf = read_snowflake(
    query="SELECT * FROM my_table",
    connection_kwargs=connection_parameters,
    chunksize=10000  # 每个分区1万行，适配10万行的表规模
)

# 执行计算得到Pandas DataFrame
result_df = ddf.compute()

该方法的优势 #

• 自动处理Snowflake连接的创建与销毁，彻底避免序列化问题。
• 基于Snowflake元数据自动分块，或手动指定chunksize实现并行读取。
• 支持复杂SQL查询、谓词下推等优化，性能远优于手动编写delayed任务。

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