pyspark выбирает определенные строки, которые имеют больше совпадающих полей столбцов

Question

У меня есть пример таблицы, как показано ниже (у меня есть 1 миллион таких строк), из этого мне нужно выбрать строки для нового кадра данных на основе условия ниже,

1. Я должен выбрать 1000 лучших учеников, которые посещали больше классов

2. Лучшие 1000 студентов, которые посещали 1-й класс, в 2,3,4 раза больше по сравнению с другими

, поэтому в моем примере мне нужно сохранить все строки студента 123 и 678 в другом фрейме данных

Я не мог получить правильную логику

Answer 1

Ниже приведено решение вашей проблемы, пожалуйста, дайте мне знать, если это вам поможет или нет

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

attended_more_classes = df.filter(
    F.col("check_in") == "y"
).groupby(
    "id"
).agg(
    F.countDistinct(F.col("class")).alias("class_count")
)

win = Window.partitionBy("id").orderBy(F.col("class_count").desc())

attended_more_classes = attended_more_classes.withColumn(
    "rank",
    F.rank().over(win)
).withColumn(
    "attended_more_class",
    F.when(
        F.col("rank")<=1000,
        F.lit("Y")
    )
)

# result of first part
attended_more_classes.show()

# answer start for second question

win2 = Window.partitionBy("id", "class").orderBy(F.col("class_count").desc())

filtered_students = df.filter(F.col("class").isin(1,2,3,4)).select("id").distinct()

aggregated_data2 = df.filter(
    F.col("check_in") == "y"
).groupby(
    "id",
    "class"
).agg(
    F.count(F.col("check_in")).alias("class_count")
).withColumn(
    "max_class",
    F.first(F.col("class")).over(win)
)

attend_more_class2 = aggregated_data2.join(
    filtered_students,
    on = "id",
    how = "inner"
)

attend_more_class23 = aggregated_data2.filter(
    F.col("max_class").isin(1,2,3,4)
).withColumn(
    "rank",
    F.rank().over(win2)
).withColumn(
    "attended_more_class",
    F.when(
        F.col("rank")<=1000,
        F.lit("Y")
    )
)

# answer of second part
attend_more_class23.show()