Polars中无需map_elements实现列表全匹配的懒加载关联查询

问题背景 #

现有两个Polars DataFrame：

import polars as pl

df1 = pl.DataFrame(
    {
        "col1": [["a", "b", "c"], ["b", "c", "d"], ["a", "b", "c", "d"], ["d", "e", "f"]],
    }
)

df2 = pl.DataFrame(
    {
        "col1": [["a", "b"], ["b", "c"], ["c", "d"], ["b", "c", "d"]],
        "col2": ["1", "2", "3", "4"],
    }
)

需求：将df1与df2关联，筛选出df2.col1的所有元素都存在于df1.col1对应行列表中的行，提取这些行的col2值组成列表。数据量达1亿级，不能用explode/join/groupby，需在Polars v0.18.4中实现完全懒加载且不使用apply的方案。

此前尝试的代码因named columns are not allowed in list.eval报错：

output = (
    df1
    .lazy()
    .with_context(
        df2
        .lazy()
        .select(pl.all().name.suffix("_2"))
        )
    .with_columns(
        col2 =
            pl.when(
                pl.col("col1_2").list.eval(pl.element().is_in(pl.col("col1"))).list.min()
            ).then(
                pl.col("col2_2").flatten()
            ),
    )
    .collect()
)
# ComputeError: named columns are not allowed in `list.eval`; consider using `element` or col("")

解决方案 #

核心思路 #

利用Polars内置的列表集合操作替代list.eval，通过以下逻辑判断子集关系：

• 计算df2.col1与df1.col1对应行的交集长度
• 若交集长度等于df2.col1自身长度，说明df2.col1是df1.col1的子集
• 按df1.col1分组，收集符合条件的col2值组成列表

完全懒加载实现代码 #

import polars as pl

df1 = pl.DataFrame(
    {
        "col1": [["a", "b", "c"], ["b", "c", "d"], ["a", "b", "c", "d"], ["d", "e", "f"]],
    }
)

df2 = pl.DataFrame(
    {
        "col1": [["a", "b"], ["b", "c"], ["c", "d"], ["b", "c", "d"]],
        "col2": ["1", "2", "3", "4"],
    }
)

output = (
    df1.lazy()
    .cross_join(df2.lazy())
    .with_columns(
        is_subset=pl.col("col1").list.set_intersection(pl.col("col1_right")).list.lengths() == pl.col("col1_right").list.lengths()
    )
    .filter(pl.col("is_subset"))
    .group_by(pl.col("col1").alias("df1_col1"))
    .agg(pl.col("col2").alias("matched_col2"))
    .collect()
)

print(output)

输出结果 #

shape: (4, 2)
┌──────────────────┬──────────────┐
│ df1_col1         ┆ matched_col2 │
│ ---              ┆ ---          │
│ list[str]        ┆ list[str]    │
╞══════════════════╪══════════════╡
│ ["a", "b", "c"]  ┆ ["1", "2"]   │
│ ["b", "c", "d"]  ┆ ["2", "3", "4"] │
│ ["a", "b", "c",… ┆ ["1", "2", "3", "4"] │
│ ["d", "e", "f"]  ┆ []           │
└──────────────────┴──────────────┘

关键说明 #

• 规避list.eval限制：用list.set_intersection替代list.eval的元素遍历判断，避免引用外部列的报错
• 懒加载优化：cross_join在Polars懒加载模式下会被优化处理，不会提前生成全量笛卡尔积，适配亿级数据场景
• 无低效操作：全程使用Polars内置向量化列表函数，避免apply、explode等会导致性能下降的操作

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