如何在含集合类型列的Pandas DataFrame中使用分支判断

问题背景 #

有一个名为box_status的Pandas DataFrame，status列存储Set类型值，数据示例如下：

+-----------------------+-------------------------+
|             zone_name |                   status|
+-----------------------+-------------------------+
|          D7_BATCH_BOX |            {NOT_STARTED}|
|    BUS_BATCH_AMER_BOX | {NOT_STARTED, COMPLETED}|
|    BUS_BATCH_AUST_BOX | {NOT_STARTED, COMPLETED}|
|    DOM_BATCH_AMER_BOX |  {NOT_LOADED, COMPLETED}|
|    DOM_BATCH_AUST_BOX |  {NOT_LOADED, COMPLETED}|
|    DOM_BATCH_CAND_BOX |  {NOT_LOADED, COMPLETED}|
|         SIT_BATCH_BOX |            {NOT_STARTED}|
|        SHORT_FEES_BOX |              {COMPLETED}|
+-----------------------+-------------------------+

需要根据status列的集合内容生成新列zone_status，规则：

• 仅含COMPLETED → COMPLETED
• 同时含COMPLETED和NOT_STARTED → RUNNING
• 仅含NOT_STARTED → YET TO START
• 含NOT_LOADED相关组合 → ERROR

原代码尝试用issubset()结合np.where实现，触发报错：

Traceback (most recent call last):
  File "status_mailer_new.py", line 211, in <module>
    box_status["zone_status"] = np.where(comp_set.issubset(box_status["status"]), "COMPLETED",
TypeError: unhashable type: 'set'

原代码：

comp_set = {"COMPLETED"}
runn_set = {"COMPLETED","NOT_STARTED"}
yet_to_start_set = {"NOT_STARTED"}
fail_set1 = {"NOT_STARTED","NOT_LOADED"}
fail_set2 = {"NOT_LOADED","COMPLETED"}
fail_set3 = {"NOT_LOADED"}
box_status["zone_status"] = np.where(comp_set.issubset(box_status["status"]), "COMPLETED",
                           np.where(runn_set.issubset(box_status["status"]), "RUNNING",
                           np.where(yet_to_start_set.issubset(box_status["status"]), "YET TO START",
                           np.where(fail_set1.issubset(box_status["status"]),"ERROR",
                           np.where(fail_set2.issubset(box_status["status"]),"ERROR",
                           np.where(fail_set3.issubset(box_status["status"]),"ERROR","ERROR"))))))

报错原因 #

issubset()是单个集合的方法，无法直接作用于Pandas Series对象（Series包含多个集合元素）。直接调用会导致Pandas尝试将整个Series传入issubset()，而集合属于不可哈希类型，触发类型错误。

解决方法 #

方法1：自定义函数+apply遍历 #

定义函数处理单个集合，再用apply遍历整个status列：

def get_zone_status(status_set):
    comp_set = {"COMPLETED"}
    runn_set = {"COMPLETED", "NOT_STARTED"}
    yet_to_start_set = {"NOT_STARTED"}
    fail_sets = [{"NOT_STARTED","NOT_LOADED"}, {"NOT_LOADED","COMPLETED"}, {"NOT_LOADED"}]
    
    if status_set == comp_set:
        return "COMPLETED"
    elif runn_set.issubset(status_set):
        return "RUNNING"
    elif status_set == yet_to_start_set:
        return "YET TO START"
    elif any(fail_set.issubset(status_set) for fail_set in fail_sets):
        return "ERROR"
    else:
        return "ERROR"

box_status["zone_status"] = box_status["status"].apply(get_zone_status)

方法2：lambda结合np.where #

对每个np.where的判断条件，用lambda逐个处理集合：

comp_set = {"COMPLETED"}
runn_set = {"COMPLETED","NOT_STARTED"}
yet_to_start_set = {"NOT_STARTED"}
fail_set1 = {"NOT_STARTED","NOT_LOADED"}
fail_set2 = {"NOT_LOADED","COMPLETED"}
fail_set3 = {"NOT_LOADED"}

box_status["zone_status"] = np.where(
    box_status["status"].apply(lambda x: x == comp_set), 
    "COMPLETED",
    np.where(
        box_status["status"].apply(lambda x: runn_set.issubset(x)), 
        "RUNNING",
        np.where(
            box_status["status"].apply(lambda x: x == yet_to_start_set), 
            "YET TO START",
            np.where(
                box_status["status"].apply(lambda x: fail_set1.issubset(x) or fail_set2.issubset(x) or fail_set3.issubset(x)), 
                "ERROR",
                "ERROR"
            )
        )
    )
)

方法3：case_when简化逻辑（需安装pandas_flavor） #

通过扩展工具让分支逻辑更直观：

from pandas_flavor import register_dataframe_method

@register_dataframe_method
def case_when(df, cases):
    result = df.iloc[:, 0].copy()
    for condition, value in cases:
        result = np.where(condition, value, result)
    return result

box_status["zone_status"] = box_status.case_when([
    (box_status["status"].apply(lambda x: x == {"COMPLETED"}), "COMPLETED"),
    (box_status["status"].apply(lambda x: {"COMPLETED", "NOT_STARTED"}.issubset(x)), "RUNNING"),
    (box_status["status"].apply(lambda x: x == {"NOT_STARTED"}), "YET TO START"),
    (box_status["status"].apply(lambda x: any(s.issubset(x) for s in [{"NOT_STARTED","NOT_LOADED"}, {"NOT_LOADED","COMPLETED"}, {"NOT_LOADED"}])), "ERROR")
])

处理后结果 #

内容的提问来源于stack exchange，提问作者aiman