使用SQL/Snowflake连接器抽取数据时内存开销过高的优化咨询

我尝试通过SQL查询从Snowflake数据库抽取数据到DataFrame（后续要转成torch tensor），遇到以下情况：

• 结果集是3.85亿行 × (1个整数 + 1个浮点数)，理论大小约4.6GB，Snowflake GUI显示压缩后实际结果更小，查询本身仅耗时40秒；
• 但代码读取数据时内存占用飙升至75GB以上，耗时超40分钟（Snowflake查询历史仍显示查询仅40秒）；加载完成后DataFrame大小符合预期，用malloc没找到内存开销来源。

核心需求：如何降低数据读取过程中的内存开销？

我试过两种代码实现：

代码尝试1（SQLAlchemy） #

import pandas as pd
from sqlalchemy import Engine, text

def load_data(date: date, engine: Engine) -> pd.DataFrame:
    """Load all data and case_ids."""
    query = text(my query) # 依赖传入的date参数

    with engine.connect() as connection:
        df = pd.read_sql_query(query, con=connection)
    return df

代码尝试2（Snowflake Connector） #

import pandas as pd
import snowflake
def load_data(date : date) -> pd.DataFrame:
    query = my query
    with snowflake.connector.connect(**credentials) as conn:
       with conn.cursor() as cur:
           rows = cur.execute(query).fetchall()
    # 将rows转换为DataFrame
    return df

1. 分批读取数据，避免一次性加载全量 #

通过分批拉取结果集，把峰值内存控制在单批次数据的大小范围内，两种连接器都支持该方式：

Snowflake Connector分批读取实现

import pandas as pd
import snowflake.connector

def load_data(date: date) -> pd.DataFrame:
    query = my query
    chunk_size = 1_000_000  # 每次拉取100万行，可根据内存情况调整
    df_list = []
    
    with snowflake.connector.connect(**credentials) as conn:
        with conn.cursor() as cur:
            cur.execute(query)
            while True:
                chunk = cur.fetchmany(chunk_size)
                if not chunk:
                    break
                chunk_df = pd.DataFrame(chunk, columns=['int_col', 'float_col'])
                df_list.append(chunk_df)
    
    return pd.concat(df_list, ignore_index=True)

SQLAlchemy + pandas分批读取实现

import pandas as pd
from sqlalchemy import Engine, text

def load_data(date: date, engine: Engine) -> pd.DataFrame:
    query = text(my query)
    chunk_size = 1_000_000
    df_list = []
    
    with engine.connect() as connection:
        for chunk in pd.read_sql_query(query, con=connection, chunksize=chunk_size):
            df_list.append(chunk)
    
    return pd.concat(df_list, ignore_index=True)

2. 提前指定紧凑数据类型，减少内存浪费 #

pandas默认推断的int64/float64类型内存占用大，可根据数据范围指定更小的类型（比如int32/float32），直接在读取阶段设置：

SQLAlchemy指定数据类型

df = pd.read_sql_query(
    query, 
    con=connection,
    dtype={
        'int_column': 'int32',  # 替换为实际列名
        'float_column': 'float32'
    }
)

Snowflake Connector指定数据类型

chunk_df = pd.DataFrame(
    chunk, 
    columns=['int_col', 'float_col'],
    dtype={'int_col': 'int32', 'float_col': 'float32'}
)

3. 使用Snowflake原生的fetch_pandas_all方法 #

Snowflake Connector提供了直接生成DataFrame的方法，比手动fetchall再转换更高效，能减少中间数据的内存开销：

import pandas as pd
import snowflake.connector

def load_data(date: date) -> pd.DataFrame:
    query = my query
    with snowflake.connector.connect(**credentials) as conn:
        with conn.cursor() as cur:
            cur.execute(query)
            df = cur.fetch_pandas_all()
            # 可选：直接转换为紧凑类型
            df['int_col'] = df['int_col'].astype('int32')
            df['float_col'] = df['float_col'].astype('float32')
    return df

4. 直接转换为torch Tensor，跳过全量DataFrame #

如果最终目标是生成torch tensor，可以在分批读取时直接转换，无需先存储全量DataFrame，进一步压缩内存占用：

import torch
import snowflake.connector

def load_data_to_tensor(date: date) -> tuple[torch.Tensor, torch.Tensor]:
    query = my query
    chunk_size = 1_000_000
    int_tensors = []
    float_tensors = []
    
    with snowflake.connector.connect(**credentials) as conn:
        with conn.cursor() as cur:
            cur.execute(query)
            while True:
                chunk = cur.fetchmany(chunk_size)
                if not chunk:
                    break
                int_col = torch.tensor([row[0] for row in chunk], dtype=torch.int32)
                float_col = torch.tensor([row[1] for row in chunk], dtype=torch.float32)
                int_tensors.append(int_col)
                float_tensors.append(float_col)
    
    return torch.cat(int_tensors), torch.cat(float_tensors)

5. 检查Snowflake结果集返回格式 #

Snowflake Connector默认用Python原生类型存储行数据，内存开销远大于pandas的紧凑类型。使用fetch_pandas_all或指定返回格式为pandas类型，可避免此类额外开销。

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