上下文

表达式几乎可以在任何地方使用，但是表达式需要一个上下文，这些上下文包括：

• 选择: df.select([..])
• 分组集合: df.groupby(..).agg([..])
• 横向堆叠(hstack) 或者增加列: df.with_columns([..])

语法糖

需要上下文的主要原因是：即使在即时执行中，你也在使用 Polars 的延迟执行API。 比如如下代码实例：

df.groupby("foo").agg([pl.col("bar").sum()])

去掉语法糖后：

(df.lazy().groupby("foo").agg([pl.col("bar").sum()])).collect()

这种设计可以让 Polars 把表达式推送给查询引擎，进行一些优化和缓存操作。

select 上下文

在 select 上下文中，选择操作是按照列进行的。在选择向下文的表达式必须要返回 Series 并且这些 Series 需要有相同的长度或者长度为1。

一个长度为 1 的 Series 会将 DataFrame 的一列赋予完全一样的值(这个值来自Series)。 注意，select 可能会返回一个新的列，这个列可能是一些聚合的结果、一些表达式的组合或者常量。

选择上下文

out = df.select(
    [
        pl.sum("nrs"),
        pl.col("names").sort(),
        pl.col("names").first().alias("first name"),
        (pl.mean("nrs") * 10).alias("10xnrs"),
    ]
)
print(out)

shape: (5, 4)
┌─────┬───────┬────────────┬────────┐
│ nrs ┆ names ┆ first name ┆ 10xnrs │
│ --- ┆ ---   ┆ ---        ┆ ---    │
│ i64 ┆ str   ┆ str        ┆ f64    │
╞═════╪═══════╪════════════╪════════╡
│ 11  ┆ null  ┆ foo        ┆ 27.5   │
│ 11  ┆ egg   ┆ foo        ┆ 27.5   │
│ 11  ┆ foo   ┆ foo        ┆ 27.5   │
│ 11  ┆ ham   ┆ foo        ┆ 27.5   │
│ 11  ┆ spam  ┆ foo        ┆ 27.5   │
└─────┴───────┴────────────┴────────┘

添加列

采用 with_columns 给 DataFrame 增加列同样也是选择上下文。

df = df.with_columns(
    [
        pl.sum("nrs").alias("nrs_sum"),
        pl.col("random").count().alias("count"),
    ]
)
print(out)

shape: (5, 6)
┌──────┬───────┬──────────┬────────┬─────────┬───────┐
│ nrs  ┆ names ┆ random   ┆ groups ┆ nrs_sum ┆ count │
│ ---  ┆ ---   ┆ ---      ┆ ---    ┆ ---     ┆ ---   │
│ i64  ┆ str   ┆ f64      ┆ str    ┆ i64     ┆ u32   │
╞══════╪═══════╪══════════╪════════╪═════════╪═══════╡
│ 1    ┆ foo   ┆ 0.154163 ┆ A      ┆ 11      ┆ 5     │
│ 2    ┆ ham   ┆ 0.74005  ┆ A      ┆ 11      ┆ 5     │
│ 3    ┆ spam  ┆ 0.263315 ┆ B      ┆ 11      ┆ 5     │
│ null ┆ egg   ┆ 0.533739 ┆ C      ┆ 11      ┆ 5     │
│ 5    ┆ null  ┆ 0.014575 ┆ B      ┆ 11      ┆ 5     │
└──────┴───────┴──────────┴────────┴─────────┴───────┘

Groupby 上下文

在 groupby 上下文中的表达式主要作用域分组上，因此他们会返回任意长度（每个组可能有不同数量的成员）。

out = df.groupby("groups").agg(
    [
        pl.sum("nrs"),  # 通过groups列对nrs求和
        pl.col("random").count().alias("count"),  # 记录组数
        # 如果name != null记录random列的和
        pl.col("random").filter(pl.col("names").is_not_null()).sum().suffix("_sum"),
        pl.col("names").reverse().alias(("reversed names")),
    ]
)
print(out)

shape: (3, 5)
┌────────┬─────┬───────┬────────────┬────────────────┐
│ groups ┆ nrs ┆ count ┆ random_sum ┆ reversed names │
│ ---    ┆ --- ┆ ---   ┆ ---        ┆ ---            │
│ str    ┆ i64 ┆ u32   ┆ f64        ┆ list[str]      │
╞════════╪═════╪═══════╪════════════╪════════════════╡
│ C      ┆ 0   ┆ 1     ┆ 0.533739   ┆ ["egg"]        │
│ A      ┆ 3   ┆ 2     ┆ 0.894213   ┆ ["ham", "foo"] │
│ B      ┆ 8   ┆ 2     ┆ 0.263315   ┆ [null, "spam"] │
└────────┴─────┴───────┴────────────┴────────────────┘

除了标准的 groupby，还有 groupby_dynamic 和 groupby_rolling 也属于 Groupby 上下文。