Polars中逐行生成DataFrame并定期刷盘的高效方案

问题背景 #

现有一个逐行生成数据的generation_mechanism()，以及能将单行数据拆解为特征的decompose()函数，需要实现基于Polars的定期捕获数据并刷盘到CSV的逻辑，以下是两种候选方案：

方案一：基于列表迭代构建，批量生成DataFrame刷盘 #

import polars as pl
# 仅示例两个特征
feature_a_list = []
feature_b_list = []
flush_threshold = 100
count = 0
for row in generation_mechanism():
    feature_a, feature_b = decompose(row)
    feature_a_list.append(feature_a)
    feature_b_list.append(feature_b)
    count += 1
    if count == flush_threshold:
        data = pl.DataFrame({"feature_a": feature_a_list, "feature_b": feature_b_list})
        with open('filename.csv', 'a') as file:
            # 首次写入才加表头，后续追加跳过
            include_header = file.tell() == 0
            data.write_csv(file, include_header=include_header)
        # 刷盘后清空列表，避免重复写入
        feature_a_list.clear()
        feature_b_list.clear()
        count = 0
# 处理剩余未达到阈值的数据
if count > 0:
    data = pl.DataFrame({"feature_a": feature_a_list, "feature_b": feature_b_list})
    with open('filename.csv', 'a') as file:
        data.write_csv(file, include_header=False)

方案二：逐行构建小DataFrame，通过vstack拼接后刷盘 #

import polars as pl
flush_threshold = 100
count = 0
data = pl.DataFrame(schema={"feature_a": pl.String, "feature_b": pl.String})
for row in generation_mechanism():
    feature_a, feature_b = decompose(row)
    new_data = pl.DataFrame({"feature_a": feature_a, "feature_b": feature_b})
    data = data.vstack(new_data)
    count += 1
    if count == flush_threshold:
        with open('filename.csv', 'a') as file:
            include_header = file.tell() == 0
            data.write_csv(file, include_header=include_header)
        # 重置空DataFrame
        data = pl.DataFrame(schema={"feature_a": pl.String, "feature_b": pl.String})
        count = 0
# 处理剩余数据
if count > 0:
    with open('filename.csv', 'a') as file:
        data.write_csv(file, include_header=False)

方案对比与最优选择 #

• 方案一效率更高：Polars的DataFrame是基于列的内存结构，直接用列表收集列数据再批量构造DataFrame，避免了频繁vstack带来的内存拷贝开销。vstack每次拼接都需要重新分配内存并复制数据，当阈值较大时，性能差距会非常明显。
• 方案一更规范：符合Polars推荐的"批量处理"思路，Polars对批量数据的处理优化远优于逐行操作，逐行构建小DataFrame属于反模式，会浪费Polars的性能优势。

其他优化思路 #

1. 直接写入CSV（跳过DataFrame）：如果不需要对中间数据做Polars的计算操作，可以直接将拆解后的特征按CSV格式写入文件，完全避免DataFrame的内存开销：

import csv
flush_threshold = 100
count = 0
buffer = []
with open('filename.csv', 'a', newline='') as file:
    writer = csv.writer(file)
    # 若文件为空，先写入表头
    if file.tell() == 0:
        writer.writerow(["feature_a", "feature_b"])
    for row in generation_mechanism():
        feature_a, feature_b = decompose(row)
        buffer.append([feature_a, feature_b])
        count += 1
        if count == flush_threshold:
            writer.writerows(buffer)
            buffer.clear()
            count = 0
    # 写入剩余数据
    if buffer:
        writer.writerows(buffer)

2. 使用Polars的pl.Series批量构造：相比普通列表，pl.Series可以提前指定数据类型，减少构造DataFrame时的类型推断开销：

import polars as pl
flush_threshold = 100
count = 0
feature_a_series = pl.Series(dtype=pl.String)
feature_b_series = pl.Series(dtype=pl.String)
for row in generation_mechanism():
    feature_a, feature_b = decompose(row)
    feature_a_series = feature_a_series.append(pl.Series([feature_a]))
    feature_b_series = feature_b_series.append(pl.Series([feature_b]))
    count += 1
    if count == flush_threshold:
        data = pl.DataFrame({"feature_a": feature_a_series, "feature_b": feature_b_series})
        with open('filename.csv', 'a') as file:
            include_header = file.tell() == 0
            data.write_csv(file, include_header=include_header)
        feature_a_series = pl.Series(dtype=pl.String)
        feature_b_series = pl.Series(dtype=pl.String)
        count = 0
if count > 0:
    data = pl.DataFrame({"feature_a": feature_a_series, "feature_b": feature_b_series})
    with open('filename.csv', 'a') as file:
        data.write_csv(file, include_header=False)

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