pyspark应用技巧

1. spark sdf和pandas pdf相互转化

一般spark sdf转化为pandas pdf使用sdf.toPandas(), pdf转化为sdf使用spark.createDataFrame(pdf)，但是直接转化中间的序列化和反序列化耗时很长，所以在执行转化的时候使用apache arrow进行加速

pyarrow版本 >= 0.8.0

spark-defaults.conf文件添加：

spark.sql.execution.arrow.enabled true

或者在设置spark conf时设置：

conf = SparkConf().setAppName("Test").setMaster("local[*]")
conf.set("spark.sql.execution.arrow.enabled", True)

别人的对比：

execution.arrow.enabled	pdf -> sdf	sdf -> pdf
false	4980ms	722ms
true	72ms	79ms

tips: 尽管转化速度提高了，但pdf是单核运算，并没有用到分布式处理，所以最好不要处理大数据量。
当计算不适用于用arrow优化的时候可以自动退回非arrow优化的方式，这是配置参数为spark.sql.execution.arrow.fallback.enabled

每批进行向量化计算的数据量由spark.sql.execution.arrow.maxRecordsPerBatch参数控制，默认10000条

2. sdf构建自定义函数时优先使用pandas_udf而不是udf

pandas udf建立在Apache arrow之上，带来了低开销， 高性能的udf，并且使用了pandas的向量化操作；而spark的udf是对每一条数据进行操作，这样就带来了性能的问题。但是pandas udf有一些数据类型不支持，例如：BinaryType，MapType, TimestampType 和嵌套的 StructType。

注意：有些低级的pyarrow版本在使用pandas_udf时会出错，因此最好使用比较高一点的版本
下面所有代码运行于linux系统中，python3.5包：numpy (1.17.0)，pandas (0.25.2)，pyarrow (0.13.0)

from pyspark import SparkConf
from pyspark.sql import SparkSession, Row
from pyspark.sql.functions import pandas_udf, PandasUDFType
import pyspark.sql.functions as F
from pyspark.sql.types import StringType

conf = SparkConf().setAppName("test").setMaster("local")
spark = SparkSession.builder.config(conf=conf).getOrCreate()

SCALAR

one or more pandas.Series -> one pandas.Series, 长度必须和原来的一致，2.4.3不支持MapType和StructType.
与dataframe.withColumn或dataframe.select一起使用

df = spark.createDataFrame([(1, 'goods'), (1, 'good'), (1, 'god'), (2, 'thanks'), (2, 'thank')], schema=['x', 'y'])
# to upper strings
@pandas_udf(StringType(), PandasUDFType.SCALAR)
def to_upper(s):
    return s.str.upper()

df.select(df.x, to_upper(df.y)).show()  # 1


df = spark.createDataFrame([[1, 2, 4], [-1, 2, 2]], ['a', 'b', 'c'])
# input multi-pandas.Series, pay attention to the returnType
@pandas_udf('double', PandasUDFType.SCALAR)
def fun_function(a, b, c):
    clip = lambda x: x.where(a >= 0, 0)
    return (clip(a) - clip(b)) / clip(c)

df.withColumn('d', fun_function(df.a, df.b, df.c)).show()  # 2


df = spark.createDataFrame([(1, [1, 2, 3]), (2, [3, 4, 5])], schema=['x', 'y'])
# process ArrayType
@pandas_udf(ArrayType(IntegerType()), PandasUDFType.SCALAR)
def lens(s):
    a = s.apply(lambda x: x * 2)
    return a

df.select(df.x, lens(df.y)).show()  # 3

1.
+---+-----------+                                                               
|  x|to_upper(y)|
+---+-----------+
|  1|      GOODS|
|  1|       GOOD|
|  1|        GOD|
|  2|     THANKS|
|  2|      THANK|
+---+-----------+
2.
+---+---+---+-----+
|  a|  b|  c|    d|
+---+---+---+-----+
|  1|  2|  4|-0.25|
| -1|  2|  2| null|
+---+---+---+-----+
3.
+---+----------+
|  x|   lens(y)|
+---+----------+
|  1| [2, 4, 6]|
|  2|[6, 8, 10]|
+---+----------+

GROUPED_MAP

one DataFrame -> one transformed DataFrame, 字段类型必须和原来数据一致对应，字段标签也必须一致对应

一般与GroupedData.apply一起使用

df = spark.createDataFrame([(1, 'goods'), (1, 'good'), (1, 'god'), (2, 'thanks'), (2, 'thank')], schema=['x', 'y'])
# the return type should be same with df
@pandas_udf("x int, y string", PandasUDFType.GROUPED_MAP)
def lens(pdf):
    y = pdf.y
    return pdf.assign(y=str(len(y)))

df.groupBy('x').apply(lens).show()  # 1


df = spark.createDataFrame([(1, [1, 2, 3]), (2, [3, 4, 5])], schema=['x', 'y'])
# use schema as returnType
@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def lens(pdf):
    y = pdf.y
    return pdf.assign(y=y*2)

df.groupBy('x').apply(lens).show()  # 2

   1.
   +---+---+                                                                       
   |  x|  y|
   +---+---+
   |  1|  3|
   |  1|  3|
   |  1|  3|
   |  2|  2|
   |  2|  2|
   +---+---+
   2.
   +---+----------+                                                                
   |  x|         y|
   +---+----------+
   |  1| [2, 4, 6]|
   |  2|[6, 8, 10]|
   +---+----------+

   

GROUPED_AGG

One or more pandas.Series -> A scalar，returnType必须是主类型，例如DoubleType，返回的常量可以是python的主类型(int, float)或者是numpy的数据类型(numpy.int64, numpy.float64)，2.4.3不支持MapType和StructType.

