numba - cuDF - groupby UDF 支持日期时间

Question

我有一个包含以下列的 cuDF 数据框：

columns = ["col1", "col2", "dt"]

( dt) 的形式datetime64[ns]。

我想编写一个 UDF 以应用于此数据框中的每个组，并dt为每个组获取最大值。这是我正在尝试的，但似乎numba不支持datetime64[ns]UDF 中的值。

def f1(dt, out):
   l = len(dt)
   maxvalue = dt[0]
   for i in  range(cuda.threadIdx.x, l, cuda.blockDim.x):
      if dt[i] > maxvalue:
         maxvalue = dt[i]
   out[:0] = maxvalue

gdf = df.groupby(["col1", "col2"], method="cudf")
df = gdf.apply_grouped(f1, incols={"dt": "dt"}, outcols=dict(out=numpy.datetime64))

这是我得到的错误：

This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: resolving callee type: Function(<numba.cuda.compiler.DeviceFunctionTemplate object at 0x7effda063510>)
[2] During: typing of call at <string> (10)

我有类似的函数，可以很好地处理整数和浮点数。这是否意味着 numba 不支持日期时间？

标签： numbanumba-procudf

Answer 1

Apply_groups不会给你我认为你追求的东西，即 max 上的 groupby dt。您需要使用aggswith max on dt。cudf 的 groupby 函数将完成其余的工作。要获取您的值datetime64[ms]，请使用astype()，并将其保存回数据框（非常快）。看我的例子：

import cudf
a = cudf.DataFrame({"col1": [1, 1, 1, 2, 2, 2], "col2": [1, 2, 1, 1, 2, 1], "dt": [10000000, 2000000, 3000000, 100000, 2000000, 40000000]}) 
a['dt'] = a['dt'].astype('datetime64[ns]')
print(a)
a['dt'] = a['dt'].astype('datetime64[ms]')
print(a)
gdf = a.groupby(["col1", "col2"]).agg({'dt':'max'})
print(gdf.head())

dt从 1970 年 1 月 1 日起，列值将被格式化为 0.1-40 毫秒（以纳秒计），为您提供打印输出

   col1  col2                         dt
0     1     1 1970-01-01 00:00:00.010000
1     1     2 1970-01-01 00:00:00.002000
2     1     1 1970-01-01 00:00:00.003000
3     2     1 1970-01-01 00:00:00.000100
4     2     2 1970-01-01 00:00:00.002000
5     2     1 1970-01-01 00:00:00.040000

   col1  col2                      dt
0     1     1 1970-01-01 00:00:00.010
1     1     2 1970-01-01 00:00:00.002
2     1     1 1970-01-01 00:00:00.003
3     2     1 1970-01-01 00:00:00.000
4     2     2 1970-01-01 00:00:00.002
5     2     1 1970-01-01 00:00:00.040

                               dt
col1 col2                        
1    1    1970-01-01 00:00:00.010
     2    1970-01-01 00:00:00.002
2    1    1970-01-01 00:00:00.040
     2    1970-01-01 00:00:00.002