如何使用PySpark并行处理同一张图片并避免不必要的数据副本？

针对你用PySpark处理S3图像数据时遇到的这些问题——既要并行处理单张图片的多个多边形，又要避免不必要的数据副本，还要解决数据倾斜——我整理了一套实用的方案，一步步来：

数据倾斜的核心原因是单条记录里绑定了最多500个多边形，导致部分分区数据量远超其他分区。我们可以先把每个(ImagePath, poly_list)拆分成多条(ImagePath, single_poly)的记录，但要注意只提前解析图像元数据（尺寸、通道数），不加载完整图像，避免提前占用内存：

from pyspark.sql import SparkSession
from pyspark.sql.functions import explode, udf
from pyspark.sql.types import StructType, StructField, IntegerType
from PIL import Image

# 假设你的原始数据是RDD格式
spark = SparkSession.builder.appName("ImagePolyProcessing").getOrCreate()
image_path_poly_rdd = spark.sparkContext.parallelize(image_path_poly_list)

# 转成DataFrame并拆分多边形列表
image_df = image_path_poly_rdd.toDF(["image_path", "poly_list"])
exploded_df = image_df.withColumn("poly", explode("poly_list")).drop("poly_list")

# 自定义UDF提取图像元数据（仅读取头部信息，不加载像素）
def extract_image_meta(s3_path):
    with Image.open(s3_path) as img:
        return (img.width, img.height, len(img.getbands()))

meta_schema = StructType([
    StructField("width", IntegerType(), True),
    StructField("height", IntegerType(), True),
    StructField("channels", IntegerType(), True)
])

extract_meta_udf = udf(extract_image_meta, meta_schema)
exploded_df = exploded_df.withColumn("image_meta", extract_meta_udf("image_path"))

要避免同一张图像被多次加载产生数据副本，关键是把同一图像的所有多边形记录分到同一个Executor分区里。这样在处理分区时，只需要加载一次图像，就能复用它处理所有关联的多边形：

# 按图像路径分区，确保同图的多边形在同一分区
partitioned_df = exploded_df.repartition("image_path")

接下来用mapPartitions处理每个分区，在分区内复用加载好的图像，并通过线程池并行处理多个多边形（线程池共享内存，不会拷贝图像数据，完美避免不必要的副本）：

from concurrent.futures import ThreadPoolExecutor
import numpy as np

def process_image_partition(partition):
    partition_records = list(partition)
    if not partition_records:
        return
    
    # 同一分区的所有记录属于同一张图，取第一个路径即可
    s3_image_path = partition_records[0]["image_path"]
    # 从S3加载图像（懒加载，仅当访问像素时才会读取）
    with Image.open(s3_image_path) as img:
        img_np = np.array(img)
        
        # 定义单多边形处理逻辑
        def process_single_poly(record):
            poly = record["poly"]
            meta = record["image_meta"]
            
            # 计算多边形的边界框（根据你的多边形格式调整）
            x_coords = [p[0] for p in poly]
            y_coords = [p[1] for p in poly]
            x_min, x_max = min(x_coords), max(x_coords)
            y_min, y_max = min(y_coords), max(y_coords)
            
            # 裁剪多边形对应区域，复用已加载的图像数组
            cropped_region = img_np[y_min:y_max, x_min:x_max]
            
            # 示例：提取指定通道并计算统计值
            target_channel = cropped_region[:, :, 0] if meta["channels"] >=1 else cropped_region
            return (s3_image_path, poly, target_channel.mean())
        
        # 用线程池并行处理该图像的所有多边形
        with ThreadPoolExecutor() as executor:
            yield from executor.map(process_single_poly, partition_records)

# 应用分区处理，得到最终结果RDD
processed_results_rdd = partitioned_df.rdd.mapPartitions(process_image_partition)

如果你的数据里存在重复的图像路径（比如同一张图被多次关联），可以在Executor本地做图像缓存，避免重复从S3下载：

import os
import boto3
from botocore.exceptions import ClientError

def load_cached_image(s3_path, cache_dir="/tmp/spark_image_cache"):
    # 创建本地缓存目录
    os.makedirs(cache_dir, exist_ok=True)
    # 提取S3路径中的文件名作为缓存键
    s3_parts = s3_path.replace("s3://", "").split("/")
    local_filename = s3_parts[-1]
    local_path = os.path.join(cache_dir, local_filename)
    
    if not os.path.exists(local_path):
        # 从S3下载到本地缓存
        s3_client = boto3.client("s3")
        bucket = s3_parts[0]
        s3_key = "/".join(s3_parts[1:])
        try:
            s3_client.download_file(bucket, s3_key, local_path)
        except ClientError as e:
            print(f"Failed to download {s3_path}: {str(e)}")
            raise
    
    return Image.open(local_path)

把这个函数替换之前的Image.open(s3_image_path)即可，同一Executor上再次处理同一张图时，直接读取本地缓存，省去S3下载开销。

对于11000×11000的超大图像，加载全图会占用大量内存。好在PIL的Image.open是懒加载的，我们可以直接裁剪多边形对应的区域，不需要加载完整图像：

def load_poly_region(s3_path, poly):
    with Image.open(s3_path) as img:
        x_coords = [p[0] for p in poly]
        y_coords = [p[1] for p in poly]
        x_min, x_max = min(x_coords), max(x_coords)
        y_min, y_max = min(y_coords), max(y_coords)
        # 直接裁剪区域，仅加载所需像素
        cropped = img.crop((x_min, y_min, x_max, y_max))
        return np.array(cropped)

如果你的多边形处理只需要区域内的像素，用这个方法可以大幅降低内存占用，避免不必要的全图数据副本。

内容的提问来源于stack exchange，提问作者Igor Franzoni Okuyama