PySpark读取CSV时如何抽取指定数量的随机行？

核心问题解答 #

• 无法直接通过spark.read.csv一步读取指定数量的随机行，因为Spark读取CSV时会先扫描全量文件以确定元数据（如schema、分区信息）；但可以通过后续操作高效获取指定数量的随机样本，无需将全量数据持久化或进行冗余处理。
• 要获取全局随机样本，必须遍历所有分区的数据（随机基于整个数据集），但可以避免全量数据的内存加载或全量转换。

具体实现方法 #

方法1：sample + limit 组合

通过sample按比例抽取基础样本，再用limit固定目标行数。如果不知道总数据量，可设置一个稍大的比例（确保抽样结果行数≥目标数），再截断到指定数量。

from pyspark.sql import SparkSession

# 初始化SparkSession
spark = SparkSession.builder.getOrCreate()

# 读取CSV文件（仅扫描元数据，不会全量加载）
df = spark.read.csv(
    '/content/student_data.csv',
    sep=',',
    inferSchema=True,
    header=True
)

# 抽取100行随机样本：先抽样比例0.1（可根据数据量调整），再截断到100行
random_sample = df.sample(withReplacement=False, fraction=0.1).limit(100)

# 查看结果
random_sample.show()

方法2：RDD takeSample 直接指定数量

利用底层RDD的takeSample方法，直接指定样本数量，无需计算比例，抽样完成后转回DataFrame。该方法会在各分区独立抽样后合并结果，效率较高。

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

df = spark.read.csv(
    '/content/student_data.csv',
    sep=',',
    inferSchema=True,
    header=True
)

# 抽取100个不重复的随机样本（seed用于固定抽样结果，可选）
rdd_sample = df.rdd.takeSample(withReplacement=False, num=100, seed=42)
# 将RDD转回DataFrame，保留原schema
random_sample_df = spark.createDataFrame(rdd_sample, schema=df.schema)

random_sample_df.show()

方法3：randomSplit 拆分数据集

将原数据集按比例拆分为两部分，取其中小比例的部分再截断到目标行数，适合需要同时保留原数据集的场景。

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

df = spark.read.csv(
    '/content/student_data.csv',
    sep=',',
    inferSchema=True,
    header=True
)

# 按1:9比例拆分数据集，取第一部分后截断到100行
sample_df, _ = df.randomSplit([0.1, 0.9], seed=42)
random_sample = sample_df.limit(100)

random_sample.show()

补充说明 #

• 若无需严格随机，仅需快速获取N行数据，可直接用df.limit(100)，但这会返回文件中的前100行，不具备随机性。
• 若提前明确CSV的schema，读取时通过schema参数指定（而非inferSchema=True），可以减少Spark扫描全量文件的开销，但仍需遍历所有分区才能获取全局随机样本。

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