Polars实现类似Pandas df.loc的行更新/插入操作

场景与需求 #

现有如下Polars DataFrame：

┌────────────┬─────────┬────────────┬─────────┬──────────┬─────────────┐
│ gDate      ┆ pl      ┆ redemption ┆ pc      ┆ ppc      ┆ tval        │
│ ---        ┆ ---     ┆ ---        ┆ ---     ┆ ---      ┆ ---         │
│ date       ┆ f64     ┆ f64        ┆ f64     ┆ f64      ┆ f64         │
╞════════════╪═════════╪════════════╪═════════╪══════════╪═════════════╡
│ 2024-01-30 ┆ 10069.0 ┆ 10068.0    ┆ 10068.0 ┆ 0.07     ┆ 1492.033    │
│ 2024-01-31 ┆ 10082.0 ┆ 10075.0    ┆ 10082.0 ┆ 0.14     ┆ 1236.318    │
│ 2024-02-03 ┆ 10095.0 ┆ 10095.0    ┆ 10095.0 ┆ 0.13     ┆ 2266.253    │
│ 2024-02-04 ┆ 10102.0 ┆ 10102.0    ┆ 10103.0 ┆ 0.08     ┆ 1418.583    │
│ 2024-02-05 ┆ 10110.0 ┆ 10109.0    ┆ 10109.0 ┆ 0.06     ┆ 1940.722    │
│ …          ┆ …       ┆ …          ┆ …       ┆ …        ┆ …           │
│ 2024-03-26 ┆ 10113.0 ┆ 10112.0    ┆ 10112.0 ┆ 0.07     ┆ 1268.171    │
│ 2024-03-27 ┆ 10126.0 ┆ 10119.0    ┆ 10126.0 ┆ 0.14     ┆ 1240.438    │
│ 2024-03-30 ┆ 10149.0 ┆ 10141.0    ┆ 10148.0 ┆ 0.22     ┆ 1317.312    │
│ 2024-04-02 ┆ 10163.0 ┆ 10162.0    ┆ 10162.0 ┆ 0.14     ┆ 1675.402    │
│ 2024-04-03 ┆ 10177.0 ┆ 10169.0    ┆ 10176.0 ┆ 0.137768 ┆ 1546.426394 │
└────────────┴─────────┴────────────┴─────────┴──────────┴─────────────┘

需要实现：

• 为指定日期的redemption列赋值，若该日期已存在则更新对应行的redemption值
• 若日期不存在，则新增一行，仅redemption列设为新值，其余列设为null

在Pandas中可以直接用df.loc[date, 'redemption'] = new_value实现，但Polars中直接使用类似索引赋值会抛出polars.exceptions.OutOfBoundsError: indices are out of bounds错误，需要采用符合Polars风格的实现方式。

当前实现代码 #

from polars import *
from datetime import date

df = from_repr("""
┌────────────┬─────────┬────────────┬─────────┬──────────┬─────────────┐
│ gDate      ┆ pl      ┆ redemption ┆ pc      ┆ ppc      ┆ tval        │
│ ---        ┆ ---     ┆ ---        ┆ ---     ┆ ---      ┆ ---         │
│ date       ┆ f64     ┆ f64        ┆ f64     ┆ f64      ┆ f64         │
╞════════════╪═════════╪════════════╪═════════╪══════════╪═════════════╡
│ 2024-01-30 ┆ 10069.0 ┆ 10068.0    ┆ 10068.0 ┆ 0.07     ┆ 1492.033    │
│ 2024-01-31 ┆ 10082.0 ┆ 10075.0    ┆ 10082.0 ┆ 0.14     ┆ 1236.318    │
│ 2024-02-03 ┆ 10095.0 ┆ 10095.0    ┆ 10095.0 ┆ 0.13     ┆ 2266.253    │
│ 2024-02-04 ┆ 10102.0 ┆ 10102.0    ┆ 10103.0 ┆ 0.08     ┆ 1418.583    │
│ 2024-02-05 ┆ 10110.0 ┆ 10109.0    ┆ 10109.0 ┆ 0.06     ┆ 1940.722    │
│ …          ┆ …       ┆ …          ┆ …       ┆ …        ┆ …           │
│ 2024-03-26 ┆ 10113.0 ┆ 10112.0    ┆ 10112.0 ┆ 0.07     ┆ 1268.171    │
│ 2024-03-27 ┆ 10126.0 ┆ 10119.0    ┆ 10126.0 ┆ 0.14     ┆ 1240.438    │
│ 2024-03-30 ┆ 10149.0 ┆ 10141.0    ┆ 10148.0 ┆ 0.22     ┆ 1317.312    │
│ 2024-04-02 ┆ 10163.0 ┆ 10162.0    ┆ 10162.0 ┆ 0.14     ┆ 1675.402    │
│ 2024-04-03 ┆ 10177.0 ┆ 10169.0    ┆ 10176.0 ┆ 0.137768 ┆ 1546.426394 │
└────────────┴─────────┴────────────┴─────────┴──────────┴─────────────┘""")
target_date = date(2024, 4, 3)
new_redemption = 7777