一般与pyspark.sql.GroupedData.agg()或pyspark.sql.Window一起使用

df = spark.createDataFrame([(1, 10), (1, 20), (1, 30), (2, 15), (2, 35)], schema=['x', 'y'])

@pandas_udf('float', PandasUDFType.GROUPED_AGG)
def gro(x):
    return x.mean()

df.groupBy('x').agg(gro(df.y)).show()  # 1

df = spark.createDataFrame([(1, [2, 3, 4], [1, 2, 3]), (1, [2, 3, 4], [2, 3, 4]), (2, [2, 3, 4], [3, 4, 5])], schema=['x', 'y', 'z'])

@pandas_udf("float", PandasUDFType.GROUPED_AGG)
def lens(y, z):
    a = 0
    b = 0
    for i in y:
        a += i.sum()
    for i in z:
        b += i.sum()
    return a + b

df.groupBy('x').agg(lens(df.y, df.z)).show()  # 2

   1.
   +---+------+                                                                    
   |  x|gro(y)|
   +---+------+
   |  1|  20.0|
   |  2|  25.0|
   +---+------+
   2.
   +---+----------+
   |  x|lens(y, z)|
   +---+----------+
   |  1|      33.0|
   |  2|      21.0|
   +---+----------+

向UDF传入其他参数

由于一些实际应用上的原因，需要向pandas_udf传入其他的参数，第一想到的就是使用偏函数functools.partial，但是使用functools.partial封装pandas_udf是一种错误的方法，例如，我想在pandas_udf中传入一个额外的z参数：

df = spark.createDataFrame([(1, 2), (1, 4), (2, 6), (2, 4)], schema=["x", "y"])

@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def f(pdf, z):
    y = pdf.y * 2 + z
    return pdf.assign(y=y)

df.groupBy(df.x).apply(partial(f, z=100)).show()

在函数f中有两个参数，而在pandas_udf装饰器中只有一个参数的返回类型，在使用functools.partial时会出现AttributeError: 'functools.partial' object has no attribute 'evalType'这个错误。

一种正确的方法就是使用另一个函数封装这个pandas_udf，并返回它：

df = spark.createDataFrame([(1, 2), (1, 4), (2, 6), (2, 4)], schema=["x", "y"])

def f(z):
    @pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
    def _internal_udf(pdf):
        y = pdf.y * 2 + z
        return pdf.assign(y=y)
    return _internal_udf

df.groupBy(df.x).apply(f(z=100)).show()

3. 使用Java UDF

PySpark: Java UDF Integration，建立好Java udf，生成jar包xxx.jar，运行spark-submit -jars xxx.jar pyspark_demo.py

4. 分发文件至spark的各个worker

当运行一个python项目的时候，特别是在linux系统下运行项目时，运行中找不到自定义模块可能是比较大的一个问题(解决这个问题最简单的方法就是以包的方式把整个项目安装到python中去。当然还要考虑项目包的冲突问题，但这个容易解决)。一般的方式就是添加文件执行路径，在linux shell中运行python文件，它是以当前路径进行文件查找的，为了适应在各个路径运行该python文件能够成功查找到它所依赖的文件，则需要在该python文件添加绝对路径，例如：

project
   |------base
   |	  |------__init__.py
   |	  |------a.py
   |		     |------class A
   |------utils
   |        |------__init__.py
   |        |------b.py
   |               |------class B
   |------main
   |        |------__init__.py
   |        |------test.py

# a.py
from utils.b import B

class A(object):
    def __init__(self):
        c = B()
        cc = c.get(4)
        self.b = cc

    def get(self):
        return self.b

# b.py
class B(object):
    def get(self, x):
        return x

# test.py
from base.a import A

if __name__ == "__main__":
    aa = A()
    print(aa.get())

project项目下有base和utils和main三个包，并把project部署到了linux的/home/aaa下(/home/aaa/project)，现在a.py用到了b.py下的class B，现在要在test.py下测试a.py，如果在main下直接运行test.py会出现ImportError: No module named base’，因为现在它只搜索mian路径下有没有XXX，而不是从project下搜索。那么现在添加sys.path.append("../"),再次在main下运行test.py则会成功。

现在的a.py:

import sys

sys.path.append("../")
from base.a import A

if __name__ == "__main__":
    aa = A()
    print(aa.get())

现在在main的上一级目录(即project下)运行test.py怎么样哪？运行python ./main/test.py这时仍出现ImportError: No module named 'base'，现在即使加上sys.path.append("../")也无用，因为他会从当前project路径向上一级路径查找。这时可以使用绝对路径：

import sys
import os
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../"))
print("当前路径{}".format(os.getcwd()))
print("查找路径{}".format(sys.path))
from base.a import A

if __name__ == "__main__":
    aa = A()
    print(aa.get())

这样，在任意路径下运行python /arbitrary/path/project/main/test.py都能成功运行。

以上是在linux中查找依赖文件的问题的解决方法，但是在spark中又出现了新的问题，当使用这种方法时，spark无法把这种路径传递到各个worker中去，如果在各个worker中需要一些其他依赖文件的时候，上述方法仍然失效，仍会出现ImportError: No module named 'XXX'，这时就需要把所依赖的文件分发到各个worker中去，在pyspark中使用的是addFile和addPyFile方法。

首先添加依赖的文件：

# 使用默认的sc, spark
sc.addPyFile("your/pyFile/path/a.py")  # 也可以是包含多个py文件的zip文件，省的一个一个添加

然后在map或其他算子函数内添加文件路径及导入包：

sys.path.insert(0, pyspark.SparkFiles.getRootDirectory())
from a import A

这样就能在worker中成功运行。

综上，最简单的方法就是把所有的包直接安装到python上，无需添加路径及使用addPyFile了！