Polars LazyFrame连接操作内存溢出问题求助

问题场景 #

执行以下Polars代码时，在Mac环境中出现内存耗尽、交换空间占满至90GB后程序崩溃的情况：

df1 = pl.scan_parquet("output_data/tmpbase.parquet").cast({"owning_account_id": pl.Categorical}).sort("owning_account_id")
df2 = pl.scan_parquet("output_data/tmpintermediate.parquet").cast({"transaction_id": pl.Categorical}).sort("transaction_id")
df1.join(
    df2,
    on=["owning_account_id", "transaction_id"],
    how="inner",
).sink_parquet("output_data/table.parquet")

原本预期Lazy API的Join操作支持流式处理，可应对内存无法容纳的大数据集，但实际出现内存溢出，核心问题在于对Polars Join执行逻辑的理解偏差。

核心原因 #

Polars的Lazy API虽支持分块流式处理，但归并Join（Polars默认优先选择的高效Join方式）要求两个表必须按相同的连接键顺序排序。

在上述代码中：

• df1仅按owning_account_id排序
• df2仅按transaction_id排序
  两者的排序键与连接键的顺序完全不匹配，Polars无法执行流式归并Join，只能退化为将其中一个表（甚至两个表）全量加载到内存中执行哈希Join，直接导致内存爆炸。

解决方案 #

1. 统一两表的排序键（与连接键顺序一致） #

将两个表都按照连接键的完整顺序排序，让Polars可以执行高效的归并Join，全程流式处理数据，无需加载全量数据到内存：

# 两表均按连接键顺序排序：owning_account_id 优先，其次 transaction_id
df1 = pl.scan_parquet("output_data/tmpbase.parquet")\
    .cast({"owning_account_id": pl.Categorical, "transaction_id": pl.Categorical})\
    .sort(["owning_account_id", "transaction_id"])

df2 = pl.scan_parquet("output_data/tmpintermediate.parquet")\
    .cast({"owning_account_id": pl.Categorical, "transaction_id": pl.Categorical})\
    .sort(["owning_account_id", "transaction_id"])

df1.join(
    df2,
    on=["owning_account_id", "transaction_id"],
    how="inner",
).sink_parquet("output_data/table.parquet")

注意：确保两个表的所有连接键都已转为Categorical类型（如果适合业务场景），且排序顺序完全一致。

2. 强制使用分块哈希Join（无法排序时） #

如果无法预先按连接键排序，可通过指定join_strategy="hash"并开启streaming=True（需Polars 0.19及以上版本），让Polars采用分块哈希Join，限制单块数据的内存占用：

df1.join(
    df2,
    on=["owning_account_id", "transaction_id"],
    how="inner",
    join_strategy="hash",
    streaming=True
).sink_parquet("output_data/table.parquet")

该方式会将其中一个表拆分为多个小数据块，逐个加载到内存执行哈希匹配，降低峰值内存消耗，但性能略低于归并Join。

3. 提前过滤冗余数据 #

若数据集中存在大量无需参与Join的行，可先执行过滤操作缩减数据量，再进行Join：

# 示例：过滤掉连接键为空的无效行
df1 = pl.scan_parquet("output_data/tmpbase.parquet")\
    .filter(pl.col("owning_account_id").is_not_null() & pl.col("transaction_id").is_not_null())\
    .cast({"owning_account_id": pl.Categorical, "transaction_id": pl.Categorical})\
    .sort(["owning_account_id", "transaction_id"])

df2 = pl.scan_parquet("output_data/tmpintermediate.parquet")\
    .filter(pl.col("owning_account_id").is_not_null() & pl.col("transaction_id").is_not_null())\
    .cast({"owning_account_id": pl.Categorical, "transaction_id": pl.Categorical})\
    .sort(["owning_account_id", "transaction_id"])

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