如何使用Python迭代地从大型.csv文件中提取数据并保存为.csv块？

import csv
import os

def extract_data_from_large_csv(input_file, output_folder, block_size):
    with open(input_file, 'r') as file:
        reader = csv.reader(file)
        header = next(reader)  # 读取并保存头部信息

        output_count = 1
        output_file = os.path.join(output_folder, f'output_{output_count}.csv')
        output = open(output_file, 'w', newline='')
        writer = csv.writer(output)
        writer.writerow(header)  # 写入头部信息

        row_count = 0
        for row in reader:
            writer.writerow(row)
            row_count += 1

            if row_count >= block_size:
                output.close()
                output_count += 1
                output_file = os.path.join(output_folder, f'output_{output_count}.csv')
                output = open(output_file, 'w', newline='')
                writer = csv.writer(output)
                writer.writerow(header)  # 写入头部信息
                row_count = 0

        output.close()

    return output_folder

input_file = 'path/to/large_file.csv'
output_folder = 'path/to/output_folder'
block_size = 10000

extract_data_from_large_csv(input_file, output_folder, block_size)