我可以用 Polars 优化这个 CPU 密集型 Pandas 代码吗？

Question

我有这个熊猫代码：

def last_non_null(s):
    return s.dropna().iloc[-1] if not s.dropna().empty else np.nan

def merge_rows_of_final_df(df_final):
    # Group by columns A, B, and C
    cols = ['A', 'B', 'C']

    # Apply function to each column not in the subset
    # NOTE: the other columns have np.float32 as dtype
    agg_dict = {col: last_non_null for col in df_final.columns.difference(cols)}

    # Group and aggregate
    return df_final.groupby(cols).agg(agg_dict).reset_index()

df_merged = merge_rows_of_final_df(df_final)

我可以用极坐标来优化它吗？如果是这样，有人可以为我提供具有相同输出的极地片段吗？这个想法是，final_df 是连接在子集列 A、B、C 上具有重复项的数据帧列表的结果。我希望通过获取每列的最后一个/第一个非空值来合并这些组的行。

Answer 1

要使用 Polars 优化给定的 Pandas 代码，您可以使用以下 Polars 代码片段：

import polars as pl

def merge_rows_of_final_df(df_final):
    subset_cols = ['A', 'B', 'C']
    other_cols = df_final.columns.difference(subset_cols)
    
    df_final = df_final.lazy() \
        .groupby(subset_cols) \
        .agg(pl.last(other_cols).ignore_nulls()) \
        .collect()
    
    return df_final

说明：

1. 我们定义了

   merge_rows_of_final_df

   函数，它将

   df_final

   DataFrame 作为输入。
2. 我们指定子集列

   ['A', 'B', 'C']

   并使用

   df_final.columns.difference(subset_cols)

   确定其他列。
3. 我们使用

   df_final

   对

   .lazy()

   开始惰性计算。
4. 我们使用

   .groupby(subset_cols)

   对子集列执行 groupby 操作。
5. 我们使用

   .agg(pl.last(other_cols).ignore_nulls())

   聚合其他列。这采用

   other_cols

   中每列的最后一个非空值。
6. 我们使用

   .collect()

   将结果收集到新的 DataFrame 中。
7. 最后，我们返回合并后的DataFrame。

Polars 代码实现了与 pandas 代码相同的结果，但它应该更高效并且针对 CPU 密集型任务进行了优化。 Polars 利用惰性评估和优化算法来高效地执行分组和聚合操作。

注意：在运行 Polars 代码之前，请确保已安装 Polars 库 (

pip install polars

)。