介绍一下 pyspark 的 DataFrame 基础操作。
一、选择和访问数据
PySpark DataFrame 是惰性计算的，简单地选择一列不会触发计算，但它会返回一个 Column 实例。并且，大多数按列操作都返回 Column 实例。
df.a
# 输出
Column<b'a'>
from pyspark.sql import Column
from pyspark.sql.functions import upper
type(df.c) == type(upper(df.c)) == type(df.c.isNull())
# 输出
True
- DataFrame.select() : 获取指定的 Column。
df.select(df.c).show()
# 输出
+-------+
| c|
+-------+
|string1|
|string2|
|string3|
+-------+
- DataFrame.withColumn() : 更新或添加 Column。
df.withColumn('upper_c', upper(df.c)).show()
# 输出
+---+---+-------+----------+-------------------+-------+
| a| b| c| d| e|upper_c|
+---+---+-------+----------+-------------------+-------+
| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|STRING1|
| 2|3.0|string2|2000-02-01|2000-01-02 12:00:00|STRING2|
| 3|4.0|string3|2000-03-01|2000-01-03 12:00:00|STRING3|
+---+---+-------+----------+-------------------+-------+
- DataFrame.filter() : 筛选数据。
df.filter(df.a == 1).show()
# 输出
+---+---+-------+----------+-------------------+
| a| b| c| d| e|
+---+---+-------+----------+-------------------+
| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|
+---+---+-------+----------+-------------------+
二、应用函数
PySpark 支持各种 UDF 和 API，以允许用户执行 Python 原生函数。例如，下面的示例允许用户在 Python 原生函数中直接使用pandas Series中的 API。
import pandas as pd
from pyspark.sql.functions import pandas_udf
@pandas_udf('long')
def pandas_plus_one(series: pd.Series) -> pd.Series:
# Simply plus one by using pandas Series.
return series + 1
df.select(pandas_plus_one(df.a)).show()
# 输出
+------------------+
|pandas_plus_one(a)|
+------------------+
| 2|
| 3|
| 4|
+------------------+
另一个例子是DataFrame.mapInPandas允许用户直接使用pandas DataFrame中的 API，而不受结果长度等任何限制。
def pandas_filter_func(iterator):
for pandas_df in iterator:
yield pandas_df[pandas_df.a == 1]
df.mapInPandas(pandas_filter_func, schema=df.schema).show()
# 输出
+---+---+-------+----------+-------------------+
| a| b| c| d| e|
+---+---+-------+----------+-------------------+
| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|
+---+---+-------+----------+-------------------+
三、数据分组操作
PySpark DataFrame 还提供了一种使用常用方法拆分-应用-组合策略来处理分组数据的方法。它按特定条件对数据进行分组，对每个组应用一个函数，然后将它们组合回 DataFrame。
df = spark.createDataFrame([
['red', 'banana', 1, 10], ['blue', 'banana', 2, 20], ['red', 'carrot', 3, 30],
['blue', 'grape', 4, 40], ['red', 'carrot', 5, 50], ['black', 'carrot', 6, 60],
['red', 'banana', 7, 70], ['red', 'grape', 8, 80]], schema=['color', 'fruit', 'v1', 'v2'])
df.show()
# 输出
+-----+------+---+---+
|color| fruit| v1| v2|
+-----+------+---+---+
| red|banana| 1| 10|
| blue|banana| 2| 20|
| red|carrot| 3| 30|
| blue| grape| 4| 40|
| red|carrot| 5| 50|
|black|carrot| 6| 60|
| red|banana| 7| 70|
| red| grape| 8| 80|
+-----+------+---+---+
分组，然后用avg()函数应用于结果组，输出平均值。同样的也可以来计算max()、min()等等。
df.groupby('color').avg().show()
# 输出
+-----+-------+-------+
|color|avg(v1)|avg(v2)|
+-----+-------+-------+
| red| 4.8| 48.0|
|black| 6.0| 60.0|
| blue| 3.0| 30.0|
+-----+-------+-------+
还可以使用 pandas API 对每个分组应用 Python 自定义函数。
def plus_mean(pandas_df):
return pandas_df.assign(v1=pandas_df.v1 - pandas_df.v1.mean())
df.groupby('color').applyInPandas(plus_mean, schema=df.schema).show()
# 输出
+-----+------+---+---+
|color| fruit| v1| v2|
+-----+------+---+---+
| red|banana| -3| 10|
| red|carrot| -1| 30|
| red|carrot| 0| 50|
| red|banana| 2| 70|
| red| grape| 3| 80|
|black|carrot| 0| 60|
| blue|banana| -1| 20|
| blue| grape| 1| 40|
+-----+------+---+---+
cogroup 和 applyInPandas 操作。
df1 = spark.createDataFrame(
[(20000101, 1, 1.0), (20000101, 2, 2.0), (20000102, 1, 3.0), (20000102, 2, 4.0)],
('time', 'id', 'v1'))
df2 = spark.createDataFrame(
[(20000101, 1, 'x'), (20000101, 2, 'y')],
('time', 'id', 'v2'))
def asof_join(l, r):
return pd.merge_asof(l, r, on='time', by='id')
df1.groupby('id').cogroup(df2.groupby('id')).applyInPandas(
asof_join, schema='time int, id int, v1 double, v2 string').show()
# 输出
+--------+---+---+---+
| time| id| v1| v2|
+--------+---+---+---+
|20000101| 1|1.0| x|
|20000102| 1|3.0| x|
|20000101| 2|2.0| y|
|20000102| 2|4.0| y|
+--------+---+---+---+
