使用单元格值的字典以并行方式更新极坐标数据框中的每个单元格？

Question

我有一个像这样的数据框：

test_df = pl.DataFrame(
    {
        "row_id": ["a", "b", "c", "d", "e", "f"],
        "s1": [None, 1, 2, 2, 33, 3],
        "s2": [None, 32, 12, 2, 0, None],
        "s3": [10, 20, 30, 40, 50, 60],
        "s4": [0, 34, None, 34, 50, 60],
        "s5": [10, 2, None, 123, 3, 432],
    }
)

以及像这样的单元格值的字典：

import random
rows = count_table["row_id"].to_list()
r_idx = {row: i for i, row in enumerate(rows)}
columns = [col for col in count_table.columns if col != "row_id"]
adjustments = {
        (i, j): random.randint(3, 9) for i in rows for j in columns if count_table[r_idx[i], j] is not None
    }

我想找到更新每个单元格值的最有效方法。

谢谢！

目前我正在这样做：

for (i, j), var in adjustments.items():
            count_table = count_table.with_columns(
                pl.when(pl.col(id_col) == i).then(count_table[j] + var.value()).otherwise(count_table[j]).alias(j)
            )

但是我觉得还有更快的方法，尤其是Polars的并行计算能力？

Answer 1

data = [list(k) + [v] for k,v in adjustments.items()]

adjustments_df = (
    pl
    .DataFrame(data, orient="row", schema=["row_id","column_name","value"])
    .pivot(on="column_name",values="value")
)

┌────────┬─────┬──────┬──────┬──────┬──────┐
│ row_id ┆ s3  ┆ s4   ┆ s5   ┆ s1   ┆ s2   │
│ ---    ┆ --- ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ str    ┆ i64 ┆ i64  ┆ i64  ┆ i64  ┆ i64  │
╞════════╪═════╪══════╪══════╪══════╪══════╡
│ a      ┆ 4   ┆ 7    ┆ 8    ┆ null ┆ null │
│ b      ┆ 7   ┆ 5    ┆ 8    ┆ 3    ┆ 4    │
│ c      ┆ 6   ┆ null ┆ null ┆ 6    ┆ 4    │
│ d      ┆ 8   ┆ 8    ┆ 8    ┆ 8    ┆ 6    │
│ e      ┆ 9   ┆ 8    ┆ 5    ┆ 5    ┆ 7    │
│ f      ┆ 5   ┆ 6    ┆ 9    ┆ 4    ┆ null │
└────────┴─────┴──────┴──────┴──────┴──────┘

cols = [x for x in test_df.schema if x != "row_id"]

(
    test_df
    .join(adjustments_df, on="row_id", how="left")
    .with_columns(
        pl.coalesce(pl.col(f"{col}_right"), pl.col(col)).alias(col)
        for col in cols
    )
    .select(pl.col.row_id, pl.col(cols))
)

┌────────┬──────┬──────┬─────┬──────┬──────┐
│ row_id ┆ s1   ┆ s2   ┆ s3  ┆ s4   ┆ s5   │
│ ---    ┆ ---  ┆ ---  ┆ --- ┆ ---  ┆ ---  │
│ str    ┆ i64  ┆ i64  ┆ i64 ┆ i64  ┆ i64  │
╞════════╪══════╪══════╪═════╪══════╪══════╡
│ a      ┆ null ┆ null ┆ 4   ┆ 7    ┆ 8    │
│ b      ┆ 3    ┆ 4    ┆ 7   ┆ 5    ┆ 8    │
│ c      ┆ 6    ┆ 4    ┆ 6   ┆ null ┆ null │
│ d      ┆ 8    ┆ 6    ┆ 8   ┆ 8    ┆ 8    │
│ e      ┆ 5    ┆ 7    ┆ 9   ┆ 8    ┆ 5    │
│ f      ┆ 4    ┆ null ┆ 5   ┆ 6    ┆ 9    │
└────────┴──────┴──────┴─────┴──────┴──────┘