Spark 3.0以下版本如何计算DataFrame大小以确定Shuffle分区数？

一、计算DataFrame总大小的方法 #

Spark 3.0及以上版本提供了DataFrame.approxTotalSize()方法直接估算大小，但3.0以下版本需要通过以下几种方式实现：

1. 缓存后通过RDD获取准确大小 #

这种方法会将数据缓存到内存/磁盘，能得到较准确的实际存储大小，注意缓存资源的占用：

from pyspark import StorageLevel

# 缓存DataFrame（可选磁盘存储避免内存压力）
df.persist(StorageLevel.DISK_ONLY)
# 触发action让缓存生效
df.count()

# 获取总字节数
total_bytes = df.rdd.getSizeInBytes()
# 转换为可读单位（比如GB）
total_gb = total_bytes / (1024**3)
print(f"DataFrame总大小约为 {total_gb:.2f} GB")

# 计算完成后释放缓存
df.unpersist()

2. 抽样估算总大小 #

适合快速预估，无需缓存，通过抽样计算单行平均大小再乘以总行数：

# 获取总行数（数据量极大时，可用近似计数加快速度）
total_rows = df.count()
# 近似计数方式：total_rows = df.rdd.countApprox(1000, 0.95)

# 抽取样本（根据数据量调整抽样比例，保证样本代表性）
sample_fraction = 0.000003  # 3亿条数据约抽取900条
sample_df = df.sample(withReplacement=False, fraction=sample_fraction, seed=42)

# 定义计算单行大小的函数
def get_row_size(row):
    import sys
    return sys.getsizeof(row)

# 计算平均单行大小
avg_row_size = sample_df.rdd.map(get_row_size).mean()

# 估算总大小
total_bytes = total_rows * avg_row_size
total_gb = total_bytes / (1024**3)
print(f"DataFrame估算总大小约为 {total_gb:.2f} GB")

注意：sys.getsizeof计算的是Python对象大小，与Spark序列化后的实际存储大小有一定误差，仅作参考。

3. 通过Spark UI查看实际大小 #

运行一个轻量action（比如df.limit(1).count()）后，进入Spark UI：

• 若已缓存数据，在Storage页面可直接查看DataFrame的总存储大小
• 未缓存的话，在Jobs页面查看对应job的每个task输入大小，累加后得到总数据量

二、根据DataFrame大小计算Shuffle分区数 #

Spark官方推荐Shuffle分区的单个文件大小在128MB-256MB之间，可按以下公式计算：

# 设置目标分区大小（这里以256MB为例）
target_partition_size = 256 * 1024 * 1024  # 字节

# 计算所需Shuffle分区数
shuffle_partitions = max(1, int(total_bytes / target_partition_size))

# 设置Spark参数
spark.conf.set("spark.sql.shuffle.partitions", str(shuffle_partitions))

若涉及Join操作，建议以较大的那个DataFrame的大小为基准计算分区数，保证Shuffle后的分区大小合理。

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