SQL、Pandas、Spark：窗口函數(shù)的3種實(shí)現(xiàn)
導(dǎo)讀
窗口函數(shù)是數(shù)據(jù)庫(kù)查詢中的一個(gè)經(jīng)典場(chǎng)景，在解決某些特定問(wèn)題時(shí)甚至是必須的。個(gè)人認(rèn)為，在單純的數(shù)據(jù)庫(kù)查詢語(yǔ)句層面【即不考慮DML、SQL調(diào)優(yōu)、索引等進(jìn)階】，窗口函數(shù)可看作是考察求職者SQL功底的一個(gè)重要方面。
前期個(gè)人以求職者身份參加面試時(shí)被問(wèn)及窗口函數(shù)的問(wèn)題，近期在作為面試官也提問(wèn)過(guò)這一問(wèn)題，但回答較為理想者居少。所以本文首先窗口函數(shù)進(jìn)行講解，然后分別從SQL、Pandas和Spark三種工具平臺(tái)展開(kāi)實(shí)現(xiàn)。
模擬問(wèn)題描述：
給定一組中學(xué)生的歷次語(yǔ)文月考成績(jī)表（每名學(xué)生含有4次成績(jī)），需要實(shí)現(xiàn)以下3個(gè)需求：
對(duì)每名學(xué)生的4次成績(jī)表分別進(jìn)行排序，排序后每人的成績(jī)排名1-2-3-4
求每名學(xué)生歷次月考成績(jī)的變化幅度，即本月較上個(gè)月的成績(jī)差值
求每名學(xué)生歷次月考成績(jī)中近3次平均分
注：文末有送書活動(dòng)！
MySQL8.0官方手冊(cè)中關(guān)于窗口函數(shù)的介紹
當(dāng)然，為了形象表達(dá)上述定義所言何物，這里還是進(jìn)一步給出一些配套插圖以便于理解。在給出具體配圖之前，首先要介紹與窗口函數(shù)相關(guān)的3個(gè)關(guān)鍵詞：
partition by：用于對(duì)全量數(shù)據(jù)表進(jìn)行切分（與SQL中的groupby功能類似，但功能完全不同），直接體現(xiàn)的是前面窗口函數(shù)定義中的“有關(guān)”，即切分到同一組的即為有關(guān)，否則就是無(wú)關(guān)；
order by：用于指定對(duì)partition后各組內(nèi)的數(shù)據(jù)進(jìn)行排序；
rows between：用于對(duì)切分后的數(shù)據(jù)進(jìn)一步限定“有關(guān)”行的數(shù)量，此種情景下即使partition后分到一組，也可能是跟當(dāng)前行的計(jì)算無(wú)關(guān)。
相應(yīng)的，這3個(gè)關(guān)鍵字在前面的數(shù)據(jù)樣表中可作如下配套解釋：
當(dāng)然，到這里還不是很理解窗口函數(shù)以及相應(yīng)的3個(gè)關(guān)鍵字也問(wèn)題不大，后續(xù)結(jié)合前述的三個(gè)實(shí)際需求再返過(guò)來(lái)看此圖多半會(huì)豁然開(kāi)朗。
上面是窗口函數(shù)的邏輯解釋，那么具體能用于實(shí)現(xiàn)什么功能呢？其實(shí)，窗口函數(shù)能實(shí)現(xiàn)什么功能則要取決于能搭配什么函數(shù)。仍然引用MySQL8.0官方文檔中的一幅圖例：
其中，上表所述的窗口函數(shù)主要分為兩大類：
排序類，包括row_number、rank、dense_rank等，也包括percent_rank、cume_dist等分布排序類
相對(duì)引用類，如lag、lead、first_value、last_value、nth_value等
除了這兩類專用窗口函數(shù)之外，還有廣義的聚合函數(shù)也可配套窗口函數(shù)使用，例如sum、avg、max、min等。
所以，現(xiàn)在來(lái)看前面提到的三個(gè)需求，就剛好是分別應(yīng)用這三類窗口函數(shù)的例子。【哪有什么剛好，不過(guò)是特意設(shè)計(jì)而已】
圍繞這三個(gè)需求，下面分別應(yīng)用SQL、Pandas和Spark三個(gè)工具予以實(shí)現(xiàn)。
既然窗口函數(shù)起源于數(shù)據(jù)庫(kù)，那么下面就首先應(yīng)用SQL予以實(shí)現(xiàn)。
注：以下所有SQL查詢語(yǔ)句實(shí)現(xiàn)均基于MySQL8.0。
Q1：求解每名同學(xué)歷次成績(jī)的排名。
A1：由于是區(qū)分每名同學(xué)進(jìn)行排序，所以需要依據(jù)uid字段進(jìn)行partition；進(jìn)一步地，按照成績(jī)進(jìn)行排序，所以order by字段即為score；最后，由于是要進(jìn)行排名，所以配套函數(shù)選擇row_number即可。注：row_number、rank和dense_rank的具體區(qū)別可參考?xì)v史文章：一文解決所有MySQL分類排名問(wèn)題。
查詢語(yǔ)句及查詢結(jié)果如下：
SELECT *, row_number() over(partition by uid order by score desc) as `rank` from score
Q2：求解每名同學(xué)歷次月考成績(jī)的差值，即本月成績(jī)-上月成績(jī)。
A2：首先，仍然是依據(jù)uid字段進(jìn)行partition；而后由于是要計(jì)算本月成績(jī)與上月成績(jī)的差值，所以此次的排序依據(jù)應(yīng)該是date；進(jìn)一步地，由于要計(jì)算差值，所以需要對(duì)每次月考成績(jī)計(jì)算其前一行的成績(jī)（在按照uid進(jìn)行切分并按照date排序后，上月成績(jī)即為當(dāng)前行的前一條記錄），所以配套函數(shù)即為lag。
給出具體實(shí)現(xiàn)SQL語(yǔ)句及查詢結(jié)果如下：
SELECT *, score - lag(score) over(partition by uid order by date) as score_diff from score
Q3：求解每名學(xué)生近3次月考成績(jī)的平均分。
A3：在前兩個(gè)需求的基礎(chǔ)上，易見(jiàn)，仍然是依據(jù)uid進(jìn)行partition、依據(jù)date進(jìn)行排序，并選用avg聚合函數(shù)作為配套窗口函數(shù)。進(jìn)一步地，由于此處限定計(jì)算近3次成績(jī)的平均分，所以除了partition和order by 兩個(gè)關(guān)鍵字外，還需增加rows between的限定。
具體SQL語(yǔ)句和查詢結(jié)果如下：
SELECT *, avg(score) over(partition by uid order by date rows between 2 preceding and current row) as avg_score3 from score
值得指出的是，對(duì)于每名學(xué)生，當(dāng)切分窗口不足指定窗口大?。茨繕?biāo)行數(shù)）時(shí)會(huì)按實(shí)際的數(shù)據(jù)進(jìn)行聚合，例如學(xué)生A，1月31日對(duì)應(yīng)的近3次平均分即為本月成績(jī)自身；2月28日對(duì)應(yīng)近3次平均分即為本月成績(jī)和上月成績(jī)的平均分，而3月31日和4月30日計(jì)算的近3次平均分則為真正意義上的3次成績(jī)均值。
Pandas作為Python數(shù)據(jù)分析與處理的主力工具，自然也是支持窗口函數(shù)的，而且花樣只會(huì)比SQL更多。對(duì)于上述三個(gè)需求，Pandas分別實(shí)現(xiàn)如下：
Q1：求解每名同學(xué)歷次成績(jī)的排名。
A1：雖然Pandas接口非常豐富，但用其實(shí)現(xiàn)分組排名貌似卻并不方便。不過(guò)也是可以的?；舅悸啡缦拢菏紫热匀环謩e用uid和score字段進(jìn)行分組和排序，而后通過(guò)對(duì)取值=1的常數(shù)列num進(jìn)行cumsum，即累加，即可獲取分組排名結(jié)果。其中，還可進(jìn)一步應(yīng)用assign函數(shù)實(shí)現(xiàn)鏈?zhǔn)秸{(diào)用，最終整個(gè)需求實(shí)現(xiàn)下來(lái)也是一行代碼即可！
具體Pandas實(shí)現(xiàn)代碼即結(jié)果如下：
df.assign(rank=df.assign(num=1).sort_values("score", ascending=False).groupby("uid")['num'].cumsum())
注：上述代碼應(yīng)用了assign實(shí)現(xiàn)鏈?zhǔn)秸{(diào)用，具體可參考文章Pandas用了一年，這3個(gè)函數(shù)是我的最愛(ài)……
Q2：求解每名同學(xué)歷次月考成績(jī)的差值，即本月成績(jī)-上月成績(jī)。
A2：對(duì)于這一特定需求，Pandas中實(shí)際上是內(nèi)置了偏移函數(shù)shift，專門用于求解當(dāng)前行的相對(duì)引用值。進(jìn)一步地，對(duì)于求解差分結(jié)果，還可直接用diff實(shí)現(xiàn)，其中diff就相當(dāng)于當(dāng)前行-shift(1)。
兩種API實(shí)現(xiàn)代碼即執(zhí)行結(jié)果分別如下：
# shift函數(shù)實(shí)現(xiàn)df.assign(score_diff=df["score"]-df.sort_values("date").groupby("uid")['score'].shift(1))# diff函數(shù)實(shí)現(xiàn)df.assign(score_diff=df.sort_values("date").groupby("uid")['score'].diff(1))
Q3：求解每名學(xué)生近3次月考成績(jī)的平均分。
A3：如果說(shuō)前兩個(gè)需求用Pandas實(shí)現(xiàn)都沒(méi)有很好體現(xiàn)窗口函數(shù)的話，那么這個(gè)需求可能才更貼近Pandas中窗口函數(shù)的標(biāo)準(zhǔn)用法——那就是用關(guān)鍵字rolling。rolling原義即有滾動(dòng)的意思，用在這里即表達(dá)滑動(dòng)窗口的意思，所以自然也就可以設(shè)置滑動(dòng)窗口的大小。至于SQL中窗口函數(shù)的另外兩個(gè)關(guān)鍵字partition和order則仍然需要借助Pandas的sort_values和gropupby來(lái)實(shí)現(xiàn)。另外，與SQL中類似，這里仍然是要用求均值函數(shù)來(lái)做為配套。
具體Pandas實(shí)現(xiàn)代碼如下：
df.assign(avg_score3=df.sort_values("date").groupby("uid").rolling(window=3, min_periods=1)['score'].mean().reset_index().set_index("level_1")['score'])
注：上述實(shí)現(xiàn)中用到了reset_index和set_index，其目的是為了保證滑窗聚合后保證順序不變。為了追求單行代碼實(shí)現(xiàn)，這里的寫法不再優(yōu)雅，并不提倡。
最后，選用Spark予以實(shí)現(xiàn)。應(yīng)該講，Spark.sql組件幾乎是完全對(duì)標(biāo)SQL語(yǔ)法的實(shí)現(xiàn)，這在窗口函數(shù)中也例外，包括over以及paritionBy、orderBy和rowsbetween等關(guān)鍵字的使用上。
注：在使用Spark窗口函數(shù)前，首先需要求引入窗口函數(shù)類Window。即
import org.apache.spark.sql.expressions.WindowQ1：求解每名同學(xué)歷次成績(jī)的排名。
A1：直接沿用SQL思路即可，需要注意Spark中的相應(yīng)表達(dá)。
代碼實(shí)現(xiàn)及相應(yīng)執(zhí)行結(jié)果如下：
df.select($"uid", $"date", $"score", row_number().over(Window.partitionBy("uid").orderBy($"score".desc)).as("rank"))
Q2：求解每名同學(xué)歷次月考成績(jī)的差值，即本月成績(jī)-上月成績(jī)。
A2：首先應(yīng)用lag算子求出上月成績(jī)，而后直接相減即可。
代碼及執(zhí)行結(jié)果如下：
df.select($"uid",$"date", $"score", ($"score"-lag($"score", 1).over(Window.partitionBy("uid").orderBy("date"))).as("score_diff"))
Q3：求解每名學(xué)生近3次月考成績(jī)的平均分。
A3：仍然沿用SQL中思路即可，只需增加rowsBetween函數(shù)。
代碼實(shí)現(xiàn)及執(zhí)行結(jié)果如下：
df.select($"uid",$"date", $"score", avg("score").over(Window.partitionBy("uid").orderBy("date").rowsBetween(-2, 0)).as("avg_score3"))
本文首先對(duì)窗口函數(shù)進(jìn)行了介紹，通過(guò)模擬設(shè)定3個(gè)實(shí)際需求問(wèn)題，分別基于SQL、Pandas和Spark三個(gè)工具平臺(tái)予以分析和實(shí)現(xiàn)。總體來(lái)看，SQL和Spark實(shí)現(xiàn)窗口函數(shù)的方式和語(yǔ)法更為接近，而Pandas雖然擁有豐富的API，但對(duì)于具體窗口函數(shù)功能的實(shí)現(xiàn)上卻不盡統(tǒng)一，而需靈活調(diào)用相應(yīng)的函數(shù)。當(dāng)然，窗口函數(shù)的功能還有很多，三個(gè)工具平臺(tái)的使用也遠(yuǎn)不止這些，但其核心原理則是大體相通的。
最后，感謝清華大學(xué)出版社為本公眾號(hào)讀者贊助《數(shù)據(jù)科學(xué)實(shí)用算法》一本，截止本周五（4月16日）早9點(diǎn)，公眾號(hào)后臺(tái)查看分享最多的前3名讀者隨機(jī)指定一人，中獎(jiǎng)讀者將在【小數(shù)志】讀者微信群中公布，若還未加群的可在公眾號(hào)菜單-"關(guān)于"中添加小編微信聯(lián)系入群。
相關(guān)閱讀：