背景

pandas dataFrame 无法支持大量数据的计算，可以尝试 spark df 来解决这个问题。

一. xgboost 预测的例子

优化前

import xgboost as xgb
import pandas as pd
import numpy as np

# 加载模型
bst = xgb.Booster()
bst.load_model("xxx.model")

# 变量列表
var_list=[...]
df.rdd.map(lambda x : cal_xgb_score(x,var_list,ntree_limit=304)).toDF()

# 计算分数
def cal_xgb_score(x,var_list,ntree_limit=50):
    feature_count = len(var_list)
    x1 = pd.DataFrame(np.array(x).reshape(1,feature_count),columns=var_list)
    # 数据变化操作
    y1 = transformFun(x1)
    
    test_x = xgb.DMatrix(y1.drop(['mobile','mobile_md5'],xais=1),missing=float('nan'))
    y1['score'] = bst.predict(test_x,ntree_limit=ntree_limit)
    y2 = y1[['mobile','mobile_md5','score']]
    return {'mobile': str(y2['mobile'][0]),'mobille_md5':str(y2['mobile_md5'][0]),'score':float(y2['score'][0])}

每条数据都转化为 pd，增加了额外开销。

优化后:

df.rdd.mapPartitions(lambda x : cal_xgb_score(x,var_list,ntree_limit=304)).toDF()

def cal_xgb_score(x,var_list,ntree_limit=50):
    feature_count = len(var_list)
    //将 iterator 转为list 
    x1 = pd.DataFrame(list(x),columns=var_list)
    ...
    //将 pdf 转为字典
    return y1[['mobile','mobile_md5','score']].to_dict(orient='record')

二. toPandas 的例子

优化前：

df.toPandas()

优化后:

import pandas as pd
def _map_to_pandas(rdds):
    return [pd.DataFrame(list(rdds))]
    
def toPandas(df, n_partitions=None):
    if n_partitions is not None: df = df.repartition(n_partitions)
    df_pand = df.rdd.mapPartitions(_map_to_pandas).collect()
    df_pand = pd.concat(df_pand)
    df_pand.columns = df.columns
    return df_pand

# 98列，22W行，类型 array/string/Long/Int，分区数 200
df = spark.sql("...").sample(False,0.002)

df.cache()
df.count()

# 原生的 toPandas 方法
%timeit df.toPandas()

# 分布式的 toPandas
%timeit toPandas(df)

#使用 apache arrow，spark 版本2.3以上
spark.sql("set spark.sql.execution.arrow.enabled=true")
%timeit df.toPandas()

总结

一. xgboost 预测

数据处理速度从 120 record / min 提高到 3278 record / min

tips: 如果一个分区数据量过大将会导致 executor oom

二. spark dataframe 转 pandas dataframe

toPandas 返回的数据归根结底还是缓存在 driver 的内存中的，不建议返回过大的数据。