如何高效读取PySpark DataFrame列中GCS路径对应的二进制数据？

我有一个包含GCS Blob路径的Spark DataFrame，示例代码如下：

df = spark.createDataFrame(
    [("gs://my/bucket/image.jpg",), ("gs://my/bucket/image2.jpg",)], 
    ["image_path"]
)

我想读取这些URI对应的二进制数据并存储到新列中，预期效果类似：

import pyspark.sql.functions as F

df.withColumn("image_blob", magic_read_function(F.col("image_path")))

需要一种高效的实现方式：Spark尽量批量读取数据，同时避免shuffle，也不要把所有路径收集到单个Worker节点。

我已经考虑过以下几种方案，但都有缺陷：

• 通过paths = [row.image_path for row in df.collect()]获取路径列表后读取：会强制shuffle，把所有路径汇聚到单个节点，不符合需求。
• 编写UDF读取路径数据：每条路径单独发起请求，无法实现Worker级批量读取，网络开销大。
• 读取父存储桶下所有文件再与URI列内连接：会引发shuffle，且依赖所有路径来自同一存储桶的假设，不想加这个限制。

推荐两种无shuffle、支持批量读取的实现方式：

方案1：使用mapPartitions批量读取 #

利用Spark的mapPartitions算子，让每个Worker节点处理自己分区内的所有路径，实现分区级批量读取，既避免单条请求的网络开销，也不会触发shuffle或路径汇聚到单点。

示例代码：

import pyspark.sql.functions as F
from pyspark.sql.types import BinaryType, StringType, StructType, StructField
from google.cloud import storage

def read_blob_batch(partition):
    # 每个分区仅初始化一次GCS客户端，减少连接开销
    client = storage.Client()
    for row in partition:
        path = row.image_path
        # 解析GCS路径，拆分存储桶和Blob名称
        bucket_name, blob_name = path.replace("gs://", "").split("/", 1)
        blob = client.bucket(bucket_name).blob(blob_name)
        yield (path, blob.download_as_bytes())

# 定义输出数据结构
output_schema = StructType([
    StructField("image_path", StringType(), nullable=False),
    StructField("image_blob", BinaryType(), nullable=True)
])

# 应用mapPartitions，保持原分区结构，无shuffle
result_df = df.rdd.mapPartitions(read_blob_batch).toDF(output_schema)

优势： #

• 分区级批量处理，Worker自主读取数据，无跨节点数据传输
• 每个分区仅初始化一次GCS客户端，降低连接开销
• 不依赖所有路径在同一存储桶的假设

方案2：利用Spark binaryFile数据源（Spark 2.4+） #

借助Spark内置的binaryFile数据源读取指定路径的二进制数据，再通过广播关联原DataFrame，避免shuffle的同时利用Spark原生的批量读取优化。

示例代码：

# 提取唯一路径，避免重复读取
unique_paths_df = df.select("image_path").distinct()

# 读取所有唯一路径的二进制数据
binary_df = spark.read.format("binaryFile").load(unique_paths_df.rdd.map(lambda x: x.image_path).collect())

# 用广播关联原DataFrame，避免shuffle
result_df = df.join(
    F.broadcast(binary_df.select(F.col("path").alias("image_path"), F.col("content").alias("image_blob"))),
    on="image_path"
)

注意： #

• 如果唯一路径数量极大，collect()可能导致Driver压力，此时可以将路径写入临时文本文件，再让binaryFile读取该文件中的路径
• 广播关联仅适用于唯一路径数据集较小的场景，若数据量过大，可调整分区策略实现无shuffle关联

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