優化 PySpark 數據幀中的過濾器 + 更新連接循環

Question

我們有一個 PySpark dataframe 代表一所擁有 35 名學生和 3 個班級的學校，每行代表一個學生。 Class #1 有 20 名學生，Class #2 有 10 名學生，Class #3 有 5 名學生。

我們想要比較這三個班級，因此我們將為每個 class 中的至少 4 名學生和最大 class 大小的一半分配考試。 這些學生必須隨機選擇。

如何盡可能高效地執行此任務？ 我在這里分享我迄今為止構建的代碼。

students = students.withColumn("exam", sf.lit(None))
for class_ in students.select("classroom").distinct().collect():

    for group_ in ['TEST', 'NO_TEST']:
        subdf = students.filter(sf.col("classroom") == class_[0])

            if group_ == 'TEST':
                subdf_group = subdf.sample(False, 0.5).limit(4) \
                    .withColumn("exam", sf.lit("EXAM"))

             else:
                 subdf_group = subdf.filter(sf.isnull(sf.col("exam"))) \
                    .withColumn("exam", sf.lit("NO_EXAM"))

         students = self.update_df(students, subdf_group)


def update_df(self, df_, new_df_):
    """
    A left join that updates the values from the students df with the new
    values on exam column.
    """
    out_df = df_.alias('l') \
        .join(new_df_.select("student_id", "exam").alias('r'),
              on="student_id", how="left").select(
        "student_id",
        self.update_column("exam")
    )
    return out_df


def update_column(column_name: str, left: str ='l', right: str ='r'):
    """
    When joining two dfs with same column names, we keep the column values from
    the right dataframe when values on right are not null, else, we keep the
    values on the left column.
    """
    return sf.when(~sf.isnull(sf.col(f'{right}.{column_name}')),
                   sf.col(f'{right}.{column_name}')) \
        .otherwise(sf.col(f'{left}.{column_name}')).alias(column_name)

這是一個玩具示例。 實際上，我們在 dataframe 中有 135 個類和總共 400 萬行，並且使用我上面共享的代碼，任務運行不佳。

Answer 1

您共享的代碼本質上非常不活躍。
最佳性能是在不使用 python 循環和多次隨機播放時。 在共享的代碼中，python 循環創建了一個 Spark 計划，其中包含數據集中的類數 X 過濾器和連接（連接 = 重洗牌操作），因此性能不佳。

假設你有一個 dataframe 有兩列["class", "student"]我將如何使用window 函數來做到這一點。
Spark 會將每個 window function 分區（在我們的案例類中）發送到不同的執行器，因此您將並行化采樣，而無需每次都過濾大 df。

from pyspark.sql.functions import col, row_number, rand, count
from pyspark.sql.window import Window

df \
.select(
    "*",
    row_number().over(Window.partitionBy('class').orderBy(rand(123))).alias('random_position'),
    count("*").over(Window.partitionBy('class')).alias('num_students_in_class'),
    ) \
.withColumn(
    "takes_test",
    (col("random_position") <= 4) &
     ((col("num_students_in_class") / 2) > col("random_position")))