首页 > 解决方案 > Pandas 相当于 SQL case when 语句创建新列

问题描述

我有这个df:

df = pd.DataFrame.from_dict(
    {
        'Name': ['Jane', 'Melissa', 'John', 'Matt', 'Abernethy', 'Annie', 'Brook', 'Brian', 'Carrie'],
        'Tag': ['tag1,tag2', 'tag1', 'tag4,tag3,tag7', 'tag2,tag9', 'tag1,tag3', 'tag3,tag4', 'tag9,tag2', 'tag3,tag2', 'tag1,tag5'],
    }
)

看起来像这样:

姓名 标签
标签1,标签2
梅丽莎 标签9,标签错误1
约翰 标签4,标签3,标签7
马特 标签2,标签9
阿伯内西 标签1,标签3
安妮 标签3,标签4,标签5
标签9,标签2
布赖恩 标签3,标签2
嘉莉 标签1,标签5

我的目标是创建第三列“Tag_after”。简单的 SQL 案例语句将是:

UPDATE table SET Tag_after =
CASE
    WHEN Tag LIKE '%tag1%'           THEN 'tag1'
    WHEN Tag LIKE '%tag2%'           THEN 'tag2'
    WHEN Tag LIKE '%tag3%'           THEN 'tag3'
    WHEN Tag LIKE '%tag4%'           THEN 'tag4'
    WHEN Tag LIKE '%tag5%'           THEN 'tag5'
    WHEN Tag LIKE '%tag_wrong1%'     THEN 'tag_right1'
    ELSE Tag
END

tag1 的优先级高于 tag2,以此类推

tag_wrong1 将更改为 tag_right1

所需的输出是这样的:

姓名 标签 标记后
标签1,标签2 标签1
梅丽莎 标签9,标签错误1 tag_right1
约翰 标签4,标签3,标签7 标签3
马特 标签2,标签9 标签2
阿伯内西 标签1,标签3 标签1
安妮 标签3,标签4,标签5 标签3
标签9,标签2 标签2
布赖恩 标签3,标签7 标签3
嘉莉 标签1,标签5 标签1

我的(错误的)try1:

import pandas as pd

tag_1 = ['tag1', 'tag2', 'tag3', 'tag4', 'tag5', 'tag6', 'tag7', 'tag8', 'tag_wrong1', 'tag9']
tag_2 = ['tag1', 'tag2', 'tag3', 'tag4', 'tag5', 'tag6', 'tag7', 'tag8', 'tag_right1', 'tag9']


df['Tag_after'] = ''


def set_visitor_tag(df, tag_before, tag_after, col_tag, add_col_tag):
    i = 0
    while i < len(tag_before):
        df.loc[~df[col_tag].isnull() & df[col_tag].str.contains(tag_before[i]), [add_col_tag]] = tag_after[i]
        i = i + 1

set_visitor_tag(df, tag_1, tag_2, 'Tag', 'Tag_after')

这个处理的结果和我设置的权限的优先级不一样。

我认为该函数对每一行数据进行多次匹配和赋值操作,我想要的是处理一次之后,就不会再处理它了。

我的(错误的) try2 :

def set_visitor_tag(df, tag_before, tag_after, col_tag, add_col_tag):
    i = 0
    while i < len(tag_before):
        if tag_before[i] in df[col_tag]:
            df.loc[df[col_tag].str.contains(tag_before[i]), [add_col_tag]] = tag_after[i]
        else:
            continue
        i = i + 1

很多谢谢。

标签: pythonsqlpandas

解决方案


一种选择是使用 pyjanitor 中的case_when函数,这类似于 SQL 的 case when:

# pip install git+https://github.com/pyjanitor-devs/pyjanitor.git
import pandas as pd
import janitor
df.case_when(
    df.Tag.str.contains('tag1'), 'tag1', # condition, result
     df.Tag.str.contains('tag2'), 'tag2',
     df.Tag.str.contains('tag3'), 'tag3',
     df.Tag.str.contains('tag4'), 'tag4',
     df.Tag.str.contains('tag5'), 'tag5',
     df.Tag.str.contains('tag_wrong1'), 'tag_right1',
     df.Tag, # default if none of the conditions evaluate to True
    column_name = 'Tag_after')
Out[11]: 
        Name              Tag   Tag_after
0       Jane        tag1,tag2        tag1
1    Melissa  tag9,tag_wrong1  tag_right1
2       John   tag4,tag3,tag7        tag3
3       Matt        tag2,tag9        tag2
4  Abernethy        tag1,tag3        tag1
5      Annie   tag3,tag4,tag5        tag3
6      Brook        tag9,tag2        tag2
7      Brian        tag3,tag2        tag2
8     Carrie        tag1,tag5        tag1

另一种选择是使用 numpy 的选择功能:

condlist = [df.Tag.str.contains('tag1'), df.Tag.str.contains('tag2'),
            df.Tag.str.contains('tag3'), df.Tag.str.contains('tag4'),
            df.Tag.str.contains('tag5'), df.Tag.str.contains('tag_wrong1') ]
choicelist = ['tag1', 'tag2', 'tag3', 'tag4', 'tag5', 'tag_right1']
df.assign(Tag_after = np.select(condlist, choicelist, df.Tag))

        Name              Tag   Tag_after
0       Jane        tag1,tag2        tag1
1    Melissa  tag9,tag_wrong1  tag_right1
2       John   tag4,tag3,tag7        tag3
3       Matt        tag2,tag9        tag2
4  Abernethy        tag1,tag3        tag1
5      Annie   tag3,tag4,tag5        tag3
6      Brook        tag9,tag2        tag2
7      Brian        tag3,tag2        tag2
8     Carrie        tag1,tag5        tag1


推荐阅读