如何用Polars高效计算两列的笛卡尔积并哈希处理？

Polars的apply方法本质是逐行调用Python函数，会带来巨大性能开销，必须用向量化操作替代。以下是针对不同场景的高效实现方案：

场景1：A、B为列表列（每行是一组元素） #

假设你的DataFrame结构如下：

import polars as pl

df = pl.DataFrame({
    "A": [[1, 2], [3, 4]],
    "B": [[5, 6], [7, 8]]
})

高效实现代码 #

# 生成展开后的结果（每行对应一个笛卡尔积元素）
df2 = (
    df
    .with_row_index("row_id")  # 标记原行索引，用于后续聚合（可选）
    .with_columns(
        # 生成每行A、B的笛卡尔积元素对数组
        cartesian_pairs=pl.struct("A", "B").arr.cartesian_product()
    )
    .explode("cartesian_pairs")  # 展开数组为单独行
    .with_columns(
        # 计算元素对的哈希值模100
        C=pl.hash(pl.col("cartesian_pairs")) % 100,
        # 若需计算元素乘积的哈希模100，替换为：
        # C=pl.hash(pl.col("cartesian_pairs").struct.field("A") * pl.col("cartesian_pairs").struct.field("B")) % 100
    )
    # 若需还原为原行结构（C列为列表），取消以下注释：
    # .group_by("row_id")
    # .agg(
    #     pl.all().exclude("cartesian_pairs"),
    #     C=pl.col("C")
    # )
    # .drop("row_id")
)

场景2：A、B为标量列（每行单个元素） #

如果A、B是标量，笛卡尔积即为每行的元素对，直接计算即可：

df2 = df.with_columns(
    # 元素对的哈希模100
    C=pl.hash(pl.struct("A", "B")) % 100,
    # 或乘积的哈希模100
    # C=pl.hash(pl.col("A") * pl.col("B")) % 100
)

性能优势说明 #

• arr.cartesian_product是Polars底层优化的向量化操作，完全避免Python级别的循环，比apply快数十倍。
• pl.hash是内置的向量化哈希函数，性能远高于Python手动计算哈希。
• 即使需要聚合回原行结构，group_by+agg的向量化操作也比逐行apply高效得多。

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