Python实现从CSV文件随机抽取指定行数生成新样本文件

场景1：优先抽取包含CG列的指定列数据并导出 #

import csv
import random

# 配置参数
INPUT_CSV = "test.csv"
OUTPUT_CSV = "sampled_output.csv"
SAMPLE_SIZE = 1000
# 要保留的列可自行增减，此处默认保留CG列
REQUIRED_COLUMNS = ["CG"]

# 读取CSV文件，自动识别表头
with open(INPUT_CSV, 'r', encoding='utf-8', newline='') as f:
    reader = csv.DictReader(f)
    if "CG" not in reader.fieldnames:
        raise ValueError("CSV文件中未找到名为CG的列，请检查表头")
    all_rows = list(reader)

# 随机抽样：需要无重复样本的话，把random.choices换成random.sample
sampled_rows = random.choices(all_rows, k=SAMPLE_SIZE)

# 导出抽样结果
with open(OUTPUT_CSV, 'w', encoding='utf-8', newline='') as f:
    writer = csv.DictWriter(f, fieldnames=REQUIRED_COLUMNS)
    writer.writeheader()
    for row in sampled_rows:
        writer.writerow({col: row[col] for col in REQUIRED_COLUMNS})

print(f"抽样完成，结果已导出到{OUTPUT_CSV}")

场景2：抽取完整行数据并导出 #

如果不需要筛选列，直接导出全量字段，使用下方代码即可：

import csv
import random

INPUT_CSV = "test.csv"
OUTPUT_CSV = "sampled_full_rows.csv"
SAMPLE_SIZE = 1000

with open(INPUT_CSV, 'r', encoding='utf-8', newline='') as f:
    reader = csv.DictReader(f)
    all_rows = list(reader)
    full_columns = reader.fieldnames

# 无重复抽样请替换为 random.sample(all_rows, k=SAMPLE_SIZE)
sampled_rows = random.choices(all_rows, k=SAMPLE_SIZE)

with open(OUTPUT_CSV, 'w', encoding='utf-8', newline='') as f:
    writer = csv.DictWriter(f, fieldnames=full_columns)
    writer.writeheader()
    writer.writerows(sampled_rows)

print(f"全量行抽样完成，结果已导出到{OUTPUT_CSV}")

问题对应说明 #

• 针对第一个问题：原有代码的冗余逻辑在chosen_values = [random.choice(row) for row in chosen_rows]这一行，这段逻辑是从选中行里随机取单个值，直接删除即可保留完整行数据，使用csv.DictReader可以方便筛选指定列内容，不需要手动索引对应列。
• 针对第二个问题：上述代码已经内置导出逻辑，会自动保留原表头，格式和原CSV一致。

注意事项 #

• 若CSV数据量远超内存容量，可以先统计文件行数，再生成随机行号逐行读取筛选，避免内存溢出。
• random.choices是有放回抽样，可能出现重复行，需要无重复样本的话替换为random.sample即可，注意样本量不能超过总数据行数。

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