如何在Polars中筛选存储UUID对象的列

问题 #

我有一个Polars DataFrame，其中id列存储UUID类型数据。希望始终将这些数据以UUID对象形式保存，不在创建DataFrame前转换为字符串/整数等类型。创建df的过程正常，但尝试筛选id列时遇到错误。请问如何通过指定的UUID来筛选id列？

相关代码及报错信息如下：

import polars as pl
from uuid import UUID

df = pl.DataFrame(
    [
        {"message": "foo", "id": UUID("925afb28-f437-4c46-b9b2-ae9ebb7899a9")},
        {"message": "bar", "id": UUID("f9278c55-4c11-4b68-a1aa-de9d6086efb6")},
    ],
    schema={"message": pl.String, "id": pl.Object},
)

df.shape
>> (2, 2)

# 直接用UUID对象比较报错
df.filter(pl.col("id").eq(UUID("925afb28-f437-4c46-b9b2-ae9ebb7899a9")))
>> TypeError: invalid literal value: '925afb28-f437-4c46-b9b2-ae9ebb7899a9'

# 转成字符串比较也报错
df.filter(pl.col("id").eq(str(UUID("925afb28-f437-4c46-b9b2-ae9ebb7899a9"))))
>> ComputeError: cannot cast 'Object' type

解决方案 #

方法1：用pl.lit()包装UUID对象 #

Polars处理Object类型列的相等性比较时，需要将目标UUID用pl.lit()包装，确保Polars以对象类型而非字面量进行判断：

target_uuid = UUID("925afb28-f437-4c46-b9b2-ae9ebb7899a9")
result = df.filter(pl.col("id").eq(pl.lit(target_uuid)))
print(result)

输出结果：

shape: (1, 2)
┌─────────┬──────────────────────────────────────┐
│ message ┆ id                                   │
│ ---     ┆ ---                                  │
│ str     ┆ object                               │
╞═════════╪══════════════════════════════════════╡
│ foo     ┆ 925afb28-f437-4c46-b9b2-ae9ebb7899a9 │
└─────────┴──────────────────────────────────────┘

方法2：使用map_elements自定义比较逻辑 #

如果第一种方法不生效，可通过map_elements调用Python原生的对象相等性判断：

target_uuid = UUID("925afb28-f437-4c46-b9b2-ae9ebb7899a9")
result = df.filter(
    pl.col("id").map_elements(
        lambda x: x == target_uuid,
        return_dtype=pl.Boolean
    )
)
print(result)

关键说明 #

Polars的Object类型是通用容器，默认会尝试将字面量转换为列的类型，但UUID对象无法直接被转换为字面量。使用pl.lit()包装目标UUID后，Polars会将其识别为Object类型的字面量，从而正确执行对象间的相等性比较。

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