Polars中利用列值作为字典键进行数据过滤的策略与优化

在polars中，直接使用列表达式作为python字典的键会导致`typeerror: unhashable type: 'expr'`。本文将深入探讨两种有效解决此问题的方法：一是利用`map_elements`进行行级别转换，这种方法直观但效率较低；二是将嵌套字典扁平化为polars dataframe，并通过`join`操作实现高效过滤，这是处理大规模数据的推荐方案。文章将详细阐述每种方法的实现细节、适用场景及其性能考量。

在Polars进行数据处理时，我们有时会遇到需要根据DataFrame中的列值去查询一个外部Python字典的情况，特别是当字典是多层嵌套时。例如，尝试使用pl.col("cliente")和pl.col("cluster")作为nested_dict的键来过滤数据，如下所示：

df_x = (
    df_x
    .filter(pl.col("score") == nested_dict[pl.col("cliente")][pl.col("cluster")])
)

这段代码会抛出TypeError: unhashable type: 'Expr'错误。这是因为pl.col("cliente")和pl.col("cluster")返回的是Polars表达式（Expr对象），而不是实际的列值。Python字典的键必须是可哈希的（hashable），而Expr对象不可哈希，因此无法直接用作字典键进行查询。

为了解决这个问题，我们可以采用两种主要策略：一种是利用map_elements在行级别应用Python函数，另一种是将外部字典转换为Polars DataFrame并进行连接（join）操作。

策略一：使用 map_elements 进行行级别转换

map_elements方法允许我们对DataFrame中的元素应用一个Python函数。通过这种方式，我们可以在函数内部解析列值，并使用这些实际值来查询Python字典。

实现步骤

1. 组合相关列： 使用pl.struct()将需要作为字典键的列（例如'cliente'和'cluster'）组合成一个结构体（Struct）。
2. 应用 map_elements： 对这个结构体应用map_elements函数。该函数会接收每一行的结构体作为输入，其中包含各列的实际值。
3. 字典查询： 在map_elements内部的lambda函数中，使用结构体中的实际值来查询nested_dict。

示例代码

import polars as pl

# 示例数据和嵌套字典
df_x = pl.DataFrame({
    "cliente": ["A", "B", "A", "C"],
    "cluster": ["X", "Y", "Z", "X"],
    "score": [10, 20, 30, 40]
})

nested_dict = {
    "A": {"X": 10, "Z": 25},
    "B": {"Y": 20},
    "C": {"X": 40}
}

# 使用 map_elements 进行过滤
df_filtered_map = (
    df_x
    .filter(
        pl.col('score').eq(
            pl.struct('cliente', 'cluster')
                .map_elements(lambda x: (
                    nested_dict.get(x['cliente'], {}).get(x['cluster']) # 使用 .get 避免 KeyError
                ), return_dtype=pl.Int64) # 指定返回类型
        )
    )
)

print("使用 map_elements 过滤后的 DataFrame:")
print(df_filtered_map)

注意事项

• 性能： map_elements会强制Polars将数据传递给Python函数进行处理，这会引入Python解释器的开销，通常比纯Polars的向量化操作效率低。对于大规模数据集，这可能成为性能瓶颈。
• 错误处理： 在lambda函数中，建议使用字典的.get()方法来安全地访问键，以防止当cliente或cluster的组合在nested_dict中不存在时引发KeyError。
• 返回类型： 必须为map_elements指定return_dtype，否则Polars可能无法推断出正确的列类型。

策略二：扁平化字典并进行连接（Join）

更高效且Polars-idiomatic 的方法是将嵌套的Python字典转换为一个Polars DataFrame，然后通过join操作将其与主DataFrame连接起来，最后再进行过滤。这种方法将字典查询转换为Polars的向量化操作，从而显著提高性能。

怪兽AI数字人

数字人短视频创作，数字人直播，实时驱动数字人

44

查看详情

实现步骤

1. 扁平化嵌套字典： 将nested_dict转换为一个包含cliente、cluster和cluster_value（即对应的分数）的Polars DataFrame。
2. 执行连接操作： 使用join方法将主DataFrame与扁平化的字典DataFrame连接起来，连接键为cliente和cluster。
3. 进行过滤： 连接后，cluster_value列将可用，可以直接用于过滤条件。

示例代码

import polars as pl

# 示例数据和嵌套字典
df_x = pl.DataFrame({
    "cliente": ["A", "B", "A", "C", "D"],
    "cluster": ["X", "Y", "Z", "X", "Y"],
    "score": [10, 20, 30, 40, 50]
})

nested_dict = {
    "A": {"X": 10, "Z": 25},
    "B": {"Y": 20},
    "C": {"X": 40}
}

# 1. 扁平化嵌套字典为 Polars DataFrame
df_nested_prelim = pl.from_dict(nested_dict)

