Databricks 第4篇:pyspark.sql 分组统计和窗口

对数据分析时,通常需要对数据进行分组,并对每个分组进行聚合运算。在一定意义上,窗口也是一种分组统计的方法。

分组数据

DataFrame.groupBy()回的是GroupedData类,可以对分组数据应用聚合函数、apply()函数和pivot()函数。

常用的聚合函数是:

  • count():统计数量
  • mean(*cols), avg(*cols):计算均值
  • max(*cols),min(*cols):计算最大值和最小值
  • sum(*cols):计算累加和

举个例子,对DataFrame计算最大的age和height列的值:

df.groupBy().max('age', 'height').collect() 

一,聚合函数

对于聚合函数,还可以使用pandas中的函数,这需要使用GroupedData类的agg(*exprs)函数,该函数的作用是计算聚合值,并返回DataFrame对象。

可以用于agg()函数中的聚合函数主要分为两类:

  • 内置的聚合函数:avg, max, min, sum, count
  • 分组聚合的pandas UDF:pyspark.sql.functions.pandas_udf()

对于内置的聚合函数,可以通过pyspark.sql.functions来导入:

gdf = df.groupBy(df.name)
from pyspark.sql import functions as F
sorted(gdf.agg(F.min(df.age)).collect())
#[Row(name='Alice', min(age)=2), Row(name='Bob', min(age)=5)]

这里重点介绍如何创建一个pandas UDF,Pandas UDF由Spark使用Arrow来传输数据,并通过Pandas对数据进行矢量化操作。在创建Pandas UDF时,需要通过pandas_udf作为修饰器或包装函数。

pyspark.sql.functions.pandas_udf(f=None, returnType=None, functionType=None)

参数注释:

  • f:UDF
  • returnType:UDF的返回值类型
  • functionType:一个枚举值,它的枚举类型是:pyspark.sql.functions.PandasUDFType,默认值是SCALAR,返回标量值。

举个例子,创建一个UDF,统计字符的个数。

在修饰器中定义函数的返回值类型是int,参数的模式是接收一个序列,返回一个序列,序列中的元素的数据类型是由修饰器决定的。

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf(IntegerType())
def slen(s: pd.Series) -> pd.Series:
    return s.str.len()

在定义函数时,显式指定输入参数的类型是MyType,函数返回值的类型是str:

# 输入参数类型提示为MyType,函数返回类型提示为str
def foo(name: MyType) -> str: 
    return str(name)

二,用户自定义的pandas函数

pyspark共支持5种UDF的模式,分别表示从形参到返回值的模式

模式1:从DataFrame到DataFrame

@pandas_udf("col1 string, col2 long")
def func(s1: pd.Series, s2: pd.Series, s3: pd.DataFrame) -> pd.DataFrame:
    s3['col2'] = s1 + s2.str.len()
    return s3

模式2:从Series到Series

@pandas_udf("string")
def to_upper(s: pd.Series) -> pd.Series:
    return s.str.upper()

模式3:从Series到Scalar,这种模式就是聚合函数,把多个值按照公式转换为标量值。

@pandas_udf("double")
def mean_udf(v: pd.Series) -> float:
    return v.mean()

模式4:Iterator[pandas.Series] -> Iterator[pandas.Series]

from typing import Iterator
@pandas_udf("long")
def plus_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    for s in iterator:
        yield s + 1

df = spark.createDataFrame(pd.DataFrame([1, 2, 3], columns=["v"]))
df.select(plus_one(df.v)).show()
+-----------+
|plus_one(v)|
+-----------+
|          2|
|          3|
|          4|
+-----------+

模式5:Iterator[Tuple[pandas.Series, …]] -> Iterator[pandas.Series]

from typing import Iterator, Tuple
from pyspark.sql.functions import struct, col
@pandas_udf("long")
def multiply(iterator: Iterator[Tuple[pd.Series, pd.DataFrame]]) -> Iterator[pd.Series]:
    for s1, df in iterator:
        yield s1 * df.v

