nodejs-polars如何定义封装Polars表达式的可链式DataFrame函数

问题背景 #

实现动机 #

需要抽象支持参数配置（通过自定义函数传参）、可链式调用的参数化Polars表达式，在DataFrame上提供高层级、可复用、支持链式调用的自定义数据分析函数。

失败复现 #

直接修改DataFrame.prototype挂载自定义方法会触发运行时错误，示例代码：

import pl from "nodejs-polars"
const { DataFrame, col } = pl

// 尝试在原型上挂载自定义增量方法
DataFrame.prototype.inc = function inc(column, x = 1, alias = `${column}Inc`) {
  return this.withColumn(col(column).add(x).alias(alias))
}
const df = new DataFrame({ a: [1, 2, 3] })
// 原生Polars写法可正常运行，但复用性差，重复代码多
console.log(df.withColumn(col("a").add(1).alias("aInc")))
// 链式调用时报错：TypeError: df.inc is not a function
console.log(df.inc("a").inc("a", 2, "aInc2"))

报错核心原因：nodejs-polars的DataFrame是Rust编译生成的WASM绑定类，直接修改JS侧暴露的原型对象不会生效，链式调用返回的所有新DataFrame实例都由原生侧构造，不会继承JS侧手动给原型添加的方法。

推荐实现方式 #

方式1：使用原生pipe方法串联自定义函数（官方推荐，零侵入） #

Polars的DataFrame原生提供pipe方法，专门用于接入自定义的可复用逻辑，不需要修改任何原生对象，兼容性最好：

import pl from "nodejs-polars"
const { DataFrame, col } = pl

// 自定义函数格式：第一个参数固定为DataFrame实例，后续接收业务参数
const inc = (df, column, x = 1, alias = `${column}Inc`) => {
  return df.withColumn(col(column).add(x).alias(alias))
}

// 链式调用通过pipe串联自定义函数和原生方法
const df = new DataFrame({ a: [1, 2, 3] })
const result = df
  .pipe(inc, "a")
  .pipe(inc, "a", 2, "aInc2")
  .filter(col("a").gt(1)) // 可以任意穿插原生Polars方法

console.log(result)

这种方式下所有自定义函数都是纯函数，测试、跨项目复用都很方便，和Polars生态完全兼容，是首选方案。

方式2：Proxy包装类（接近原生调用体验） #

如果不想每次链式调用都写pipe，可以用Proxy做一层轻量包装，自动转发原生方法，同时挂载自定义方法，调用体验和原生DataFrame几乎无差异：

import pl from "nodejs-polars"
const { DataFrame, col } = pl

const wrapDf = (rawDf) => {
  return new Proxy(rawDf, {
    get(target, prop) {
      // 注册自定义方法
      if (prop === "inc") {
        return (column, x = 1, alias = `${column}Inc`) => {
          const newDf = target.withColumn(col(column).add(x).alias(alias))
          return wrapDf(newDf)
        }
      }
      // 自动转发原生方法，如果返回值是DataFrame则自动包装
      const value = target[prop]
      if (typeof value === "function") {
        return (...args) => {
          const res = value.apply(target, args)
          return res instanceof DataFrame ? wrapDf(res) : res
        }
      }
      return value
    }
  })
}

// 使用示例
const df = wrapDf(new DataFrame({ a: [1, 2, 3] }))
df
  .inc("a")
  .inc("a", 2, "aInc2")
  .head(2) // 原生方法可直接调用
  .writeCSV("./res.csv")

不要尝试直接修改DataFrame.prototype，除了前面提到的实例继承问题，后续Polars版本升级时原型结构变动也会导致自定义方法失效，维护成本极高。

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