Question

我有PySpark代码，可以有效地对行进行数字分组，并在满足特定条件时递增。我在弄清楚如何有效地将此代码转换为可应用于组的代码时遇到了麻烦。

获取此示例数据帧df

df = sqlContext.createDataFrame(
    [
        (33, [], '2017-01-01'),
        (33, ['apple', 'orange'], '2017-01-02'),
        (33, [], '2017-01-03'),
        (33, ['banana'], '2017-01-04')
    ],
    ('ID', 'X', 'date')
)

此代码实现了我想要的此示例df的功能，该功能按日期排序并创建当size列回到0时递增的组（'grp'）。

df \
.withColumn('size', size(col('X'))) \
.withColumn(
    "grp", 
    sum((col('size') == 0).cast("int")).over(Window.orderBy('date'))
).show()

这部分基于Pyspark - Cumulative sum with reset condition

现在我要尝试的是将相同的方法应用于具有多个ID的数据框-实现看起来像

的结果

df2 = sqlContext.createDataFrame(
    [
        (33, [], '2017-01-01', 0, 1),
        (33, ['apple', 'orange'], '2017-01-02', 2, 1),
        (33, [], '2017-01-03', 0, 2),
        (33, ['banana'], '2017-01-04', 1, 2),
        (55, ['coffee'], '2017-01-01', 1, 1),
        (55, [], '2017-01-03', 0, 2)
    ],
    ('ID', 'X', 'date', 'size', 'group')
)

为清晰起见

1）对于每个ID的第一个日期，该组应为1-无论其他列中显示什么。

2）但是，对于以后的每个日期，我都需要检查大小列。如果size列为0，则我增加组号。如果它是任何非零的正整数，那么我将继续使用先前的组号。

我已经看到了一些在Pandas中处理此问题的方法，但是我很难理解pyspark中的应用程序以及Pandas vs Spark中分组数据不同的方式（例如，我是否需要使用称为UADF的东西）？）

Answer 1

通过检查zero_or_first是零还是该行是第一行来创建列size。然后sum。

df2 = sqlContext.createDataFrame(
    [
        (33, [], '2017-01-01', 0, 1),
        (33, ['apple', 'orange'], '2017-01-02', 2, 1),
        (33, [], '2017-01-03', 0, 2),
        (33, ['banana'], '2017-01-04', 1, 2),
        (55, ['coffee'], '2017-01-01', 1, 1),
        (55, [], '2017-01-03', 0, 2),
        (55, ['banana'], '2017-01-01', 1, 1)
    ],
    ('ID', 'X', 'date', 'size', 'group')
)


w = Window.partitionBy('ID').orderBy('date')
df2 = df2.withColumn('row', F.row_number().over(w))
df2 = df2.withColumn('zero_or_first', F.when((F.col('size')==0)|(F.col('row')==1), 1).otherwise(0))
df2 = df2.withColumn('grp', F.sum('zero_or_first').over(w))
df2.orderBy('ID').show()

这里是输出。您可以看到group列== grp。其中group是预期的结果。

+---+---------------+----------+----+-----+---+-------------+---+
| ID|              X|      date|size|group|row|zero_or_first|grp|
+---+---------------+----------+----+-----+---+-------------+---+
| 33|             []|2017-01-01|   0|    1|  1|            1|  1|
| 33|       [banana]|2017-01-04|   1|    2|  4|            0|  2|
| 33|[apple, orange]|2017-01-02|   2|    1|  2|            0|  1|
| 33|             []|2017-01-03|   0|    2|  3|            1|  2|
| 55|       [coffee]|2017-01-01|   1|    1|  1|            1|  1|
| 55|       [banana]|2017-01-01|   1|    1|  2|            0|  1|
| 55|             []|2017-01-03|   0|    2|  3|            1|  2|
+---+---------------+----------+----+-----+---+-------------+---+

Answer 2

我添加了一个窗口函数，并在每个ID中创建了一个索引。然后，我将条件语句扩展为也引用该索引。以下内容似乎产生了我想要的输出数据帧-但我想知道是否有一种更有效的方法来完成此操作。

window = Window.partitionBy('ID').orderBy('date')
df \
.withColumn('size', size(col('X'))) \
.withColumn('index', rank().over(window).alias('index')) \
.withColumn(
    "grp", 
    sum(((col('size') == 0) | (col('index') == 1)).cast("int")).over(window)
).show()

产生

+---+---------------+----------+----+-----+---+
| ID|              X|      date|size|index|grp|
+---+---------------+----------+----+-----+---+
| 33|             []|2017-01-01|   0|    1|  1|
| 33|[apple, orange]|2017-01-02|   2|    2|  1|
| 33|             []|2017-01-03|   0|    3|  2|
| 33|       [banana]|2017-01-04|   1|    4|  2|
| 55|       [coffee]|2017-01-01|   1|    1|  1|
| 55|             []|2017-01-03|   0|    2|  2|
+---+---------------+----------+----+-----+---+

如何将条件计数（重置）应用于PySpark中的分组数据？

2 个答案: