使用极坐标自定义每行转换

Question

我有一个数据集，其中包含一个工作解决方案，该解决方案使用带有 Polars 的 for 循环。但是，我想使用内置的 Polars 函数（例如 map_elements）来应用基于 id 列创建可变长度结构的函数。

工作解决方案：

import polars as pl
pl.Config.set_fmt_str_lengths(1000)

import warnings
warnings.filterwarnings("ignore")

df = pl.DataFrame({
    "id": [1, 1, 2, 3],
    "dog": ["labrador", "labrador", "boxer", "airedale"],
    "age": [2, 2, 3, 4],
    "owner": ["Jim", "Paul", "Kim", "Lynne"],
    "spots": [None, None, True, None],
    "food": ["Raw", "Kibble", None, "Raw"],
    "leash": [True, False, None, True]
})

id_struct_dict = {
    1: ["owner", "food"],
    2: ["owner", "spots", "food"],
    3: ["owner", "food", "leash"]
}

run = 0
for id, lst in id_struct_dict.items():
    print(id)
    struct_expr = pl.struct(lst).map_elements(lambda x: str(x)).alias("struct")
    if run == 0:
        df = df.with_columns([pl.when(pl.col("id") == id).then(struct_expr).otherwise(None)])
        run += 1
    else:
        df = df.with_columns([pl.when(pl.col("id") == id).then(struct_expr).otherwise(pl.col('struct'))])

df

我真的不喜欢我的解决方案，但这是我能够达到我想要的结果的唯一方法。我尝试通过不同的方法映射函数/udf，但最终的结构案例将应用于所有行，尽管 ID 有所不同。

任何想法/帮助将不胜感激。谢谢你。

Answer 1

您可以将多个when/then表达式传递给

pl.coalesce

以生成单列结果。

如果 JSON 可接受，可以使用

.struct.json_encode()

对结构进行“字符串化”。

df.with_columns(
   pl.coalesce(
      pl.when(pl.col("id") == _id)
        .then(pl.struct(cols).struct.json_encode())
      for _id, cols in id_struct_dict.items()
   )
   .alias("struct")
)

shape: (4, 8)
┌─────┬──────────┬─────┬───────┬───────┬────────┬───────┬─────────────────────────────────────────────┐
│ id  ┆ dog      ┆ age ┆ owner ┆ spots ┆ food   ┆ leash ┆ struct                                      │
│ --- ┆ ---      ┆ --- ┆ ---   ┆ ---   ┆ ---    ┆ ---   ┆ ---                                         │
│ i64 ┆ str      ┆ i64 ┆ str   ┆ bool  ┆ str    ┆ bool  ┆ str                                         │
╞═════╪══════════╪═════╪═══════╪═══════╪════════╪═══════╪═════════════════════════════════════════════╡
│ 1   ┆ labrador ┆ 2   ┆ Jim   ┆ null  ┆ Raw    ┆ true  ┆ {"owner":"Jim","food":"Raw"}                │
│ 1   ┆ labrador ┆ 2   ┆ Paul  ┆ null  ┆ Kibble ┆ false ┆ {"owner":"Paul","food":"Kibble"}            │
│ 2   ┆ boxer    ┆ 3   ┆ Kim   ┆ true  ┆ null   ┆ null  ┆ {"owner":"Kim","spots":true,"food":null}    │
│ 3   ┆ airedale ┆ 4   ┆ Lynne ┆ null  ┆ Raw    ┆ true  ┆ {"owner":"Lynne","food":"Raw","leash":true} │
└─────┴──────────┴─────┴───────┴───────┴────────┴───────┴─────────────────────────────────────────────┘