如何在Kedro流水线中为pandas.SQLTableDataSet使用chunk_size参数

根因说明 #

你当前配置的chunk_size放在save_args下，仅作用于数据写入数据库的分块逻辑，不会影响数据集读取输出的类型。如果运行流水线时拿到的是生成器而非DataFrame，说明你在数据集的load_args中也配置了chunk_size，此时Pandas会返回分块迭代器，Kedro会直接透传该结果到节点输入。

对应解决方法 #

场景1：仅需要写入时分块，读取需返回完整DataFrame #

仅保留save_args中的chunk_size配置即可，移除load_args下的chunk_size参数，你给出的现有catalog配置已经满足写入分块的要求，不需要额外调整，此时读取该数据集得到的就是完整的pd.DataFrame对象，写入时会自动按每块1000行批量写入Redshift。

场景2：表体积过大，必须分块读取处理 #

如果需要通过分块读取避免内存溢出，需要调整节点逻辑适配迭代器输入：

1. 首先在catalog中补充读取分块配置：

table_name:
  type: pandas.SQLTableDataSet
  credentials: redshift
  table_name : rs_table_name
  layer: output
  load_args:
    chunk_size: 1000
  save_args:
    if_exists: append
    schema: schema.name
    chunk_size: 1000

2. 对应处理节点的函数要适配迭代器输入，逐块处理后返回迭代器：

import pandas as pd
from typing import Iterator

def process_chunked_table(table_chunks: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
    for chunk in table_chunks:
        # 此处写单块数据的处理逻辑
        processed_chunk = chunk.drop_duplicates(subset=["unique_key"])
        yield processed_chunk

SQLTableDataSet会自动识别输入的迭代器类型，逐块写入数据库，不需要额外调整保存逻辑。

注意事项 #

• 分块读取得到的迭代器仅能遍历一次，如果处理后的结果需要供给多个下游节点使用，不要直接返回迭代器，需要先合并所有块为完整DataFrame再输出，避免下游节点读取时迭代器已耗尽。
• Redshift场景下建议将写入chunk_size调整到单块数据量1MB~10MB对应行数，避免单次请求过大触发写入限流。

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