df = df.join(
    DataFrame([[target_date], [new_redemption]]),
    left_on='gDate',
    right_on='column_0',
    how='outer_coalesce',
).with_columns(
    when(col('column_1').is_null())
    .then(col('redemption'))
    .otherwise(col('column_1'))
    .alias('redemption')
).drop('column_1', 'column_0')

print(df)

注：原代码未对新生成的redemption列做别名，也未删除临时列column_0，此处做了修正以保证结果列结构正确。

更简洁的Polars风格实现 #

可以通过构造单行更新数据，与原表合并后按日期分组保留最新值的方式实现，逻辑更直观：

from polars import *
from datetime import date

df = from_repr("""
┌────────────┬─────────┬────────────┬─────────┬──────────┬─────────────┐
│ gDate      ┆ pl      ┆ redemption ┆ pc      ┆ ppc      ┆ tval        │
│ ---        ┆ ---     ┆ ---        ┆ ---     ┆ ---      ┆ ---         │
│ date       ┆ f64     ┆ f64        ┆ f64     ┆ f64      ┆ f64         │
╞════════════╪═════════╪════════════╪═════════╪══════════╪═════════════╡
│ 2024-01-30 ┆ 10069.0 ┆ 10068.0    ┆ 10068.0 ┆ 0.07     ┆ 1492.033    │
│ 2024-01-31 ┆ 10082.0 ┆ 10075.0    ┆ 10082.0 ┆ 0.14     ┆ 1236.318    │
│ 2024-02-03 ┆ 10095.0 ┆ 10095.0    ┆ 10095.0 ┆ 0.13     ┆ 2266.253    │
│ 2024-02-04 ┆ 10102.0 ┆ 10102.0    ┆ 10103.0 ┆ 0.08     ┆ 1418.583    │
│ 2024-02-05 ┆ 10110.0 ┆ 10109.0    ┆ 10109.0 ┆ 0.06     ┆ 1940.722    │
│ …          ┆ …       ┆ …          ┆ …       ┆ …        ┆ …           │
│ 2024-03-26 ┆ 10113.0 ┆ 10112.0    ┆ 10112.0 ┆ 0.07     ┆ 1268.171    │
│ 2024-03-27 ┆ 10126.0 ┆ 10119.0    ┆ 10126.0 ┆ 0.14     ┆ 1240.438    │
│ 2024-03-30 ┆ 10149.0 ┆ 10141.0    ┆ 10148.0 ┆ 0.22     ┆ 1317.312    │
│ 2024-04-02 ┆ 10163.0 ┆ 10162.0    ┆ 10162.0 ┆ 0.14     ┆ 1675.402    │
│ 2024-04-03 ┆ 10177.0 ┆ 10169.0    ┆ 10176.0 ┆ 0.137768 ┆ 1546.426394 │
└────────────┴─────────┴────────────┴─────────┴──────────┴─────────────┘""")
target_date = date(2024, 4, 3)
new_redemption = 7777

# 构造单行更新数据
update_df = DataFrame({
    'gDate': [target_date],
    'redemption': [new_redemption]
})

# 合并原表与更新表，按gDate分组保留最新值
df = (
    concat([df, update_df])
    .group_by('gDate', maintain_order=True)
    .agg(
        pl.col('pl').first(),
        pl.col('redemption').last(),  # 取最后一条数据，即更新后的值
        pl.col('pc').first(),
        pl.col('ppc').first(),
        pl.col('tval').first()
    )
)

print(df)

这种方式利用Polars的向量化操作特性，通过concat合并数据、group_by聚合筛选，逻辑清晰且符合Polars的设计风格。

内容的提问来源于stack exchange，提问作者AXO