如何在Polars中基于List[f64]类型键连接两个DataFrame？

问题重现 #

尝试以List[pl.Float64]类型的coords列作为键对两个Polars DataFrame做左连接时，触发ComputeError，因为Polars目前未实现列表类型的哈希左连接：

import polars as pl

df = pl.DataFrame({"id": [1.0, 2.0], "coords": [(1.0, 2.0), (3.0, 4.0)]})
df_allowed = pl.DataFrame(
    {
        "coords": [(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)],
        "people": ["Alice", "Bob", "Charlie"],
    }
)

df.join(df_allowed, how="left", on="coords")
# 报错：ComputeError: not yet implemented: Hash Left Join between list[f64] and list[f64]

期望得到的结果：

df_expected = pl.DataFrame(
    {"id": [1.0, 2.0], "coords": [(1.0, 2.0), (3.0, 4.0)], "people": ["Alice", "Bob"]}
)

解决方案 #

以下两种方案均适用于数万条数据的场景，且能保留原coords列的列表格式：

方案1：将列表列转换为结构体（Struct）连接

Polars支持以结构体作为连接键，先把列表列转为结构体，连接完成后可再转回列表格式：

# 将列表列转为结构体（字段命名为x、y）
df_struct = df.with_columns(pl.col("coords").list.to_struct(fields=["x", "y"]))
df_allowed_struct = df_allowed.with_columns(pl.col("coords").list.to_struct(fields=["x", "y"]))

# 执行左连接
result = df_struct.join(df_allowed_struct, how="left", on="coords")

# 可选：将结构体转回列表列，恢复原格式
result = result.with_columns(pl.col("coords").struct.to_list())

print(result)

方案2：拆分列表为单独字段连接

把coords列表拆分为独立的数值列作为连接键，连接后再合并回列表格式：

# 拆分列表为x、y两个字段
df_split = df.with_columns(
    x=pl.col("coords").list[0],
    y=pl.col("coords").list[1]
)
df_allowed_split = df_allowed.with_columns(
    x=pl.col("coords").list[0],
    y=pl.col("coords").list[1]
)

# 用x、y字段执行左连接
result = df_split.join(df_allowed_split, how="left", on=["x", "y"])

# 合并x、y字段回列表列，并删除临时字段
result = result.with_columns(coords=pl.struct(["x", "y"]).struct.to_list()).drop(["x", "y"])

print(result)

两种方案执行后均能得到符合预期的结果，且Polars的向量化处理能力能保证数万条数据的处理效率。

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