如何在pyspark窗口分区执行自定义逻辑

Question

我有一个如下所示格式的数据框，其中我们将有多个DEPNAME 条目，如下所示，我的要求是在DEPNAME 级别设置result = Y，如果flag_1 或flag_2= Y，如果标志即flag_1 和flag_2 = N，则结果将设置为N，如DEPNAME=personnel 所示

我能够使用连接获得所需的结果，但我很好奇我们是否可以使用窗口函数来实现它，因为数据集的大小非常大。

+---------+------+------+-+------+
|  depName|flag_1|flag_2| result |
+---------+------+------+-+------+
|    sales|    N|  Y    |  Y    |
|    sales|    N|  N    |  Y    |
|    sales|    N|  N    |  Y    |
|personnel|    N|  N    |  N    |
|personnel|    N|  N    |  N    |
|  develop|    Y|  N    |  Y    |
|  develop|    N|  N    |  Y    |
|  develop|    N|  N    |  Y    |
|  develop|    N|  N    |  Y    |
|  develop|    N|  N    |  Y    |
+---------+-----+------+ +------+

Answer 1

这回答了问题的原始版本。

这看起来像 case 表达式：

select t.*,
       (case when flag_1 = 'Y' or flag_2 = 'Y'
             then 'Y' else 'N'
        end) as result

对于更新版本：

select t.*,
       max(case when flag_1 = 'Y' or flag_2 = 'Y'
                then 'Y' else 'N'
           end) over (partition by depname) as result

Answer 2

如果您正在使用 PySpark（因为您将其包含在标签中）并说您的数据框称为 df，您可以使用

import pyspark.sql.functions as F
from pyspark.sql.window import Window

w = Window.partitionBy('depName')

df = df\
  .withColumn('cnt', F.sum(F.when((F.col('flag_1') == 'Y') | (F.col('flag_2') == 'Y'), 1).otherwise(0)).over(w))\
  .withColumn('result', F.when(F.col('cnt') >= 1, 'Y').otherwise('N'))

df.show()

+---------+------+------+---+------+
|  depName|flag_1|flag_2|cnt|result|
+---------+------+------+---+------+
|  develop|     Y|     N|  1|     Y|
|  develop|     N|     N|  1|     Y|
|  develop|     N|     N|  1|     Y|
|  develop|     N|     N|  1|     Y|
|  develop|     N|     N|  1|     Y|
|personnel|     N|     N|  0|     N|
|personnel|     N|     N|  0|     N|
|    sales|     N|     Y|  1|     Y|
|    sales|     N|     N|  1|     Y|
|    sales|     N|     N|  1|     Y|
+---------+------+------+---+------+

基本上，在由depName 确定的每个分区中，您计算​​条件flag_1 == 'Y' | flag_2 == 'Y' 出现的次数，并将其存储在cnt 中以用于该分区的所有行。
然后，您使用一个简单的.when 指示所有具有cnt >= 1 的组。