Как использовать функции окна в PySpark?

Я пытаюсь использовать некоторые функции Windows (ntile и percentRank) для фрейма данных, но я не знаю, как их использовать.

Может кто-нибудь помочь мне с этим, пожалуйста? В документации по Python API нет примеров по этому поводу.

В частности, я пытаюсь получить квантили числового поля в моем фрейме данных.

Я использую спарк 1.4.0.

Ответ 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 не связан с ntile.