如何加快EMR PySpark Notebook中DataFrame写入S3的速度

• 首要问题：代码中df = df.join(df2.select('page_md'))没有指定关联键，触发了全量笛卡尔积，13万行数据的笛卡尔积会生成169亿行的超大表，完全超出合理计算量，所以任务卡在计算阶段无法进入写入步骤。小数据集测试时因数据量极小，笛卡尔积可以快速计算完成，因此未暴露问题。
• 次要问题1：mapPartitionsWithIndex返回的id_是分区索引而非行唯一标识，没有可用于关联两张表的唯一键，无法正确将提取的文本和原始行对应。
• 次要问题2：boto3调用S3未设置超时和重试限制，拉取Common Crawl文件时如果出现网络波动会导致分区任务无限卡住。
• 次要问题3：缺少必要的Spark配置，未禁用笛卡尔积、未配置S3连接参数，可能导致任务异常或性能低下。
• 次要问题4：写入前强制coalesce(4)会将全量数据集中到4个任务处理，大幅降低并行度，也可能拖慢写入速度。

1. 先添加必要的Spark配置 #

在Notebook代码最开头、导入库之后添加以下配置即可，不需要额外的magic命令：

# 禁用笛卡尔积，避免误操作触发全量交叉连接
spark.conf.set("spark.sql.crossJoin.enabled", "false")
# 配置S3连接超时与重试
spark.conf.set("spark.hadoop.fs.s3a.connection.timeout", "10000")
spark.conf.set("spark.hadoop.fs.s3a.attempts.maximum", "5")
# 调整shuffle并行度匹配你的集群规模
spark.conf.set("spark.sql.shuffle.partitions", "200")

2. 修复数据关联逻辑 #

给原始数据添加行唯一ID，处理后通过ID关联，避免笛卡尔积。

3. 优化S3调用配置 #

给boto3客户端添加超时和重试参数，避免单请求卡住整个分区。

4. 调整写入逻辑 #

移除不必要的coalesce(4)，保留合理并行度写入S3，同时添加写入模式避免残留文件影响。

import html2text
import boto3
import botocore
import os
import re
import zlib
import gzip
from bs4 import BeautifulSoup as bs
from bs4 import Comment
import logging
from pyspark.sql.functions import monotonically_increasing_id

# 基础Spark配置
spark.conf.set("spark.sql.crossJoin.enabled", "false")
spark.conf.set("spark.hadoop.fs.s3a.connection.timeout", "10000")
spark.conf.set("spark.hadoop.fs.s3a.attempts.maximum", "5")
spark.conf.set("spark.sql.shuffle.partitions", "200")

def load_index():
    input_file='s3://cc-stuff/uploads/DMOZ_bussineses_ccindex.csv'
    df = spark.read.option("header",True).csv(input_file)
    # 新增唯一行ID，用于后续关联
    df = df.withColumn("row_id", monotonically_increasing_id())
    return df
    
def process_warcs(iterator):
    html_textract = html2text.HTML2Text()
    html_textract.ignore_links = True
    html_textract.ignore_images = True
    # 新增S3超时、重试配置
    no_sign_request = botocore.client.Config(
        signature_version=botocore.UNSIGNED,
        connect_timeout=5,
        read_timeout=30,
        retries={'max_attempts': 3}
    )
    s3client = boto3.client('s3', config=no_sign_request)
    s3pattern = re.compile('^s3://([^/]+)/(.+)')
    PREFIX = "s3://commoncrawl/"

    for row in iterator:
        text = None
        try:
            start_byte = int(row['warc_record_offset'])
            stop_byte = (start_byte + int(row['warc_record_length']))
            s3match = s3pattern.match((PREFIX + row['warc_filename']))
            bucketname = s3match.group(1)
            path = s3match.group(2)
            resp = s3client.get_object(Bucket=bucketname, Key=path, Range='bytes={}-{}'.format(start_byte, stop_byte))
            content = resp['Body'].read()
            data = zlib.decompress(content, wbits = zlib.MAX_WBITS | 16).decode('utf-8',errors='ignore')
            data = data.split('\r\n\r\n',2)[2]
            soup = bs(data,'html.parser')
            for x in soup.findAll(text=lambda text:isinstance(text, Comment)):
                x.extract()  
            for x in soup.find_all(["head","script","button","form","noscript","style"]):
                x.decompose()
            text = html_textract.handle(str(soup))
        except Exception as e:
            pass
        # 返回行ID和提取的文本，用于关联
        yield (row['row_id'], text)
        
def run_job(write_out=True):
    df = load_index()
    # 改用mapPartitions，不再需要分区ID
    df2 = df.rdd.repartition(200).mapPartitions(process_warcs).toDF(["row_id", "page_md"])
    # 基于row_id关联，避免笛卡尔积
    df = df.join(df2, on="row_id", how="left")
    if write_out:             
        output = "s3://cc-stuff/emr-out/DMOZ_bussineses_ccHTML"
        # 移除coalesce(4)，保留并行度写入，添加overwrite模式避免残留文件影响
        df.write.mode("overwrite").json(output) 
    return df

df = run_job(write_out=True)

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