关于python:在PySpark的DataFrame的某些行上应用操作(其结果取决于整个DataFrame中的信息)

Applying an operation (whose result depends on the information in the entire DataFrame) on certain rows of the DataFrame in PySpark

我目前正在努力实施Spark-y方法来使用PySpark进行操作。我有一个具有以下结构的大型DataFrame(约500,000行)

1
2
3
4
5
6
7
   ID    DATE    CHANGE    POOL
   -----------------------------
 1 ID1   DATE1   CHANGE1   POOL1
 2 ID2   DATE2   CHANGE2   POOL2
 3 ID3   DATE3   CHANGE3   POOL3
 4 ID4   DATE4   CHANGE4   POOL4
 ....

其中ID是唯一的,DATE不一定彼此等距(并且有可能重复),CHANGE可以是任意数量的float类型(通常但并非总是如此) (例如-500.0到500.0之间),并且POOL可以是空列表[]或长度可变的非空列表,具体取决于特定的行。我们可以安全地假定DataFrame按DATE

排序

我想通过以下操作替换POOL列中的空列表。我只考虑在W天之内的事件(例如,W是180天)。

  • 对于POOL为空的行,请注意该行的DATE
  • 我从步骤1中提到的日期开始,选择原始DataFrame中所有过去W天的窗口中的所有行。
  • 从这个子集中,我收集所有CHANGE的列表(可能使用collect_list()函数)。
  • 我将此列表分配给与步骤1中提到的行相对应的POOL列。
  • 我试图通过窗口函数来(部分)实现这一点,如下所示(我可以用partitionBy()安全地删除行)

    1
    2
    3
    4
    5
    6
    7
    POOL_DAYS = 180
    days = lambda i: i * 86400
    window_glb = Window\\
                     .partitionBy()\\
                     .orderBy(col('DATE').cast('long'))\\
                     .rangeBetween(-days(POOL_DAYS), 0)
    df = df.withColumn('POOL', collect_list('CHANGE').over(window_glb))

    但由于此错误而崩溃

    1
    2
    3
    Py4JJavaError: An error occurred while calling o1859.collectToPython.
    : org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 203.1 failed 4 times, most recent failure: Lost task 0.3 in stage 203.1 (TID 6339, xxxxxxxxxx.com, executor 1):
    java.lang.IllegalArgumentException: Size exceeds Integer.MAX_VALUE

    在我看来,这似乎是内存不足的问题。为了对此进行测试,我对原始DataFrame的一个较小子集使用了相同的操作,并成功完成了该操作。

    请注意,我确实了解到上述实现对所有行都重复了该操作,而没有考虑POOL列是否为空,但是我首先想了解一下是否可以使其正常工作,然后再继续进行操作下一步;看来此方法不适用于大型DataFrame。

    蛮力(而不是完全使用Spark-y)方法是对具有空POOL的行进行for循环,并将以下常规(而非UDF)函数应用于原始DataFrame

    1
    2
    3
    4
    5
    6
    def get_pool(row, window_size):
        end_date = row.DATE
        start_date = end_date - datetime.timedelta(days = window_size)
        return df.where((col('DATE') < end_date) & (col('DATE') >= start_date))\\
                 .agg(collect_list('CHANGE'))\\
                 .rdd.flatMap(list).first()

    但是,预计需要很长时间才能完成。

    更新:我也尝试过使用UDF

    1
    2
    3
    4
    5
    6
    7
    def get_pool(date, window_size):
        end_date = date
        start_date = end_date - datetime.timedelta(days = window_size)
        return df.where((col('DATE') < end_date) & (col('DATE') >= start_date))\\
                 .agg(collect_list('CHANGE'))\\
                 .rdd.collect()
    get_pool_udf = udf(get_pool, ArrayType(DoubleType()))

    但是由于函数的结果取决于整个(或大部分)DataFrame上的信息,因此它会因错误而停止

    1
    2
    PicklingError: Could not serialize object: Py4JError: An error occurred while calling o1857.__getnewargs__. Trace:
    py4j.Py4JException: Method __getnewargs__([]) does not exist

    第一个基于窗口函数的方法,蛮力方法和损坏的UDF的尝试,我该如何实现?我需要能够对从感兴趣的日期起某个日期范围内的DataFrame的子集进行分组和处理;如何使用Spark进行此操作?

    我对Spark(和PySpark)非常陌生,并且很难思考这个问题。我非常感谢您的帮助。谢谢!


    我正在回答我自己的问题:我选择POOL列为空的IDDATE,并执行两个DataFrame的联接,同时限制那些属于感兴趣的日期范围。按ID分组并将每个组中的CHANGE收集为列表,从而为POOL为空列表的行将所需的列表分配给POOL列:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    renamed_df = df.select('DATE', 'CHANGE')\\
                   .withColumnRenamed('DATE', 'r_DATE')\\
                   .withColumnRenamed('CHANGE', 'r_CHANGE')

    no_pool_df = df.where(size('POOL') == 0)

    cross_df = no_pool_df\\
                .join(renamed_df,\\
                        (renamed_df.r_DATE >= date_sub(no_pool_df.DATE, POOL_DAYS)) &\\
                        (renamed_df.r_DATE < no_pool_df.DATE)\\
                    , 'right')\\
                .select('ID', 'r_CHANGE')\\
                .groupBy('ID')\\
                .agg(collect_list('r_CHANGE'))\\
                .withColumnRenamed('ID', 'no_pool__ID')\\
                .withColumnRenamed('collect_list(r_CHANGE)', 'no_pool__POOL')

    df = df.join(cross_df, cross_df.no_pool__ID == df.ID, 'left_outer')\\
           .drop('no_pool__ID')

    现在,DataFrame df具有两列,一起包含所需的信息,即原始的POOL和新创建的no_pool__POOL。我使用以下UDF将这两部分正确地混合为一列POOL_mix

    1
    2
    3
    4
    5
    6
    7
    8
    9
    def mix_pools(history, no_history):
        if not history:
            return no_history
        else:
            return history
    mix_pools_udf = sfn.udf(mix_pools, ArrayType(DoubleType()))

    df = df.withColumn('POOL_mix', mix_pools_udf(col('POOL'), col('no_pool__POOL')))\\
           .drop('POOL', 'no_pool__POOL')