如何在 Polars 中执行 pandas reindex

Question

在pandas中，我可以使用多级索引对数据框进行reindex()操作，使每个组的日期范围保持一致。在polars中是否有相同的方法来实现相同的结果？

请参考以下使用pandas的示例：

import pandas as pd
data = pd.DataFrame({
             "date":pd.date_range("2022-01-01", "2022-06-01", freq="MS"),
             "group":["A", "A", "A", "B", "B", "B"],
             "value":[10,20,30,40,50,60]
     }).set_index(["group", "date"])
new_index = pd.MultiIndex.from_product([data.index.levels[0].tolist(), data.index.levels[1].tolist()], names=["group", "date"])
data.reindex(new_index)

将数据从以下形式转换：

                  value
group date             
A     2022-01-01     10
      2022-02-01     20
      2022-03-01     30
B     2022-04-01     40
      2022-05-01     50
      2022-06-01     60

将两个组的日期范围设置为相同：

                  value
group date             
A     2022-01-01   10.0
      2022-02-01   20.0
      2022-03-01   30.0
      2022-04-01    NaN
      2022-05-01    NaN
      2022-06-01    NaN
B     2022-01-01    NaN
      2022-02-01    NaN
      2022-03-01    NaN
      2022-04-01   40.0
      2022-05-01   50.0
      2022-06-01   60.0

- JASMINE LIAW

Answer 1

正如您所读到的，polars不使用索引。

在这种情况下，您的new_index是索引列唯一值的交叉连接，可以轻松地在polars中再现。

pldata=pl.DataFrame({
            "date":pl.date_range(datetime(2022,1,1), datetime(2022,6,1),'1mo'),
            "group":["A", "A", "A", "B", "B", "B"],
            "value":[10,20,30,40,50,60]
    })

pl_index = pldata.select(pl.col('date').unique()) \
        .join(
              pldata.select(pl.col('group').unique()), 
         how='cross'
         )

然后，您可以使用另一个连接代替重建索引命令，这次是在这些列上进行外部连接，并进行排序以恢复原始顺序

pldata.join(plnew_index, on=['date','group'], how='outer').sort(['group','date'])

您可以编写一个辅助函数，用于为任意数量的索引列创建plnew_index

def make_plindex(df, indexcols):
    newdf=df.select(pl.col(indexcols[0]).unique())
    for curcol in indexcols[1:]:
        newdf=newdf.join(df.select(pl.col(curcol).unique()), how='cross')
    return newdf

当然，如果您实际上不关心中间的数据框，可以通过在函数中放置外连接并返回最终结果来扩展该函数。

def make_nullrows(df, indexcols):
    newdf=df.select(pl.col(indexcols[0]).unique())
    for curcol in indexcols[1:]:
        newdf=newdf.join(df.select(pl.col(curcol).unique()), how='cross')
    return df.join(newdf, on=indexcols, how='outer')

然后你可以这样做

make_nullrows(pldata, ['group','date'])

- Dean MacGregor