如何在 PySpark 中使用窗口函數？

Question

我正在嘗試將一些 Windows 函數（ ntile和percentRank ）用於數據框，但我不知道如何使用它們。

任何人都可以幫我解決這個問題嗎？ 在Python API 文檔中沒有關於它的示例。

具體來說，我試圖在我的數據框中獲取數字字段的分位數。

我正在使用火花 1.4.0。

Answer 1

為了能夠使用窗口函數，您必須先創建一個窗口。 定義與普通 SQL 幾乎相同，這意味着您可以定義順序、分區或兩者。 首先讓我們創建一些虛擬數據：

import numpy as np
np.random.seed(1)

keys = ["foo"] * 10 + ["bar"] * 10
values = np.hstack([np.random.normal(0, 1, 10), np.random.normal(10, 1, 100)])

df = sqlContext.createDataFrame([
   {"k": k, "v": round(float(v), 3)} for k, v in zip(keys, values)])

確保您使用的是HiveContext （僅限 Spark < 2.0）：

from pyspark.sql import HiveContext

assert isinstance(sqlContext, HiveContext)

創建一個窗口：

from pyspark.sql.window import Window

w =  Window.partitionBy(df.k).orderBy(df.v)

這相當於

(PARTITION BY k ORDER BY v) 

在 SQL 中。

根據經驗，窗口定義應始終包含PARTITION BY子句，否則 Spark 會將所有數據移動到單個分區。 某些函數需要ORDER BY ，而在不同情況下（通常是聚合）可能是可選的。

還有兩個可選項可用於定義窗口跨度 - ROWS BETWEEN和RANGE BETWEEN 。 在這種特殊情況下，這些對我們沒有用。

最后，我們可以將其用於查詢：

from pyspark.sql.functions import percentRank, ntile

df.select(
    "k", "v",
    percentRank().over(w).alias("percent_rank"),
    ntile(3).over(w).alias("ntile3")
)

請注意， ntile與分位數沒有任何關系。