本文探讨了在polars dataframe中，如何解决直接使用`expr`作为字典键导致`typeerror`的问题。我们提供了两种解决方案：一种是使用`map_elements`结合`pl.struct`实现直接但效率较低的列值到字典键映射；另一种是推荐的优化方案，通过将嵌套字典扁平化为polars dataframe，然后利用高效的`join`操作进行数据筛选，从而显著提升性能。
在Polars中进行数据处理时，有时我们需要根据DataFrame中某列或多列的值去查询一个外部的Python字典，并将查询结果用于筛选。然而，直接尝试将Polars的表达式（pl.col(...)）作为Python字典的键会导致TypeError: unhashable type: 'Expr'错误，因为Expr对象本身是不可哈希的，无法直接作为字典键。本文将详细介绍如何优雅地解决这一问题，并提供两种不同的实现方案及其优缺点。
问题阐述
假设我们有一个Polars DataFrame df_x 和一个嵌套的Python字典 nested_dict。我们希望根据 df_x 中的 cliente 和 cluster 列的值，从 nested_dict 中获取对应的值，然后用这个值来筛选 df_x 中 score 列的数据。
错误的尝试示例：
import polars as pl
# 示例数据和字典
df_x = pl.DataFrame({
"cliente": ["A", "A", "B", "B", "C"],
"cluster": ["X", "Y", "X", "Y", "X"],
"score": [10, 20, 30, 40, 50]
})
nested_dict = {
"A": {"X": 10, "Y": 25},
"B": {"X": 35, "Y": 40},
"C": {"X": 50, "Y": 55}
}
# 错误的尝试，会导致 TypeError: unhashable type: 'Expr'
try:
df_x_filtered = (
df_x
.filter(pl.col("score") == nested_dict[pl.col("cliente")][pl.col("cluster")])
)
except TypeError as e:
print(f"捕获到错误: {e}")上述代码尝试在filter表达式内部直接使用pl.col("cliente")和pl.col("cluster")作为字典键，这在Polars的表达式上下文中是无效的，因为pl.col(...)返回的是一个表达式对象，而不是实际的列值。
解决方案一：使用 map_elements 实现列值到字典键的映射
map_elements方法允许我们将Polars DataFrame中的结构化数据（例如，由多列组成的结构体）传递给一个Python函数进行处理。通过这种方式，我们可以在Python函数内部解析出列的实际值，并用它们来查询字典。
import polars as pl
# 示例数据和字典（同上）
df_x = pl.DataFrame({
"cliente": ["A", "A", "B", "B", "C"],
"cluster": ["X", "Y", "X", "Y", "X"],
"score": [10, 20, 30, 40, 50]
})
nested_dict = {
"A": {"X": 10, "Y": 25},
"B": {"X": 35, "Y": 40},
"C": {"X": 50, "Y": 55}
}
# 解决方案一：使用 map_elements
df_x_filtered_map = (
df_x
.filter(
pl.col('score').eq(
pl.struct('cliente','cluster') # 将多列组合成一个结构体
.map_elements(lambda x: ( # 对每个结构体元素应用Python函数
nested_dict[x['cliente']][x['cluster']] # 在Python函数内部解析值并查询字典
), return_dtype=pl.Int64 # 指定返回数据类型
)
)
)
)
print("使用 map_elements 过滤后的结果:")
print(df_x_filtered_map)说明：
- pl.struct('cliente','cluster') 将 cliente 和 cluster 两列组合成一个结构体（struct）列。
- .map_elements(lambda x: ..., return_dtype=...) 对这个结构体列的每一个元素（行）应用一个Python lambda函数。x 在这里是一个Python字典，其键是列名，值是当前行的列值。
- 在lambda函数内部，我们可以安全地使用 x['cliente'] 和 x['cluster'] 来访问字典 nested_dict。
- return_dtype 参数是必需的，它告诉Polars map_elements 返回的数据类型。
注意事项：
- 效率问题： map_elements 会在Polars的内部优化器和Python解释器之间进行数据传递，这引入了Python的用户定义函数（UDF）开销。对于大型数据集，这种方法可能不是最有效的。
- 类型指定： 必须为 map_elements 指定 return_dtype，以确保Polars能够正确处理返回结果。
解决方案二：优化方案——扁平化字典并使用 join 操作
为了获得更好的性能，尤其是处理大规模数据时，推荐的方法是将外部的Python字典转换为一个Polars DataFrame，然后使用Polars原生的 join 操作来合并数据并进行筛选。这种方法能够充分利用Polars的向量化和并行处理能力。
步骤一：扁平化 nested_dict 为 Polars DataFrame
我们需要将 nested_dict 转换为一个包含 cliente、cluster 和 cluster_value 三列的Polars DataFrame。
import polars as pl
# 示例数据和字典（同上）
df_x = pl.DataFrame({
"cliente": ["A", "A", "B", "B", "C"],
"cluster": ["X", "Y", "X", "Y", "X"],
"score": [10, 20, 30, 40, 50]
})
nested_dict = {
"A": {"X": 10, "Y": 25},
"B": {"X": 35, "Y": 40},
"C": {"X": 50, "Y": 55}
}
# 扁平化 nested_dict 为 Polars DataFrame
df_nested_prelim = pl.from_dict(nested_dict) # 转换为初步的DataFrame
# print("初步转换的 df_nested_prelim:")
# print(df_nested_prelim)
df_nested_parts = []
for col_name in df_nested_prelim.columns:
df_nested_parts.append(
df_nested_prelim.lazy()
.select(pl.col(col_name)).unnest(col_name) # 展开嵌套结构
.unpivot(variable_name='cluster', value_name='cluster_value') # 将列名转换为cluster，值转换为cluster_value
.with_columns(cliente=pl.lit(col_name)) # 添加cliente列，值为当前外部键
)
df_nested = pl.concat(df_nested_parts).collect()
print("\n扁平化后的 df_nested:")
print(df_nested)说明：
- pl.from_dict(nested_dict) 将字典转换为一个初步的DataFrame，其中外部键（"A", "B", "C"）成为列名，内部字典成为列表中的结构体。
- 我们遍历 df_nested_prelim 的每一列（即 cliente），然后：
- .unnest(col_name) 展开该列中的嵌套结构。
- .unpivot(variable_name='cluster', value_name='cluster_value') 将内部字典的键（"X", "Y"）转换为 cluster 列的值，将对应的值转换为 cluster_value 列。
- .with_columns(cliente=pl.lit(col_name)) 添加 cliente 列，其值为当前外部键。
- pl.concat(df_nested_parts).collect() 将所有部分DataFrame合并成最终的扁平化DataFrame df_nested。
步骤二：使用 join 进行高效筛选
有了扁平化的 df_nested，我们就可以将其与原始DataFrame df_x 进行 join 操作，然后直接进行筛选。
# 解决方案二：使用 join 进行高效筛选
df_x_filtered_join = (
df_x
.join(df_nested, on=['cliente','cluster'], how='inner') # 根据 cliente 和 cluster 进行内连接
.filter(pl.col('score')==pl.col('cluster_value')) # 筛选 score 等于 cluster_value 的行
.select(df_x.columns) # 仅保留原始 df_x 的列，移除 join 引入的 cluster_value
)
print("\n使用 join 过滤后的结果:")
print(df_x_filtered_join)说明：
- df_x.join(df_nested, on=['cliente','cluster'], how='inner') 通过 cliente 和 cluster 列将 df_x 与 df_nested 进行内连接。这意味着只有在两边都存在的 (cliente, cluster) 组合才会被保留，并且 df_nested 中的 cluster_value 列会被添加到 df_x 的行中。
- .filter(pl.col('score')==pl.col('cluster_value')) 接着使用Polars原生的筛选操作，比较 score 列和连接后得到的 cluster_value 列。
- .select(df_x.columns) 这一步是可选的，用于在筛选完成后，将结果DataFrame的列恢复到与原始 df_x 相同的结构，移除 join 操作引入的辅助列 cluster_value。
总结与选择建议
在Polars中利用列值作为字典键进行筛选时，直接使用Expr对象是不可行的。我们提供了两种有效的解决方案：
-
map_elements 方法：
- 优点： 语法上更接近于直接在Python中操作，对于小规模数据集或需要复杂Python逻辑的场景较为方便。
- 缺点： 涉及Python UDF，效率相对较低，不适合大规模数据集。
-
扁平化字典并 join 方法：
- 优点： 充分利用Polars的优化查询引擎，性能卓越，尤其适合大规模数据集。
- 缺点： 需要额外的步骤将外部字典转换为Polars DataFrame，代码相对复杂一些。
选择建议：
- 对于数据量较小，或者字典查询逻辑非常复杂难以用Polars表达式表达时，可以考虑使用 map_elements。
- 对于绝大多数生产环境和大数据场景，强烈推荐将外部字典扁平化为Polars DataFrame，并使用 join 操作进行处理。这种方法虽然初期设置稍显复杂，但能带来显著的性能提升。
理解Polars的核心理念——尽可能使用其原生的、向量化的操作，避免Python UDF的频繁调用，是编写高效Polars代码的关键。
