如何用Polars合并已排序Parquet文件以避免重复全局排序？

核心思路 #

既然每个Parquet文件已经按groupby列完成排序，我们可以让Polars直接对这些有序数据集执行归并排序，而非将所有数据加载到内存后重新全局排序。关键是要明确告知Polars每个子数据集的有序性，再调用对应的合并逻辑。

具体实现代码 #

import polars as pl

# 替换为实际的排序列名
groupby = ["your_sort_columns_here"]

# 逐个读取已排序文件，标记为有序LazyFrame
sorted_lfs = []
for path in tissue_pq_paths:
    lf = pl.scan_parquet(path)
    # 标记当前数据集已按指定列排序
    sorted_lfs.append(lf.set_sorted(groupby))

# 归并排序合并所有有序数据集并写入结果
(
    pl.concat(sorted_lfs, how="vertical_relaxed")
    .sort(groupby, maintain_order=True)  # 触发归并排序而非全局排序
    .sink_parquet(output_pq_file, compression="snappy", statistics=True)
)

关键细节说明 #

• set_sorted(groupby)：必须调用该方法标记每个LazyFrame的有序性，否则Polars无法识别分段有序的特性，仍会执行全局排序。
• concat(how="vertical_relaxed")：采用宽松模式垂直拼接，兼容少量结构差异（若所有文件结构完全一致，也可使用vertical模式）。
• sort(..., maintain_order=True)：该参数是触发归并排序的核心。当输入为分段有序的数据集时，此参数会让Polars仅在各分段间执行合并操作，无需打乱所有数据重新排序，大幅节省内存与计算资源。

额外优化提示 #

• 确保所有文件的排序方向一致（均为升序或均为降序），若存在混合方向，需在set_sorted时指定descending参数，例如lf.set_sorted(groupby, descending=True)。
• 若排序列为数值型或分布均匀的字符串型，归并排序的执行效率会进一步提升。

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