Spark DataFrame 开发指南

DataFrame 是 Spark 在 RDD 之后新推出的一个数据集，从属于 Spark SQL 模块，适用于结构化数据。对于我们这些用惯了数据库表和散列/字典结构的人来说，尤为亲切。

女神镇楼

可以直接读取关系型数据库产生 DataFrame：

from pyspark.sql import SparkSession
spark = SparkSession \
        .builder \
        .appName("myapp") \
        .config("spark.sql.shuffle.partitions", 10) \
        .getOrCreate()
table = "(select * from users where province='jiangxi') as myusers"
df = spark.read \
                .format("jdbc") \
                .option("driver", "com.mysql.jdbc.Driver") \
                .option("url", "jdbc:mysql://"+dbhost+"/"+dbname) \
                .option("dbtable", table) \
                .option("user", dbuser) \
                .option("password", dbpass) \
                .option("partitionColumn", partitionColumn) \
                .option("lowerBound", lowerBound) \
                .option("upperBound", upperBound) \
                .option("numPartitions", "4") \
                .load()

几个选项说明一下：

• url - 数据库的 JDBC 连接串
• dbtable - 可以是表名，也可以是一个子查询。如果是子查询的话，必须用括号括起来，并加别名，参见上面代码示例。
• user - 数据库用户名
• password - 数据库密码
• partitionColumn - 用于并发分区的表字段。下面几个选项都是围绕这个字段来的。Spark 会根据分区数量按这个字段的上下限把取出来的数据等分成几份，并行处理。
• lowerBound - 字段下限
• upperBound - 字段上限
• numPartitions - 分区数量

得到的 DataFrame，可以通过字段名来引用列或者某行的值，如：

username = df.name
age = df["age"]

而不像以前用 RDD，老是 split 成元组，然后 0、1、2、3... 地引用，很容易搞错位置。

调用 show 方法可以查看 DataFrame 里的数据：

df.show()

如果数据当中有中文的话，show 方法会报编码错误，需要提前设置一下环境变量：

export PYTHONIOENCODING=utf8

然后，就可以像 RDD 一样，进行各种数据变化、聚合操作了。不过遍查文档，你也找不到 DataFrame 的 map 方法，你需要用 select 和 withColumn 这两个方法来替代。其实回想 DataFrame 在 Saprk SQL 这个模块里，所以它的很多行为都会非常像关系数据库的 SQL 查询。

可以把每个 DataFrame 都想象成一个临时表，select 方法，就是从这个表里选择出一些字段，做一点变换，变成另外一个 DataFrame（另外一个临时表）：

df2 = df.select(df['name'], df['age'] + 1)
df2.show()
# +-------+---------+
# |   name|(age + 1)|
# +-------+---------+
# |Michael|     null|
# |   Andy|       31|
# | Justin|       20|
# +-------+---------+

在实际使用中，我们经常会在某个 DataFrame 基础上加上一个变换后的字段，如果用 select 来写，就是这样：

import pyspark.sql.functions as F
df3 = df2.select("*", F.abs(df2.age).alias("abs_age"))

这里有三个新知识点：一个是在 select 中使用 "*"，这样可以避免我们把 df2 的所有字段重新敲一遍。然后我们新增了一个字段，是 age 的绝对值。我们这里用到了 abs 函数，它来自 pyspark.sql.functions，是 spark 内置的众多变换函数之一。pyspark.sql.functions 里提供了相当多的变换函数，可以在文档里查到。这些变换函数结合加减乘除这些运算，基本上可以做完成任何变换了。实在搞不定的，可以用我们后面讲的 UDF（用户自定义函数）。第三点是用到了 alias 函数，来为新变换出的字段命名，类似 SQL 语句中的 "as"。下图是内置函数文档目录：

内置函数文档目录

上面这段代码使用 withColumn 会更加简洁，withColumn 就是在一个 DataFrame 的基础上增加或替换一个字段：

import pyspark.sql.functions as F
df3 = df2.withColumn("abs_age", F.abs(df2.age))

然后有些字段变换成其他字段以后就没用了，可以 drop 掉它：

df4 = df.drop("name")
df5 = df.select("*", df.age+1).drop(df.age)

我们到哪了？我们还是在说 RDD 的 map 方法，可以用 DataFrame 的 select、withColumn、alias、drop、加减乘除、pyspark.sql.functions 组合来替代。看似是用很多东西才替代了 RDD 的 map，但是实际开发的时候会很省事，省去了在 map 函数中每次都要把需要的字段逐一返回了，特别是有时候只是把某些字段在 pair 的 key 和 value 间移动。后面还会看到，DataFrame 可以直接对任意字段做 groupBy，而不用先移动摆好 pair 的 key 和 value。

如果你需要做的变换实在太复杂，无法用加减乘除和 pyspark.sql.functions 来搞定，那么 DataFrame 也支持自定义函数：udf。udf 本身是一个普通的内置函数，可以用它来包装普通的 python 函数，生成 select 和 withColumn 支持的变换函数：

import pyspark.sql.functions as F
import pyspark.sql.types as T
def do_something(col1, col2, col3):
    return col1*col1 + col2 / col3
udf_dosth = F.udf(do_something, T.IntegerType())
df6 = df.select(df.name, udf_dosth(df.age, df.born_year, df.children_count))

好了，到此可以先松口气，休息一下了。filter 方法 DataFrame 直接支持，不需要寻找替代品：

df.filter(df['age'] > 21).show()
# +---+----+
# |age|name|
# +---+----+
# | 30|Andy|
# +---+----+

DataFrame 可以很方便地做 join、groupBy 等操作，就像在写 SQL：

people = sqlContext.read.parquet("...")
department = sqlContext.read.parquet("...")

people.filter(people.age > 30).join(department, people.deptId == department.id) \
  .groupBy(department.name, "gender").agg({"salary": "avg", "age": "max"})

其实可以真的把 DataFrame 注册成临时表，然后真的写 SQL：

df.createOrReplaceTempView("people")

sqlDF = spark.sql("SELECT * FROM people")
sqlDF.show()
# +----+-------+
# | age|   name|
# +----+-------+
# |null|Michael|
# |  30|   Andy|
# |  19| Justin|
# +----+-------+

agg 这个方法是用来对 groupBy 之后的数据集做聚合操作的（对应 RDD 的 reduceByKey 方法）。它支持 avg、max、min、sum、count 等操作。但是这几个操作在实际使用中是远远不够用的，这时候我们还是需要自定义函数的。这种自定义函数叫做 UDAF( User Defined Aggregate Function)。UDAF 只在 Spark 的 scala 和 Java 中支持，pyspark并不支持。在 Scala 中，你需要重载 UserDefinedAggregateFunction 这个类即可。本文就不具体展示了，留待我稍后一篇专门介绍 Scala Spark 的文章里细说。

好了，本文就到这里，有任何疑问欢迎留言探讨。

DataFrame API 参考： http://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.DataFrame
pyspark.sql.functions 参考：
http://spark.apache.org/docs/latest/api/python/pyspark.sql.html#module-pyspark.sql.functions