Snowflake Connector fetch_pandas_all()处理大量数据失败排查

问题描述 #

尝试将Snowflake数据导出为CSV时，现有代码仅能处理少量记录，且每次运行支持的最大记录数不稳定：

• 当SQL查询limit设为100时，脚本可成功生成CSV；
• 当limit设为2000左右及以上时，df = cur.fetch_pandas_all()执行失败，有时触发内存错误，多数情况脚本直接退出。

用户代码如下：

import pandas as pd
import snowflake.connector
import time
import csv

# Calculate the start time
start = time.time()

credentials = {
    'account'    : 'my account',
    'user'     : 'my user',
    'authenticator' : 'externalbrowser',
    'database' : ' my database',
    'schema' : 'my schema',
    'warehouse' : 'my warehouse',
    'role' : 'my role'
    }

query = "fake query limit 2500"
p = r"out csv path"
with snowflake.connector.connect(**credentials) as cnx:
    print("Running...")
    cur  = cnx.cursor()
    print("Executing query...")
    cur.execute(query)
    print("Fetch...")
    df = cur.fetch_pandas_all()
    print(len(df))
    df.to_csv(p,  sep=',',  header=True,index=False)
end = time.time()
length = end - start
print("It took", round(length/60), "minutes!")
print("Done!!")

问题原因 #

• 一次性加载内存过载：fetch_pandas_all()会将查询结果全部加载到内存中，当数据量较大时，直接超出Python进程的可用内存上限，导致内存错误或进程崩溃。每次运行的最大记录数不稳定，是因为系统可用内存会随其他进程占用情况变化。
• Pandas内存开销放大：DataFrame对数据的存储有额外开销，尤其是包含字符串、日期等复杂类型的字段，实际占用内存会远大于原始数据大小。
• 连接器默认配置无分批机制：Snowflake Python连接器默认未开启分批获取，强制一次性拉取所有查询结果，进一步加重内存压力。

解决办法 #

1. 分批获取并逐批写入CSV #

避免一次性加载全量数据，通过fetchmany()分批获取并写入文件，大幅降低内存占用：

import pandas as pd
import snowflake.connector
import time
import csv

start = time.time()

credentials = {
    'account': 'my account',
    'user': 'my user',
    'authenticator': 'externalbrowser',
    'database': 'my database',
    'schema': 'my schema',
    'warehouse': 'my warehouse',
    'role': 'my role'
}

query = "fake query"
p = r"out csv path"
batch_size = 1000  # 可根据内存情况调整批次大小

with snowflake.connector.connect(**credentials) as cnx:
    print("Running...")
    cur = cnx.cursor()
    print("Executing query...")
    cur.execute(query)
    
    # 写入CSV表头
    header = [col[0] for col in cur.description]
    with open(p, 'w', newline='', encoding='utf-8') as csv_file:
        writer = csv.writer(csv_file)
        writer.writerow(header)
        
        # 分批获取并写入数据
        print("Fetching batches...")
        batch_num = 0
        while True:
            batch_data = cur.fetchmany(batch_size)
            if not batch_data:
                break
            writer.writerows(batch_data)
            batch_num += 1
            print(f"完成第{batch_num}批写入，共{len(batch_data)}条记录")

end = time.time()
print(f"耗时{round(end - start)/60}分钟！")
print("Done!!")

2. 使用Snowflake原生COPY命令导出（推荐） #

让Snowflake直接将数据导出到云存储（如S3、Azure Blob），再下载到本地，完全避免Python内存瓶颈：

-- 先创建外部阶段（若未创建）
CREATE OR REPLACE STAGE my_csv_stage
URL = 's3://your-bucket/path/'
CREDENTIALS = (AWS_KEY_ID='your-key' AWS_SECRET_KEY='your-secret');

-- 导出数据到阶段
COPY INTO '@my_csv_stage/exported_data.csv'
FROM (SELECT * FROM your_table)
FILE_FORMAT = (TYPE = CSV FIELD_OPTIONALLY_ENCLOSED_BY = '"' HEADER = TRUE);

之后可通过Snowflake连接器的get_stage_file方法将文件下载到本地。

3. 优化Pandas内存占用（适合必须用DataFrame的场景） #

通过指定数据类型减少DataFrame内存开销：

with snowflake.connector.connect(**credentials) as cnx:
    cur = cnx.cursor()
    cur.execute(query)
    # 获取字段名和自定义数据类型
    cols = [desc[0] for desc in cur.description]
    dtype_config = {
        'string_col': 'category',  # 重复率高的字符串用category类型
        'int_col': 'int32',        # 缩小整数类型
        'float_col': 'float32'     # 缩小浮点数类型
    }
    # 加载数据并转换类型
    df = pd.DataFrame(cur.fetchall(), columns=cols).astype(dtype_config)
    df.to_csv(p, sep=',', header=True, index=False)

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