如何在pyspark中使用map变换替换python中的循环,我们希望将前一行和当前行与多个条件进行比较

时间:2018-03-09 13:22:29

标签: pyspark rdd

在pyspark数据框中应用地图功能时,只是拖了一个路障类型的情况,需要你的帮助才能解决这个问题。

虽然问题更复杂,但让我用下面的例子使用字典和for循环来简化它,并且需要在pyspark中解决。

这里有关于虚拟数据的python代码的例子,我想在pyspark map转换中使用when,子句使用窗口或任何其他方式。

问题 - 我有一个pyspark数据框,列名作为下面字典中的键,并希望在此示例中添加/修改具有类似逻辑的for section列。

record=[
{'id':xyz,'SN':xyz,'miles':xyz,'feet':xyz,'MP':xyz,'section':xyz},
{'id':xyz,'SN':xyz,'miles':xyz,'feet':xyz,'MP':xyz,'section':xyz},
{'id':xyz,'SN':xyz,'miles':xyz,'feet':xyz,'MP':xyz,'section':xyz}
]

last_rec='null'
section=0
for cur_rec in record:
    if lastTrack != null:
        if (last_rec.id != cur_rec.id | last_rec.SN != cur_rec.SN):
            section+=1

        elif (last_rec.miles == cur_rec.miles & abs(last_rec.feet- cur_rec.feet) > 1):
            section+=1

        elif (last_rec.MP== 555 & cur_rec.MP != 555):
            section+=1

        elif (abs(last_rec.miles- cur_rec.miles) > 1):
            section+=1


    cur_rec['section']= section
    last_rec = cur_rec

1 个答案:

答案 0 :(得分:1)

您的窗口函数是布尔变量的累积和。 让我们从一个示例数据帧开始:

import numpy as np
record_df = spark.createDataFrame(
    [list(x) for x in zip(*[np.random.randint(0, 10, 100).tolist() for _ in range(5)])], 
    ['id', 'SN', 'miles', 'feet', 'MP'])
record_df.show()

    +---+---+-----+----+---+
    | id| SN|miles|feet| MP|
    +---+---+-----+----+---+
    |  9|  5|    7|   5|  1|
    |  0|  6|    3|   7|  5|
    |  8|  2|    7|   3|  5|
    |  0|  2|    6|   5|  8|
    |  0|  8|    9|   1|  5|
    |  8|  5|    1|   6|  0|
    |  0|  3|    9|   0|  3|
    |  6|  4|    9|   0|  8|
    |  5|  8|    8|   1|  0|
    |  3|  0|    9|   9|  9|
    |  1|  1|    2|   7|  0|
    |  1|  3|    7|   7|  6|
    |  4|  9|    5|   5|  5|
    |  3|  6|    0|   0|  0|
    |  5|  5|    5|   9|  3|
    |  8|  3|    7|   8|  1|
    |  7|  1|    3|   1|  8|
    |  3|  1|    5|   2|  5|
    |  6|  2|    3|   5|  6|
    |  9|  4|    5|   9|  1|
    +---+---+-----+----+---+

累积和是一个有序窗口函数,因此我们需要使用monotonically_increasing_id给我们的行命令:

import pyspark.sql.functions as psf
record_df = record_df.withColumn(
    'rn', 
    psf.monotonically_increasing_id())

对于布尔变量,我们需要使用lag

from pyspark.sql import Window
w = Window.orderBy('rn')
record_df = record_df.select(
    record_df.columns 
    + [psf.lag(c).over(w).alias('prev_' + c) for c in ['id', 'SN', 'miles', 'feet', 'MP']])

由于所有条件在section上产生相同的结果,因此它是or子句:

clause = (psf.col("prev_id") != psf.col("id")) | (psf.col("prev_SN") != psf.col("SN")) \
    | ((psf.col("prev_miles") == psf.col("miles")) & (psf.abs(psf.col("prev_feet") - psf.col("feet")) > 1)) \
    | ((psf.col("prev_MP") == 555) & (psf.col("MP") != 555)) \
    | (psf.abs(psf.col("prev_miles") - psf.col("miles")) > 1)
record_df = record_df.withColumn("tmp", (clause).cast('int'))

最后是累积金额

record_df = record_df.withColumn("section", psf.sum("tmp").over(w))