如何高效从Databricks的PySpark DataFrame提取大体积Base64图片数据？

问题背景 #

在Databricks中处理PySpark DataFrame，仅关注image列：该列每行存储约540,696字符的Base64编码图片字符串，当前数据集150,000行且会持续增长。目标是将这些图片字符串提取到Notebook本地Python环境（如列表），后续转换为嵌入向量存入向量数据库，要求方案高效可扩展，且需避免将数据写入云存储（如DBFS、S3）以降低成本。

已尝试的低效方案 #

• collect()：数据量过大导致Driver内存溢出
• toPandas()：大规模数据集下同样出现内存不足问题
• toLocalIterator()/循环迭代：仅适用于小批量数据，大规模场景下速度极慢
• foreachPartition()：在Worker节点执行逻辑，但无法将结果返回Driver
• 基于limit()/索引范围/分区分批：可避免崩溃，但性能表现不佳

当前采用的最优但耗时极久的实现：

from pyspark.sql.functions import monotonically_increasing_id
import pyspark.sql.functions as f

images_df_with_index = images_df.withColumn("index", monotonically_increasing_id())

batch_size = 1000
total_rows = images_df_with_index.count()
num_batches = total_rows // batch_size + 1

all_images = []

for i in range(num_batches):
    lower = i * batch_size
    upper = lower + batch_size

    batch_df = images_df_with_index.filter((f.col("index") >= lower) & (f.col("index") < upper))
    batch = batch_df.select("image").collect()
    all_images.extend(batch)

优化方案 #

1. 分区迭代优化（减少重复计算） #

原方案中使用monotonically_increasing_id()加filter的方式会触发多次Job，每次批量都要重新扫描数据，效率极低。改为先调整分区，再通过toLocalIterator()按分区拉取数据，仅需一次Shuffle操作：

# 根据Driver内存调整分区大小，示例为每个分区处理2000行（约1GB数据）
target_partitions = (images_df.count() // 2000) + 1
# 仅选择需要的image列并重新分区
images_repartitioned = images_df.select("image").repartition(target_partitions)

all_images = []
# 按分区迭代，每个分区的数据一次性拉取到Driver
for partition in images_repartitioned.toLocalIterator():
    # 直接提取image字段值，避免存储Row对象节省内存
    batch = [row.image for row in partition]
    all_images.extend(batch)

2. 分布式预处理嵌入向量（减少数据传输） #

若后续需将图片转换为嵌入向量，可利用Spark的Pandas UDF在Worker节点完成Base64解码和嵌入生成，仅将体积更小的向量传回Driver，大幅降低数据传输量和Driver内存压力：

from pyspark.sql.functions import pandas_udf, ArrayType, FloatType
import pandas as pd
import base64

# 替换为你的嵌入向量生成逻辑
def generate_embeddings(images_series: pd.Series) -> pd.Series:
    def process_single_image(base64_str):
        # Base64解码
        img_bytes = base64.b64decode(base64_str)
        # 调用嵌入模型生成向量，示例为伪代码
        embedding = your_embedding_model.encode(img_bytes)
        return embedding.tolist()
    
    return images_series.apply(process_single_image)

# 注册Pandas UDF
embedding_udf = pandas_udf(generate_embeddings, ArrayType(FloatType()))

# 在Worker节点分布式生成嵌入向量
embeddings_df = images_df.select(embedding_udf("image").alias("embedding"))

# 拉取向量到本地环境
all_embeddings = []
for partition in embeddings_df.toLocalIterator():
    all_embeddings.extend([row.embedding for row in partition])

# 存入向量数据库
your_vector_database.add_embeddings(all_embeddings)

3. 调整Driver内存配置 #

在Databricks集群配置中适当提高Driver内存（如从8GB提升至16GB/32GB），可支持一次性拉取更大的分区数据，减少迭代次数，进一步提升效率。

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