R 语言的管道
这回来介绍一下如何利用管道(pipe)风格将 Pandas 相关的代码写得更易读,不过首先让我们看看隔壁 R 语言中管道是怎么用的。假设输入是 x,经过连续四个函数的处理后得到输出 y,代码可以按顺序写:
x1 <- func1(x, arg1)
x2 <- func2(x1, arg2)
x3 <- func3(x2, arg3)
y <- func4(x3, arg4)
流程很清晰,但函数与函数之间会产生中间变量。这里为了方便取 x 加数字后缀形式的名字,日常编程时最好还是起个有意义点的名字,例如 x_after_func1 之类的。另一种简练的写法是:
y <- func4(func3(func2(func1(x, arg1), arg2), arg3), arg4)
代码更短,也没有中间变量了,但代价是重看代码时需要像剥洋葱一样从两边向中间一层层读。并且当函数名更长参数更多时,可读性会进一步恶化,列数也很容易超出屏幕的宽度。
这样看来似乎第一种风格更为妥当。不过,若是活用 magrittr 包里的管道符 %>% 的话,就能写出既清晰又简练的代码了。简单介绍一下 %>% 的功能:
x %>% f等价于f(x)。x %>% f(y)等价于f(x, y)。x %>% f(y, .)等价于f(y, x)。x %>% f(y, z = .)等价于f(y, z = x)。
即输入 x 通过管道 %>% 传给函数 f,f 里不用写 x,管道会自动把 x 作为 f 的第一个参数;如果 x 并非第一个参数,那么可以用占位符 . 代指 x。
应用了管道符后的代码风格是:
y <- x %>%
func1(arg1) %>%
func2(arg2) %>%
func3(arg3) %>%
func4(arg4)
格式整齐,代码顺序和操作顺序一致,语义清晰,没有多余的中间变量,强迫症患者感到十分舒适。这种写法的另一个好处是,增删函数就像增删空行一样简单,而前两种风格改起来就会十分烦人。
Pandas 中的管道
遗憾的是 Python 中并没有成熟的管道包,但有一种神似的写法:
x = 'fried chicken\n'
y = x.rstrip().replace('fried', 'roast').upper().rjust(20)
print(y)
ROAST CHICKEN
即对 x.rstrip() 方法返回的字符串调用 replace 方法,再对返回值调用 upper 方法,最后调用 rjust 方法,构成了方法链(method chaining)。这个写法看似简洁,实则局限很大:以一节节管道做比喻的话,R 中每节管子可以是任意函数,而 Python 中每节管子只能是输入管子的对象自带的方法。如果你想实现的操作不能用输入对象的方法达成,那么管道就连不起来,你还是得乖乖打断管道,在下一行调用函数或写表达式。
但细分到用 Pandas 包做数据分析的领域,基于方法链的管道已经完全够用了:绝大部分操作都可以用 DataFrame 或 Series 的方法实现,并且方法返回的结果依旧是 DataFrame 或 Series 对象,保证可以接着调用方法;外部函数用 map、apply、applymap 或 pipe 方法应用到数据上。下面以处理站点气象数据表格为例:
- 查询指定站点。
- 丢弃站点列。
- 将时间列转为
DatetimeIndex。 - 按时间排序。
- 去除时间上重复的记录。
- 设置时间索引。
- 将 999999 替换成 NaN。
- 重采样到逐小时分辨率并插值填充。
- 加入风速分量列。
先来个普通风格:
def wswd_to_uv(ws, wd):
'''风速风向转为uv分量.'''
wd = np.deg2rad(270 - wd)
u = ws * np.cos(wd)
v = ws * np.sin(wd)
return u, v
station = 114514
df.query('station == @station', inplace=True)
df.drop(columns='station', inplace=True)
df['time'] = pd.to_datetime(df['time'], format='%Y-%m-%d %H:%M')
df.sort_values('time', inplace=True)
df.drop_duplicates(subset='time', keep='last', inplace=True)
df.set_index('time', inplace=True)
df.mask(df >= 999999, inplace=True)
df = df.resample('H').interpolate()
df['u'], df['v'] = wswd_to_uv(df['ws'], df['wd'])
得益于很多方法自带原地修改的 inplace 参数,中间变量已经很少了。再来看看管道风格:
def set_time(df, fmt):
return df.assign(time=pd.to_datetime(df['time'], format=fmt))
def add_uv(df):
u, v = wswd_to_uv(df['ws'], df['wd'])
return df.assign(u=u, v=v)
dfa = (df
.query('station == @station')
.drop(columns='station')
.pipe(set_time, fmt='%Y-%m-%d %H:%M')
.sort_values('time')
.drop_duplicates(subset='time', keep='last')
.set_index('time')
.mask(lambda x: x >= 999999)
.resample('H').interpolate()
.pipe(add_uv)
)
个人感觉管道风格的格式更整齐,一眼就能看出每行的“动词”(方法)。去除了每行都有的 inplace 参数后,不仅视觉上更清爽,还保证了一套操作下来输入数据不会无缘无故遭到修改。接着再说说管道风格里的两个细节。
pipe
就是 Pandas 版的 %>%:
df.pipe(func)等价于func(df)。df.pipe(func, *args, **kwargs)等价于func(df, *args, **kwargs)。df.pipe((func, 'arg2'), arg1=a)等价于func(arg1=a, arg2=df)。
可以将复杂的多行运算打包成形如 func(df, *args, **kwargs) 的函数,然后结合 pipe 使用。前文的 set_time 和 add_uv 函数就是例子。
assign
assign 方法的功能就是无副作用的列赋值:复制一份对象自己,在列尾添加新列或是修改已有的列,然后返回这份拷贝:
# 相当于:
# dfa = df.copy()
# dfa['a'] = a
# dfa['b'] = b
dfa = df.assign(a=a, b=b)
# 相当于:
# df['a'] = a
# df['b'] = b
df.assign(a=a, b=b, inplace=True)
第一次看到 assign 时我只觉得多此一举,赋值不是用等号就可以吗?但后来我意识到它是搭配管道风格使用的:想要对管道内的中间变量做列赋值,同时不中断管道,就只能用 assign 方法。同时考虑到中间变量里的内容可能已经跟原始输入大不相同,assign 的参数还可以是以调用对象本身(即 self)为唯一参数的函数:
# 省略号表示略去的方法.
dfa = (df
...
.assign(u=uwind, v=vwind)
.assign(ws=lambda x: np.hypot(x['u'], x['v']))
...
)
这里不能写成 assign(ws=np.hypot(df['u'], df['v'])),因为 df 里本来是没有 u 和 v 的,但中间变量有,那么把匿名函数传给 assign 就可以解决这一问题。
不只是 assign,where 和 mask 等方法,乃至 loc 和 iloc 索引器都能接受函数(准确来说是 callable 对象),方便在管道风格中使用。
什么时候该用管道
管道并非优雅代码的万金油,而是有特定使用场景的:
-
输入经过一连串的操作得到一个输出的情况适合使用管道,输入和输出都很多时显然不太适合。
-
管道里的操作多于十个时会使 debug 变得很麻烦,因为缺少中间变量来定位 bug。建议当操作很多时适当分出中间变量,不要一个管道写到头。
-
方法链中对象的类型发生改变时建议将链条进行拆分,不然会令人迷惑。