Зачем вам считать строки ROW NUMBER

В прошлый раз мы говорили о использовании агрегирующих функциях с использованием окон (window function)  в PySpark. Сегодня поговорим об одной из ранжирующих функций ROW NUMBER, а также посмотрим, как ее можно использовать.

Данные по Титанику

Возьмем для примера датасет с затонувшим Титаником. Его можно взять тут. У него много столбцов, поэтому вы покажем только его схему.

df = spark.read.csv("titanic_dataset.csv", header=True, inferScheme=True)

"""
root
 |-- PassengerId: integer (nullable = true)
 |-- Survived: integer (nullable = true)
 |-- Pclass: integer (nullable = true)
 |-- Name: string (nullable = true)
 |-- Sex: string (nullable = true)
 |-- Age: double (nullable = true)
 |-- SibSp: integer (nullable = true)
 |-- Parch: integer (nullable = true)
 |-- Ticket: string (nullable = true)
 |-- Fare: double (nullable = true)
 |-- Cabin: string (nullable = true)
 |-- Embarked: string (nullable = true)
"""

Как использовать ROW NUMBER в PySpark

Функция ROW NUMBER присваивает каждой строке номер в порядке следования. Так, с ее помощью можно создать уникальные идентификаторы после сортировки ORDER BY, для каждой группы PARTITION BY или и то , и другое.

В Apache Spark обязательно нужно указать порядок ORDER BY, без него вылетит ошибка (хотя некоторые СУБД могут сами разрешить порядок).

Итак, допустим дадим номера в порядке следования, упорядочив данные по плате за проезд Fare. Тогда запрос SQL будет выглядеть так:

SELECT t.Name, t.Sex, t.Fare,
    ROW_NUMBER() OVER(ORDER BY t.Fare) AS rn
FROM titanic t

"""
+--------------------+------+---+
|                Name|  Fare| rn|
+--------------------+------+---+
| Leonard, Mr. Lionel|   0.0|  1|
|Harrison, Mr. Wil...|   0.0|  2|
|Tornquist, Mr. Wi...|   0.0|  3|
|"Parkes, Mr. Fran...|   0.0|  4|
|Johnson, Mr. Will...|   0.0|  5|
|Cunningham, Mr. A...|   0.0|  6|
|Campbell, Mr. Wil...|   0.0|  7|
|"Frost, Mr. Antho...|   0.0|  8|
| Johnson, Mr. Alfred|   0.0|  9|
|Parr, Mr. William...|   0.0| 10|
|Watson, Mr. Ennis...|   0.0| 11|
|Knight, Mr. Robert J|   0.0| 12|
|Andrews, Mr. Thom...|   0.0| 13|
|    Fry, Mr. Richard|   0.0| 14|
|Reuchlin, Jonkhee...|   0.0| 15|
| Betros, Mr. Tannous|4.0125| 16|
|Carlsson, Mr. Fra...|   5.0| 17|
|Nysveen, Mr. Joha...|6.2375| 18|
|Lemberopolous, Mr...|6.4375| 19|
|Holm, Mr. John Fr...|  6.45| 20|
+--------------------+------+---+
"""

Новый столбец rn и является тем самым столбцом с индивидуальными номерами.

В PySpark этот же запрос будет выглядеть так:

import pyspark.sql.functions as F
from pyspark.sql import Window

w = Window.orderBy("Fare")
df.select("Name", "Fare", F.row_number().over(w).alias("rn"))

Мы можем также задействовать группы PARTITION BY. Например, сгруппируем пассажиров Титаника по их классу Pclass:

w = Window.partitionBy("Pclass").orderBy("Fare")
df.select(
    "Name", "Pclass", "Sex", "Fare",
    F.row_number().over(w).alias("rn")
)

"""
+--------------------+------+------+-------+---+
|                Name|Pclass|   Sex|   Fare| rn|
+--------------------+------+------+-------+---+
|Harrison, Mr. Wil...|     1|  male|    0.0|  1|
|Parr, Mr. William...|     1|  male|    0.0|  2|
|Andrews, Mr. Thom...|     1|  male|    0.0|  3|
|    Fry, Mr. Richard|     1|  male|    0.0|  4|
|Reuchlin, Jonkhee...|     1|  male|    0.0|  5|
|Carlsson, Mr. Fra...|     1|  male|    5.0|  6|
|Colley, Mr. Edwar...|     1|  male|25.5875|  7|
| Baumann, Mr. John D|     1|  male| 25.925|  8|
|Leader, Dr. Alice...|     1|female|25.9292|  9|
|Swift, Mrs. Frede...|     1|female|25.9292| 10|
+--------------------+------+------+-------+---+
"""

Можем убедиться, что у другого класса своя нумерация:

df.select(
    "Name", "Pclass", "Sex", "Fare",
    F.row_number().over(w).alias("rn")
).where("Pclass == 2")

"""
+--------------------+------+----+----+---+
|                Name|Pclass| Sex|Fare| rn|
+--------------------+------+----+----+---+
|"Parkes, Mr. Fran...|     2|male| 0.0|  1|
|Cunningham, Mr. A...|     2|male| 0.0|  2|
|Campbell, Mr. Wil...|     2|male| 0.0|  3|
|"Frost, Mr. Antho...|     2|male| 0.0|  4|
|Watson, Mr. Ennis...|     2|male| 0.0|  5|
+--------------------+------+----+----+---+
"""

Можно поступить другим путем и взять несколько строк из каждой группы, как это показано далее.

Use case: взять по несколько строк из каждой группы

Что если нужно взять заданное количество строк из каждой группы? Тогда вам может понадобится функция ROW NUMBER.

В этом случае мы можем взять в виде подзапроса запрос с функцией ROW NUMBER и отфильтровать его. Например, вот так выглядит SQL-запрос, который позволит взять не более две строки из каждой группы:

SELECT * FROM (
    SELECT t.Name, t.Pclass, t.Fare,
    ROW_NUMBER() OVER(PARTITION BY Pclass ORDER BY t.Fare) AS rn
    FROM titanic t
) x
WHERE x.rn < 3;

"""
+--------------------+------+----+---+
|                Name|Pclass|Fare| rn|
+--------------------+------+----+---+
| Leonard, Mr. Lionel|     3| 0.0|  1|
|Tornquist, Mr. Wi...|     3| 0.0|  2|
|Harrison, Mr. Wil...|     1| 0.0|  1|
|Parr, Mr. William...|     1| 0.0|  2|
|"Parkes, Mr. Fran...|     2| 0.0|  1|
|Cunningham, Mr. A...|     2| 0.0|  2|
+--------------------+------+----+---+
"""

Тот же самый запрос в PySpark:

x = df.select(
    "Name", "Pclass", "Sex", "Fare",
    F.row_number().over(w).alias("rn")
)
x.where("rn < 3")

Графовые алгоритмы в Apache Spark

Код курса
GRAS
Ближайшая дата курса
20 июня, 2022
Длительность обучения
16 ак.часов
Стоимость обучения
40 000 руб.

А в следующей статье поговорим о функциях RANK. Еще больше подробностей об оконных функциях вы узнаете на наших образовательных курсах в лицензированном учебном центре обучения и повышения квалификации руководителей и IT-специалистов (менеджеров, архитекторов, инженеров, администраторов, Data Scientist’ов и аналитиков Big Data) в Москве:

Записаться на курс

Смотреть раcписание

Источники
  1. row_number

Добавить комментарий

Поиск по сайту