如何按行数拆分(.csv)文件，但保留每个拆分子文件上的第一行(列标题)？python

import csv
import os

def split_csv_file(input_file, output_dir, rows_per_file):
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    with open(input_file, 'r', newline='') as file:
        reader = csv.reader(file)
        header = next(reader)  # 读取第一行（列标题）

        file_index = 1
        row_count = 0
        output_file = None

        for row in reader:
            # 如果达到指定的行数，创建一个新的拆分子文件
            if row_count % rows_per_file == 0:
                if output_file:
                    output_file.close()
                output_file = open(os.path.join(output_dir, f'output_{file_index}.csv'), 'w', newline='')
                writer = csv.writer(output_file)
                writer.writerow(header)  # 写入列标题到每个拆分子文件
                file_index += 1

            writer.writerow(row)
            row_count += 1

        if output_file:
            output_file.close()

# 示例用法
input_file = 'input.csv'  # 输入文件路径
output_dir = 'output'  # 输出目录路径
rows_per_file = 1000  # 每个拆分子文件的行数

split_csv_file(input_file, output_dir, rows_per_file)