长 .when().then().when().then().otherwise() 链的替代方案

Question

是否有一些聪明的替代方法可以编写长的when().then().otherwise()链而不对值进行硬编码，请参见下面的示例：

\n

假设我们有以下数据框

\n

df = pl.DataFrame(\n    {\n        "Market":["AT", "AT", "DE", "DE", "CA", "DE", "UK", "US"],\n        "Number of Days":[1, 2, 3, 4, 3, 4, 2, 1],\n        \n    }\n)\n

\n

用户将一些条件定义为不同国家的字典

\n

params = {\n    "AT":{"Value": 1},\n    "DE":{"Value": 2},\n    "CA":{"Value": 3},\n    "UK":{"Value": 1},\n    "US":{"Value": 2}\n}\n

\n

然后我对国家/地区进行硬编码并使用 Polars .with_columns() 中的国家/地区，如下所示：

\n

(\n    df\n    .with_columns(\n        [\n            pl.when(pl.col("Market") == "AT").then(pl.col("Number of Days") + params["AT"]["Value"])\n            .when(pl.col("Market") == "DE").then(pl.col("Number of Days") + params["DE"]["Value"])\n            .when(pl.col("Market") == "CA").then(pl.col("Number of Days") + params["CA"]["Value"])\n            .when(pl.col("Market") == "UK").then(pl.col("Number of Days") + params["UK"]["Value"])\n            .when(pl.col("Market") == "US").then(pl.col("Number of Days") + params["US"]["Value"])\n            .otherwise(None)\n            .alias("New Column")\n        ]\n    )\n)\n

\n

有没有办法让我不对 .with_columns 中的国家/地区进行硬编码，而是以某种方式循环字典并根据提供的值创建表达式？\xc2\xa8

\n

我尝试了下面的内容，但它说我有重复的列名。

\n

exprs = []\nfor market, data in params.items():\n    condition = (pl.col("Market") == market)\n    result = (pl.col("Number of Days") + data["Value"])\n    expr = pl.when(condition).then(result)\n    exprs.append(expr)\n\ndf.with_columns(exprs)\n

\n

Answer 1

您可以使用map_dict

(
    df
        .with_columns((
            pl.col("Number of Days") + 
            pl.col("Market")
                .map_dict({x:y['Value'] for x,y in params.items()})
            )
            .alias("New Column")
        )
)

或者，您可以将 params 设置为其自己的 df，然后加入它们。map_dict工作原理是将你的 dict 转换为 df 然后加入它们，但根据你的参数来自哪里，这可能是首选。

paramsdf = pl.DataFrame({'Market':params.keys(), 
                         'Value':[x['Value'] for x in params.values()]})
(
    df.join(paramsdf, on="Market")
        .with_columns((pl.col("Number of Days")+ pl.col("Value")).alias("New Column"))
        .drop("Value")
)