将 dask 日期月份与整数进行比较时出错

Error comparing dask date month with an integer

下面代码中的 dask map_partitions 函数有一个 dask 日期字段,其中它的月份与整数进行比较。此比较失败并出现以下错误:

ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().

这是什么错误以及如何解决?

import pandas as pd
import dask
import dask.dataframe as dd
import datetime

pdf = pd.DataFrame({
    'id2': [1, 1, 1, 2, 2],
    'balance': [150, 140, 130, 280, 260],
    'date2' : [datetime.datetime(2021,3,1), datetime.datetime(2021,4,1), 
               datetime.datetime(2021,5,1), datetime.datetime(2021,1,1), 
               datetime.datetime(2021,2,1)]
})

ddf = dd.from_pandas(pdf, npartitions=1) 

def func2(obj):
    m = obj.date2.dt.month
    if m > 10:
        return 1
    else:
        return 2

ddf2 = ddf.map_partitions(func2, meta=int)
ddf2.compute()   # <-- fails here

通过使用 .map_partition,每个 dask 数据帧分区(这是一个 pandas 数据帧)都被传递给函数 func2。结果obj.date2.dt.month指的是一个Series,不是单个值,所以通过运行与整数的比较,不清楚Python如何判断比较的有效性.

作为一个选项,下面是创建新列的片段,条件是 dt.month 结果:

import pandas as pd
import dask
import dask.dataframe as dd
import datetime

pdf = pd.DataFrame({
    'id2': [1, 1, 1, 2, 2],
    'balance': [150, 140, 130, 280, 260],
    'date2' : [datetime.datetime(2021,3,1), datetime.datetime(2021,4,1), 
               datetime.datetime(2021,5,1), datetime.datetime(2021,1,1), 
               datetime.datetime(2021,2,1)]
})

ddf = dd.from_pandas(pdf, npartitions=1) 

def func2(obj):
    m = obj.date2.dt.month
    obj.loc[m>10, 'new_int']=1
    obj.loc[m<=10, 'new_int']=2
    return obj

ddf2 = ddf.map_partitions(func2)
ddf2.compute()