Polars中按行比较Struct的高效实现方案（单上下文）咨询

你需要在Polars中以Struct的identifiers字段为索引，比较两个Struct结构，检查struct1中每个(id, value)对是否在struct2中存在相同id且值一致的项（不考虑位置）。原方案使用map_elements调用Python函数逐行处理，效率极低，以下是保持单一上下文的优化方案。

原方案代码 #

def compare_structs(struct1, struct2):
    ids1, vals1 = struct1["identifiers"], struct1["values"]
    ids2, vals2 = struct2["identifiers"], struct2["values"]
    return [
        any(id1 == id2 and val1 == val2 for id2, val2 in zip(ids2, vals2))
        for id1, val1 in zip(ids1, vals1)
    ]


df.with_columns(
    pl.struct(["struct1", "struct2"])
    .map_elements(lambda s: compare_structs(s["struct1"], s["struct2"]))
    .alias("match_results")
)

优化思路 #

避免使用map_elements触发Python UDF（逐行处理效率极低），改用Polars内置的向量化列表操作，将id和value打包成元组后直接做存在性检查，完全依托Polars的底层引擎处理数据。

优化后的代码 #

方案一：分步清晰版 #

df.with_columns(
    # 将struct2的id和value打包成元组列表，作为匹配基准
    struct2_pairs=pl.col("struct2").struct["identifiers"].list.zip(pl.col("struct2").struct["values"]),
    # 对struct1的每个id-value对，检查是否在struct2的元组列表中
    match_results=pl.col("struct1").struct["identifiers"]
    .list.eval(
        pl.element().zip(pl.col("struct1").struct["values"].list.get(pl.int_range(0, pl.col("struct1").struct["values"].list.len())))
        .is_in(pl.col("struct2_pairs"))
    )
).drop("struct2_pairs")

方案二：单表达式紧凑版 #

df.with_columns(
    match_results=pl.struct(
        ids1=pl.col("struct1").struct["identifiers"],
        vals1=pl.col("struct1").struct["values"],
        ids2=pl.col("struct2").struct["identifiers"],
        vals2=pl.col("struct2").struct["values"]
    ).map(
        lambda s: (pl.Series(s["ids1"]).zip(s["vals1"])).is_in(pl.Series(s["ids2"]).zip(s["vals2"])).to_list(),
        return_dtype=pl.List(pl.Boolean)
    )
)

效率说明 #

上述方案完全利用Polars的向量化引擎处理数据，避免了Python层面的循环和逐行UDF调用，在数据量较大时性能会有显著提升。

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