df = spark.createDataFrame(pd.DataFrame([1, 2, 3], columns=["v"]))
df.withColumn('output', multiply(col("v"), struct(col("v")))).show()
+---+------+
|  v|output|
+---+------+
|  1|     1|
|  2|     4|
|  3|     9|
+---+------+

三,apply(udf)函数和applyInPandas(func, schema)

apply()和applyInPandas()函数的作用是:对每个分组应用函数,计算每个分组的聚合值。

apply(udf)函数使用 pyspark.sql.functions.pandas_udf() 作为参数,applyInPandas(func, schema)函数使用python 原生函数作为参数。

例如,apply()函数使用pandas_udf作为参数:

from pyspark.sql.functions import pandas_udf, PandasUDFType
df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))
@pandas_udf("id long, v double", PandasUDFType.GROUPED_MAP)  
def normalize(pdf):
    v = pdf.v
    return pdf.assign(v=(v - v.mean()) / v.std())
df.groupby("id").apply(normalize).show()  
+---+-------------------+
| id|                  v|
+---+-------------------+
|  1|-0.7071067811865475|
|  1| 0.7071067811865475|
|  2|-0.8320502943378437|
|  2|-0.2773500981126146|
|  2| 1.1094003924504583|
+---+-------------------+

例如,applyInPandas()使用python 原生的函数作为参数:

import pandas as pd  
from pyspark.sql.functions import pandas_udf, ceil
df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))  
def normalize(pdf):
    v = pdf.v
    return pdf.assign(v=(v - v.mean()) / v.std())
df.groupby("id").applyInPandas(
    normalize, schema="id long, v double").show()  
+---+-------------------+
| id|                  v|
+---+-------------------+
|  1|-0.7071067811865475|
|  1| 0.7071067811865475|
|  2|-0.8320502943378437|
|  2|-0.2773500981126146|
|  2| 1.1094003924504583|
+---+-------------------+

四,pivot()函数

从当前的DataFrame种透视一列,并执行指定的聚合操作。

pivot(pivot_col, values=None)

参数注释:

  • pivot_col:指定用于透视的列
  • values:被旋转为列的值列表,该参数如果为None,表示旋转列的所有值。

举个例子,按照year进行分组,把course列种的值透视为列,并计算earnings列的累加值:

df4.groupBy("year").pivot("course", ["dotNET", "Java"]).sum("earnings").collect()
#[Row(year=2012, dotNET=15000, Java=20000), Row(year=2013, dotNET=48000, Java=30000)]

df4.groupBy("year").pivot("course").sum("earnings").collect()
#[Row(year=2012, Java=20000, dotNET=15000), Row(year=2013, Java=30000, dotNET=48000)]

窗口函数

用于定义DataFrame的窗口,并对窗口进行计算。在进行窗口移动值,窗口的当前行(currentRow)的位置是0,如果position小于0,表示在当前行之前,如果position大于0,表示在当前行之后。

Window的位置属性:

  • Window.unboundedPreceding:窗口的第一行
  • Window.unboundedFollowing:窗口的最后一行
  • Window.currentRow:窗口的当前行

通过窗口函数来创建窗口:

  • partitionBy(*cols):分区
  • orderBy(*cols):排序
  • rangeBetween(start, end):start和end是相对于current row的位置,
  • rowsBetween(start, end):start和end是相对于current row的位置,

举个例子,利用这四个函数来创建窗口:

# ORDER BY date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
window = Window.orderBy("date").rowsBetween(Window.unboundedPreceding, Window.currentRow)

# PARTITION BY country ORDER BY date RANGE BETWEEN 3 PRECEDING AND 3 FOLLOWING
window = Window.orderBy("date").partitionBy("country").rangeBetween(-3, 3)

 

 

 

参考文档:

pyspark.sql module

发表评论

评论已关闭。

相关文章