下面是一个示例代码,可以按行数拆分CSV文件并保留头部:
import csv
def split_csv(input_file, output_path, rows_per_file):
with open(input_file, 'r') as file:
reader = csv.reader(file)
header = next(reader) # 读取头部
file_index = 1
row_count = 0
output_file = None
for row in reader:
if output_file is None or row_count >= rows_per_file:
if output_file is not None:
output_file.close()
output_file = open(f'{output_path}/{file_index}.csv', 'w', newline='')
writer = csv.writer(output_file)
writer.writerow(header) # 写入头部
file_index += 1
row_count = 0
writer.writerow(row)
row_count += 1
if output_file is not None:
output_file.close()
# 示例用法
split_csv('input.csv', 'output', 1000)
上述代码中,split_csv
函数接受三个参数:输入文件路径、输出文件夹路径和每个输出文件的行数限制。函数使用CSV模块读取输入文件,并逐行处理。
在处理行数达到限制或需要新的输出文件时,函数会关闭当前的输出文件(如果有),并创建一个新的输出文件。在每个输出文件中,函数会先写入头部,然后逐行写入行数据。
在示例中,split_csv('input.csv', 'output', 1000)
将会读取名为input.csv
的输入文件,将其拆分成每个文件最多包含1000行的输出文件,并保存到名为output
的文件夹中。