Polars Pivot 在求和时将空值视为 0

Question

我有这个代码：

import polars as pl

pl.DataFrame({
    'label':   ['AA', 'CC', 'BB', 'AA', 'CC'],
    'account': ['EU', 'US', 'US', 'EU', 'EU'],
    'qty':     [1.5,  43.2, None, None, 18.9]})\
  .pivot('account', index='label', aggregate_function='sum')

这给出了

shape: (3, 3)
┌───────┬──────┬──────┐
│ label ┆ EU   ┆ US   │
│ ---   ┆ ---  ┆ ---  │
│ str   ┆ f64  ┆ f64  │
╞═══════╪══════╪══════╡
│ AA    ┆ 1.5  ┆ null │
│ CC    ┆ 18.9 ┆ 43.2 │
│ BB    ┆ null ┆ 0.0  │
└───────┴──────┴──────┘

现在，当原始数据中有任何

null

值时，我希望数据透视表在相应的单元格中显示

null

。但是，AA-EU 显示 1.5（但应为空），BB-US 显示 0.0（但也应为空）。

我尝试使用

aggregate_function=lambda col: pl.when(col.has_nulls())\
                                 .then(pl.lit(None, dtype=pl.Float64))\
                                 .otherwise(pl.sum(col))

但它会出错

AttributeError: 'function' object has no attribute '_pyexpr'

。

我该如何解决这个问题？

Answer 1

您可以使用 Polars 表达式作为聚合函数：

df.pivot('account', index='label', aggregate_function=
    pl.when(~pl.element().has_nulls()).then(pl.element().sum()))

shape: (3, 3)
┌───────┬──────┬──────┐
│ label ┆ EU   ┆ US   │
│ ---   ┆ ---  ┆ ---  │
│ str   ┆ f64  ┆ f64  │
╞═══════╪══════╪══════╡
│ AA    ┆ null ┆ null │
│ CC    ┆ 18.9 ┆ 43.2 │
│ BB    ┆ null ┆ null │
└───────┴──────┴──────┘