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