pandas的apply操作

pandas的apply操作类似于Scala的udf一样方便，假设存在如下dataframe：

  id_part                  pred               pred_class v_id

0       d  [0.722817, 0.650064]                  cat,dog   d1

1       5  [0.119208, 0.215449]  other_label,other_label   d2

需要把 v_id=d1 中，pred 与 pred_class 一一对应，需要将 pred 大于0.5的pred_class取出来作为新的一列，如果小于0.5则不取出来：

import pandas as pd

# 提取类别

def get_pred_class(pred_class, pred):

    pred_class_list = pred_class.split(",")

    result_class_list = []

    for i in range(0, len(pred)):

        if float(pred[i]) >= 0.5:

            result_class_list.append(pred_class_list[pred.index(pred[i])])

    return result_class_list

# 新建一个dataframe

data = pd.DataFrame({

    'v_id': ["d1", 'd2'],

    'pred_class': ["cat,dog", 'other_label,other_label'],

    'pred': [[0.722817,0.650064], [0.119208,0.215449]],

    'id_part': ["d", '5'],

})

df = data.copy()

df["pos_labels"] = data.apply(lambda row: get_pred_class(row['pred_class'], row['pred']), axis=1)

print(df)

得到结果为：

  id_part                  pred               pred_class v_id  pos_labels

0       d  [0.722817, 0.650064]                  cat,dog   d1  [cat, dog]

1       5  [0.119208, 0.215449]  other_label,other_label   d2          []

PS：如果没有df = data.copy()将会出现错误：

A value is trying to be set on a copy of a slice from a DataFrame.

Try using .loc[row_indexer,col_indexer] = value instead