Databricks 第四篇：分組統計和視窗

悅光陰發表於2021-01-11

對資料分析時，通常需要對資料進行分組，並對每個分組進行聚合運算。在一定意義上，視窗也是一種分組統計的方法。

分組資料

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

相關文章