pandas-udf函数定义和使用

本文介绍pandas下udf函数处理单列/多列数据返回单列多列结果的demo

import pandas as pd
def func1(x):
    y = x + 1
    return y
    
def func2(x1,x2):
	y = x1+x2
	return y

def func3(x1,x2):
	y1 = x1+x2
	y2 = x1-x2
	return y1,y2


df = pd.read_csv("test.csv")

# 简单处理
df.loc[:,'y1'] = df.c1.apply(lambda x: 1 if x > 0 else 0)

# 对单列值根据udf自定义操作
df.loc[:,"y2"] = df.c1.apply(lambda x: func1(x)) 

# 对多列值根据udf自定义操作
df['y3'] = df.apply(lambda x : func2(x['c1'], x['c2']),axis = 1)

# 对多列值操作返回多列结果
df[['word_list','unit']] = df.apply(lambda x : func3(x['c1'],x['c2']),axis = 1,result_type="expand")