Kedro PartitionedDataSet是否支持惰性保存以节省内存？

问题背景 #

使用Kedro的PartitionedDataSet管理图片数据集，配置如下：

raw_images:
  type: PartitionedDataSet
  <<: *data_path_on_disk
  dataset:
    type: pillow.ImageDataSet
  filename_suffix: ".png"

定义节点处理图片并输出到同类型的分区数据集：

node(func=crop_image, inputs="raw_images", outputs="cropped_images")

处理函数实现：

def crop_image(images: dict):
    return {image_path: image().crop([10,10,20,20]) for image_path, image in images.items()}

核心问题解答 #

当前实现下，所有图片会被完整加载到内存中，处理大数据集时极易引发内存过载，且不会逐步写入磁盘。

原因说明 #

Kedro的PartitionedDataSet默认采用eager加载模式：

1. 运行节点时，会一次性遍历所有分区，将每张图片通过pillow.ImageDataSet加载到内存，构建成完整的字典传入crop_image函数。
2. 你的函数会基于这个全量字典生成包含所有处理后图片的新字典，直到整个字典构建完成后，Kedro才会批量将所有处理结果写入磁盘。

优化方案（避免内存过载） #

要实现逐张处理、逐步写入，需开启PartitionedDataSet的lazy加载模式（Kedro 0.17及以上版本支持）：

1. 修改数据集配置，添加load_args开启lazy模式：

raw_images:
  type: PartitionedDataSet
  <<: *data_path_on_disk
  dataset:
    type: pillow.ImageDataSet
  filename_suffix: ".png"
  load_args:
    lazy: True

2. 调整处理函数逻辑（配合lazy模式进一步降低内存占用）：
   开启lazy后，传入函数的字典中的值是延迟加载的加载器，只有访问时才会真正读取图片。可以改用迭代方式处理，避免一次性生成全量结果字典：

def crop_image(images: dict):
    processed = {}
    for image_path, image_loader in images.items():
        # 仅在需要时加载单张图片
        with image_loader() as img:
            processed[image_path] = img.crop([10,10,20,20])
        # 处理完单张后，原始图片内存可被自动回收
    return processed

这种模式下，内存中只会同时存在一张原始图片和一张处理后的图片，大幅降低内存压力，Kedro在接收到结果字典后会逐个写入对应分区。

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