nodejs-polars中如何简便实现DataFrame各列与对应列均值的运算

环境 #

macos:             monterey
node:              v18.1.0
nodejs-polars:     0.5.3

实现目标 #

对polars DataFrame中的每一列，减去该列自身的均值。

Pandas实现方案 #

在pandas中该操作实现非常简洁，可借助DataFrame.sub(other, axis='columns', level=None, fill_value=None)方法实现，该方法的other参数支持传入标量、序列、Series或DataFrame类型：

df.sub(df.mean())
df - df.mean()

nodejs-polars实现尝试 #

在nodejs-polars中，查看源码可见，sub方法定义为sub: (other) => wrap("sub", prepareOtherArg(other).inner())，其other参数最初看起来仅支持Series类型，实际测试中踩了不少坑。

1. 准备测试数据 #

测试用DataFrame内容如下：

console.log(df)
┌─────────┬─────────┬─────────┬─────────┐
│   A     ┆   B     ┆   C     ┆   D     │
│ ---     ┆ ---     ┆ ---     ┆ ---     │
│ i64     ┆ i64     ┆ i64     ┆ i64     │
╞═════════╪═════════╪═════════╪═════════╡
│ 13520   ┆ -16     ┆ 384     ┆ 208     │
│ 13472   ┆ -16     ┆ 384     ┆ 176     │
│ 13456   ┆ -16     ┆ 368     ┆ 160     │
│ 13472   ┆ -16     ┆ 368     ┆ 160     │
│ 13472   ┆ -16     ┆ 352     ┆ 176     │
└─────────┴─────────┴─────────┴─────────┘

各列均值计算结果：

console.log(df.mean())
┌─────────┬─────────┬─────────┬─────────┐
│   A     ┆   B     ┆   C     ┆   D     │
│ ---     ┆ ---     ┆ ---     ┆ ---     │
│ f64     ┆ f64     ┆ f64     ┆ f64     │
╞═════════╪═════════╪═════════╪═════════╡
│ 13478.4 ┆ -16.0   ┆ 371.2   ┆ 176.0   │
└─────────┴─────────┴─────────┴─────────┘

2. 第一次尝试 #

直接传入均值结果DataFrame调用sub方法，直接触发类型错误：

df.sub(df.mean())

Error: Failed to determine supertype of Int64 and Struct([Field { name: "A", dtype: Int32 }, Field { name: "B", dtype: Int32 }, Field { name: "C", dtype: Int32 }, Field { name: "D", dtype: Int32 }])

3. 第二次尝试 #

把均值行提取出来转成Series传入，直接触发内存问题导致程序崩溃：

df.sub(pl.Series(df.mean().row(0)))

Program crashes due to memory problems.

4. 第三次尝试 #

对照官方测试用例能发现，0.5.3版本的sub方法仅支持两类入参：传入标量时，全表所有元素统一减去该标量值；传入Series时，按行位置对齐做逐行减法，和pandas默认按列对齐广播的逻辑完全不匹配。

test("sub", () => {
  const actual = pl.DataFrame({
    "foo": [1, 2, 3],
    "bar": [4, 5, 6]
  }).sub(1);
  const expected = pl.DataFrame({
    "foo": [0, 1, 2],
    "bar": [3, 4, 5]
  });
  expect(actual).toFrameEqual(expected);
});
test("sub:series", () => {
  const actual = pl.DataFrame({
    "foo": [1, 2, 3],
    "bar": [4, 5, 6]
  }).sub(pl.Series([1, 2, 3]));
  const expected = pl.DataFrame({
    "foo": [0, 0, 0],
    "bar": [3, 3, 3]
  });
  expect(actual).toFrameEqual(expected);
});

最初只能用比较繁琐的方案实现：逐列单独计算减均值，再把所有列横向拼接成结果表。

pl.concat(df.columns.map((col) => df.select(col).sub(df.select(col).mean(0).toSeries())), {how:'horizontal'})

5. 更简洁的实现方案 #

后续摸索出一个更简洁的写法，不需要逐列遍历拼接：

df.select(pl.col('*').sub(pl.col('*').mean()))

这个写法的运行逻辑符合polars的表达式设计规则：pl.col('*')是通配选中所有列，表达式链上的运算会对每一列独立执行，所以.mean()会单独计算每一列自身的均值，再和原列做逐元素减法，天然实现了按列广播减均值的效果，执行效率比逐列拼接高很多。

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