我们有一个PySpark Dataframe ，它表示一所有35名学生和3个班级的学校，每行代表一个学生。1班有20名学生，2班有10名学生，3班有5名学生。

我们想要比较这三个班级，因此我们将分配一个考试给每个班级中最少4名学生和最多班级人数的一半。这些学生必须随机挑选。

如何才能尽可能高效地执行这项任务？我在这里分享我到目前为止构建的代码。

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万行，使用我在上面分享的代码，任务运行得很差。