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

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

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

  1. +---------+------+------+-+------+
  2. | depName|flag_1|flag_2| result |
  3. +---------+------+------+-+------+
  4. | sales| N| Y | Y |
  5. | sales| N| N | Y |
  6. | sales| N| N | Y |
  7. |personnel| N| N | N |
  8. |personnel| N| N | N |
  9. | develop| Y| N | Y |
  10. | develop| N| N | Y |
  11. | develop| N| N | Y |
  12. | develop| N| N | Y |
  13. | develop| N| N | Y |
  14. +---------+-----+------+ +------+
l7wslrjt

l7wslrjt1#

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

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

13z8s7eq2#

如果您正在使用pyspark(因为您将其包含在标记中),并说您的 Dataframe 被调用 df ,你可以使用

  1. import pyspark.sql.functions as F
  2. from pyspark.sql.window import Window
  3. w = Window.partitionBy('depName')
  4. df = df\
  5. .withColumn('cnt', F.sum(F.when((F.col('flag_1') == 'Y') | (F.col('flag_2') == 'Y'), 1).otherwise(0)).over(w))\
  6. .withColumn('result', F.when(F.col('cnt') >= 1, 'Y').otherwise('N'))
  7. df.show()
  8. +---------+------+------+---+------+
  9. | depName|flag_1|flag_2|cnt|result|
  10. +---------+------+------+---+------+
  11. | develop| Y| N| 1| Y|
  12. | develop| N| N| 1| Y|
  13. | develop| N| N| 1| Y|
  14. | develop| N| N| 1| Y|
  15. | develop| N| N| 1| Y|
  16. |personnel| N| N| 0| N|
  17. |personnel| N| N| 0| N|
  18. | sales| N| Y| 1| Y|
  19. | sales| N| N| 1| Y|
  20. | sales| N| N| 1| Y|
  21. +---------+------+------+---+------+

基本上,在由 depName ,您可以计算条件的次数 flag_1 == 'Y' | flag_2 == 'Y' 发生,并将其存储在 cnt 对于该分区的所有行。
然后,使用一个简单的 .when 表示 'Y' 所有拥有 cnt >= 1 .

展开查看全部

相关问题