DataFrame的assign和query方法

楔子

pandas支持的功能已经非常强大了,但是本着程序猿的"懒惰"精神,还是想少写一点代码的。pandas中的某些方法不仅可以实现我们需要的功能,还可以减少一定的代码量,我们来看一下。

df.assign

这个方法估计有人没怎么用过,因为我们不用它也完全可以实现,那么它是干什么的呢?我们来看一下:

import pandas as pd

df = pd.DataFrame(
    {"name": ["mashiro", "satori", "koishi"],
     "date": pd.date_range("2018-1-1", "2018-1-3", freq="D")}
)

print(df)
"""
      name       date
0  mashiro 2018-01-01
1   satori 2018-01-02
2   koishi 2018-01-03
"""
# 我们需要增加三个字段,分别存储date的年、月、日
# 之前的话,我们可能会这么做
df["year"], df["month"], df["day"] = (df["date"].dt.year,
                                      df["date"].dt.month,
                                      df["date"].dt.day)
print(df)
"""
      name       date  year  month  day
0  mashiro 2018-01-01  2018      1    1
1   satori 2018-01-02  2018      1    2
2   koishi 2018-01-03  2018      1    3
"""

先说一下,上面的代码没有任何问题,非常标准的写法。但如果使用assign的话会更加方便

import pandas as pd

df = pd.DataFrame(
    {"name": ["mashiro", "satori", "koishi"],
     "date": pd.date_range("2018-1-1", "2018-1-3", freq="D")}
)

print(df)
"""
      name       date
0  mashiro 2018-01-01
1   satori 2018-01-02
2   koishi 2018-01-03
"""

df = df.assign(
    year=df["date"].dt.year,
    month=df["date"].dt.month,
    day=df["date"].dt.day
)
print(df)
"""
      name       date  year  month  day
0  mashiro 2018-01-01  2018      1    1
1   satori 2018-01-02  2018      1    2
2   koishi 2018-01-03  2018      1    3
"""

df.assign里面通过关键字参数传递,参数名会作为列名,参数值会作为列的值。

df.query

df.query,从名字上也能看出来这是用于筛选的。

import pandas as pd

df = pd.DataFrame(
    {"name": ["mashiro", "satori", "koishi", "kurisu", "nagsia"],
     "age": [17, 17, 16, 18, 21],
     "height": [161, 155, 154, 172, 158],
     "where": ["樱花庄", "地灵殿", "地灵殿", "石头门", "clannad"]
     }
)
print(df)
"""
      name  age  height    where
0  mashiro   17     161      樱花庄
1   satori   17     155      地灵殿
2   koishi   16     154      地灵殿
3   kurisu   18     172      石头门
4   nagsia   21     158  clannad
"""

# 筛选age > 17的
print(df.query("age > 17"))  # 等价于df[df["age"] > 17]
"""
     name  age  height    where
3  kurisu   18     172      石头门
4  nagsia   21     158  clannad
"""

# 筛选where == '地灵殿'的
print(df.query("where == '地灵殿'"))  # 等价于df[df["where"] > "地灵殿"]
"""
     name  age  height where
1  satori   17     155   地灵殿
2  koishi   16     154   地灵殿
"""

# 筛选where == '地灵殿' 并且 age == 17的
# 等价于df[(df["where"] > "地灵殿") & (df["age"] == 17)]
print(df.query("where == '地灵殿' and age == 17"))
"""
     name  age  height where
1  satori   17     155   地灵殿
"""

# 筛选出身高大于年龄的,当然这个肯定都满足,只是演示一下query支持的功能
print(df.query("height > age"))  # 等价于df[df["height"] >df["age"]]
"""
      name  age  height    where
0  mashiro   17     161      樱花庄
1   satori   17     155      地灵殿
2   koishi   16     154      地灵殿
3   kurisu   18     172      石头门
4   nagsia   21     158  clannad
"""

# 另外如果字段名含有空格怎么办?我们将name变成na me
df = df.rename(columns={"name": "na me"})
# 这个时候应该将na me使用``包起来,告诉pandas,``里面的内容是一个整体。否则的话,是不会正确解析的
print(df.query("`na me` == 'satori'"))
"""
    na me  age  height where
1  satori   17     155   地灵殿
"""

# 查找age位于[17, 21]当中、并且where位于["樱花庄", "clannad]大当中的记录
print(df.query("age in [17, 21] and where in ['樱花庄', 'clannad']"))
"""
     na me  age  height    where
0  mashiro   17     161      樱花庄
4   nagsia   21     158  clannad
"""

# 问题来了,如果我们外面有一个变量,我们怎么在query里面使用呢?
place = "地灵殿"
try:
    print(df.query("where == place"))
except Exception as e:
    # 尽管我们外面定义了place,但是不能再query里面直接使用
    print(e)  # name 'place' is not defined
# 如果使用的话,那么应该使用@作为前缀,来标识这是一个外面存在的变量
print(df.query("where == @place"))
"""
    na me  age  height where
1  satori   17     155   地灵殿
2  koishi   16     154   地灵殿
"""

l = [16, 18]
print(df.query("age in @l"))
"""
    na me  age  height where
2  koishi   16     154   地灵殿
3  kurisu   18     172   石头门
"""

# 也可以使用索引
print(df.query("age == @l[0]"))
"""
    na me  age  height where
2  koishi   16     154   地灵殿
"""

# 我们试一下函数
age = 17
def foo(x): return x
print(df.query("age == @foo(x=@age)"))
"""
     na me  age  height where
0  mashiro   17     161   樱花庄
1   satori   17     155   地灵殿
"""
# 如果上面是可以的,那么age == @foo(@age)也是可以的,无参函数更是可以的

# 再试一下字典
d = {"age": 16}
try:
    print(df.query("age == @d['age']"))
except Exception as e:
    # 如果是通过中括号获取的话,那么[]里面不能出现''
    print(e)  # data type must provide an itemsize
# 所以我们可以通过字典的get方法获取
print(df.query("age == @d.get('age')"))
"""
    na me  age  height where
2  koishi   16     154   地灵殿
"""

# age1不存在,也自动使用了默认值
print(df.query("age == @d.get('age1', 16)"))
"""
    na me  age  height where
2  koishi   16     154   地灵殿
"""

# 回到通过中括号获取元素上面来,我们使用列表是可以通过[]获取指定元素的
# 因为列表里面指定的是数值,那么是不是意味着如果字典里面的key也是数值的话、我们就可以通过[]来获取呢
d = {16: 21}
# 显然答案是正确的
print(df.query("age == @d[16]"))
"""
    na me  age  height    where
4  nagsia   21     158  clannad
"""
# 所以对于字典来讲,如果通过[]获取的话,那么里面不能出现'',所以也就无法获取字符串
# 但是对于整型是可以的,因为不需要''括起来。
# 不过可能有人会问了,如果在外面定义一个变量age=16,然后再在query里面通过"age = @d[@age]"获取可不可以呢
# 因为此时[]里面没有出现''。答案也是不可以的,具体原因可以自己试着写一下。

# 总体来说query的功能还是能够支持我们大部分逻辑的
原文地址:https://www.cnblogs.com/traditional/p/11111327.html