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

Question

У меня есть датафрейм, который выглядит так:

Genres | Year | Number_Movies
Drama  |2015  | 705
Romance|2015  | 203
Comedy |2015  | 586
Drama  |2014  | 605
Romance|2014  | 293
Comedy |2014  | 786

Я хотел бы вернуть пол по году, в котором есть максимальное количество фильмов:

Genres | Year | Number_Movies
Drama  |2015  | 705
Comedy |2014  | 786

Пожалуйста, помогите, если это возможно. Большое спасибо.

Answer 1

Вот несколько вариантов, которые могут решить эту проблему -

df = spark.createDataFrame([('Drama',2015,705),('Romance',2015,203),('Comedy',2015,586),('Drama',2014,605),('Romance',2014,293),('Comedy ',2014,786)],['Genres','Year','Number_Movies'])

Первый вариант: Определить ранг с помощью оконной функции (разделить на - Год и упорядочить на - Number_Movies desc).Самые высокие Number_Movies каждый год получат ранг "1".

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number,desc

w = Window.partitionBy("Year").orderBy(desc("Number_Movies"))
rank = row_number().over(w).alias('rank')

df.withColumn("rank", rank)\
    .where("rank=1")\
    .drop("rank")\
    .show()

#+-------+----+-------------+
#| Genres|Year|Number_Movies|
#+-------+----+-------------+
#|Comedy |2014|          786|
#|  Drama|2015|          705|
#+-------+----+-------------+

Второй вариант: Получить максимальное количество Number_Movies для каждого года и самостоятельно объединиться с фреймом данных, чтобы получить жанры.

from pyspark.sql.functions import max,col

joining_condition = [col('a.Year') == col('b.Year'), col('a.max_Number_Movies') == col('b.Number_Movies')]

df.groupBy("Year").\
    agg(max("Number_Movies").alias("max_Number_Movies")).alias("a").\
    join(df.alias("b"),  joining_condition).\
    selectExpr("b.Genres","b.Year","b.Number_Movies").\
    show()

#+-------+----+-------------+
#| Genres|Year|Number_Movies|
#+-------+----+-------------+
#|Comedy |2014|          786|
#|  Drama|2015|          705|
#+-------+----+-------------+