df_nested_parts = []
for col_name in df_nested_prelim.columns:
    df_nested_parts.append(
        df_nested_prelim.lazy()
        .select(col_name).unnest(col_name) # 展开内部字典为列
        .unpivot(
            on=[], # 不指定on，对所有非id列进行unpivot
            index=[], # 没有id列，所有列都参与unpivot
            variable_name='cluster', 
            value_name='cluster_value'
        )
        .with_columns(cliente=pl.lit(col_name)) # 添加原始的cliente名称
    )

df_nested = pl.concat(df_nested_parts).collect()

# 2. 移除可能产生的null值（如果原始字典中没有对应的cluster）
df_nested = df_nested.filter(pl.col("cluster_value").is_not_null())

print("\n扁平化后的字典 DataFrame:")
print(df_nested)

# 3. 执行连接并过滤
df_filtered_join = (
    df_x
    .join(df_nested, on=['cliente', 'cluster'], how='inner') # 使用inner join确保只匹配存在的值
    .filter(pl.col('score') == pl.col('cluster_value'))
    .select(df_x.columns) # 仅保留原始 df_x 的列
)

print("\n使用 join 过滤后的 DataFrame:")
print(df_filtered_join)

扁平化字典的详细解释

上述扁平化字典的代码可能看起来有些复杂，我们来逐步解析：

1. df_nested_prelim = pl.from_dict(nested_dict): 将nested_dict转换为Polars DataFrame。此时，nested_dict的顶层键（如"A", "B", "C"）会成为DataFrame的列名，而它们对应的值（内层字典）会成为这些列中的结构体（Struct）。 例如，df_nested_prelim可能看起来像：

   shape: (1, 3)
   ┌───────────┬───────────┬───────────┐
   │ A         ┆ B         ┆ C         │
   │ struct[2] ┆ struct[1] ┆ struct[1] │
   ╞═══════════╪═══════════╪═══════════╡
   │ {10,25}   ┆ {20}      ┆ {40}      │
   └───────────┴───────────┴───────────┘

   登录后复制

2. for col_name in df_nested_prelim.columns:: 遍历df_nested_prelim中的每一列（即原始的cliente名称）。

3. select(col_name).unnest(col_name): 选择当前列，并将其解嵌套。例如，对于列"A"，它包含{"X": 10, "Z": 25}这个结构体。unnest("A")会将其展开为两列："X"和"Z"。

   shape: (1, 2)
   ┌─────┬─────┐
   │ X   ┆ Z   │
   │ i64 ┆ i64 │
   ╞═════╪═════╡
   │ 10  ┆ 25  │
   └─────┴─────┘

   登录后复制

4. unpivot(on=[], index=[], variable_name='cluster', value_name='cluster_value'): 将宽格式的DataFrame（X, Z作为列）转换为长格式。variable_name指定了新的列名，用于存放原始的列名（X, Z），value_name指定了存放原始列值（10, 25）的列名。 结果会是：

   shape: (2, 2)
   ┌─────────┬───────────────┐
   │ cluster ┆ cluster_value │
   │ str     ┆ i64           │
   ╞═════════╪═══════════════╡
   │ X       ┆ 10            │
   │ Z       ┆ 25            │
   └─────────┴───────────────┘

   登录后复制

5. with_columns(cliente=pl.lit(col_name)): 添加一个名为cliente的新列，其值就是当前循环的原始cliente名称（例如"A"）。

   shape: (2, 3)
   ┌─────────┬───────────────┬─────────┐
   │ cluster ┆ cluster_value ┆ cliente │
   │ str     ┆ i64           ┆ str     │
   ╞═════════╪═══════════════╪═════════╡
   │ X       ┆ 10            ┆ A       │
   │ Z       ┆ 25            ┆ A       │
   └─────────┴───────────────┴─────────┘

   登录后复制

6. pl.concat(df_nested_parts).collect(): 将所有cliente循环生成的DataFrame片段拼接在一起，形成最终的扁平化字典DataFrame。

优势

• 高性能： join和filter操作都是Polars的高度优化和向量化操作，能够充分利用多核CPU，处理大规模数据集时效率远高于map_elements。
• Polars范式： 这种方法更符合Polars的设计哲学，即尽量将操作保持在DataFrame级别，避免Python循环和UDF。

总结

当需要在Polars中根据列值查询外部Python字典时，直接使用列表达式作为字典键是不可行的。

• map_elements方法 提供了一种快速实现的方式，尤其适用于数据量较小或逻辑复杂难以用Polars表达式表达的场景。但其性能开销较大。
• 扁平化字典并进行join的方法 是处理大规模数据集时的首选方案。它将字典查询转换为高效的Polars DataFrame操作，提供了卓越的性能和可扩展性。

在实际应用中，应根据数据规模和性能要求选择最合适的方法。对于生产环境和大数据场景，强烈推荐采用扁平化字典并进行连接的策略。

以上就是Polars中利用列值作为字典键进行数据过滤的策略与优化的详细内容，更多请关注php中文网其它相关文章！