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

um6iljoc  于 2021-09-29  发布在  Java
关注(0)|答案(2)|浏览(450)

我有一个如下所示格式的数据框,其中我们将有多个 DEPNAME 如下所示,我的要求是设置 result =y,在 DEPNAME 水平,如果 flag_1flag_2 =y,如果两个标志都是。 flag_1flag_2 =n结果将设置为n,如图所示 DEPNAME =人员
我可以使用连接获得所需的结果,但我很好奇是否可以使用窗口函数,因为数据集的大小相当大。

+---------+------+------+-+------+
|  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    |
+---------+-----+------+ +------+
l7wslrjt

l7wslrjt1#

这看起来像一个 case 表达方式:

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

13z8s7eq2#

如果您正在使用pyspark(因为您将其包含在标记中),并说您的 Dataframe 被调用 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 表示 'Y' 所有拥有 cnt >= 1 .

相关问题