使用 python 数据表按组创建行号

create row number by group, using python datatable

如果我有这样的 python 数据表:

from datatable import f, dt
data = dt.Frame(grp=["a","a","b","b","b","b","c"], value=[2,3,1,2,5,9,2])

如何按组创建具有行号的新列?。即Rdata.table

相当于什么
data[, id:=1:.N, by=.(grp)]

这可行,但看起来完全荒谬

data['id'] = np.concatenate(
                [np.arange(x)
                    for x in data[:,dt.count(), dt.by(f.grp)]['count'].to_numpy()])

期望的输出:

   | grp    value     id
   | str32  int32  int64
-- + -----  -----  -----
 0 | a          2      0
 1 | a          3      1
 2 | b          1      0
 3 | b          2      1
 4 | b          5      2
 5 | b          9      3
 6 | c          2      0

一种方法是将 to_pandas, groupby (on the pandas DataFrame) and use cumcount:

import datatable as dt

data = dt.Frame(grp=["a", "a", "b", "b", "b", "b", "c"], value=[2, 3, 1, 2, 5, 9, 2])

data["id"] = data.to_pandas().groupby("grp").cumcount()
print(data)

输出

   | grp    value     id
   | str32  int32  int64
-- + -----  -----  -----
 0 | a          2      0
 1 | a          3      1
 2 | b          1      0
 3 | b          2      1
 4 | b          5      2
 5 | b          9      3
 6 | c          2      0
[7 rows x 3 columns]

datatable没有累计计数功能,实际上目前任何聚合都没有累计功能。

一种可能提高速度的方法是使用更快的 numpy 迭代,其中 for 循环在 C 中完成,并且效率更高。代码来自并为此修改:

from datatable import dt, f, by
import numpy as np

In [244]: def create_ranges(indices):
     ...:     cum_length = indices.cumsum()
     ...:     ids = np.ones(cum_length[-1], dtype=int)
     ...:     ids[0] = 0
     ...:     ids[cum_length[:-1]] = -1 * indices[:-1] + 1
     ...:     return ids.cumsum()


counts =  data[:, dt.count(), by('grp', add_columns=False)].to_numpy().ravel()
data[:, f[:].extend({"counts" : create_ranges(counts)})]

   | grp    value  counts
   | str32  int32   int64
-- + -----  -----  ------
 0 | a          2       0
 1 | a          3       1
 2 | b          1       0
 3 | b          2       1
 4 | b          5       2
 5 | b          9       3
 6 | c          2       0
[7 rows x 3 columns]

create_ranges 函数非常棒(建立在 cumsum 上的逻辑很好)并且随着数组大小的增加真正发挥作用。

当然这有它的缺点;你正在走出数据表进入 numpy 领域,然后回到数据表;另一方面是我寄希望于这样一个事实,即这些组是按词汇排序的;如果数据未排序(并且必须在分组列上排序),这将不起作用。

初步测试速度明显提升;同样,它的范围有限,如果将其烘焙到数据表库中,它会 easier/better 很多。

如果你擅长C++,可以考虑贡献这个函数到库中;我和其他很多人都会感谢你的努力。

您可以查看 pypolars,看看它是否对您的用例有帮助。从 h2o 基准测试来看,它看起来是一个非常快的工具。