Polars DataFrame中的复杂计算与Numba优化

在数据处理领域，Polars是一个高效且快速的数据框架，提供了诸如Pandas的类似功能，但性能更优。然而，当涉及到复杂的自定义函数计算时，Polars的处理方式可能不尽如人意，特别是当你需要在DataFrame中进行多列的计算并保留中间结果时。本文将探讨如何通过Numba优化和Polars的map_batches方法来解决这一问题。

问题背景

假设我们有一个包含三列的Polars DataFrame，分别为A、B和D，每列存储double类型的数据。我们需要执行以下计算：

• A列的值乘以一个初始值
• B列的值乘以初始值
• D列的值为A和B列值之和
• 初始值更新为D列的值

这样的计算需要保留A和B列的中间结果，但由于Polars当前版本（如0.20.26）的限制，Numba的gufuncs只能返回一个元素，这导致A和B列的值无法被保留。

解决方案

使用Polars的map_batches和Numba的JIT编译

Polars提供了map_batches方法，可以让我们将自定义的Numba函数应用到DataFrame的每一行上，并且可以返回一个包含多个元素的Struct系列