仅使用一个Window函数生成组ID

时间:2018-07-25 13:54:41

标签: python sql apache-spark pyspark

我正在尝试为表分配一个组ID。
对于“ part”列中的每个组,按照“ ord”列中的顺序,我在“ id”列中遇到的第一个元素将收到一个new_id 0,然后每次遇到一个不同的ID时,都会增加“ new_id”。

当前,我需要使用2个窗口函数,因此该过程非常缓慢。

df = sqlContext.createDataFrame(
  [
    (1,1,'X'),
    (1,2,'X'),
    (1,3,'X'),
    (1,4,'Y'),
    (1,5,'Y'),
    (1,6,'Y'),
    (1,7,'X'),
    (1,8,'X'),
    (2,1,'X'),
    (2,2,'X'),
    (2,3,'X'),
    (2,4,'Y'),
    (2,5,'Y'),
    (2,6,'Y'),
    (2,7,'X'),
    (2,8,'X'),
  ],
  ["part", "ord", "id"]
)

df.withColumn(
  "new_id", 
  F.lag(F.col("id")).over(Window.partitionBy("part").orderBy("ord"))!=F.col('id')
).withColumn(
  "new_id", 
  F.sum(
    F.col("new_id").cast('int')
  ).over(Window.partitionBy("part").orderBy("ord"))
).na.fill(0).show()

+----+---+---+------+
|part|ord| id|new_id|
+----+---+---+------+
|   2|  1|  X|     0|
|   2|  2|  X|     0|
|   2|  3|  X|     0|
|   2|  4|  Y|     1|
|   2|  5|  Y|     1|
|   2|  6|  Y|     1|
|   2|  7|  X|     2|
|   2|  8|  X|     2|
|   1|  1|  X|     0|
|   1|  2|  X|     0|
|   1|  3|  X|     0|
|   1|  4|  Y|     1|
|   1|  5|  Y|     1|
|   1|  6|  Y|     1|
|   1|  7|  X|     2|
|   1|  8|  X|     2|
+----+---+---+------+

我可以仅使用一个窗口函数来实现相同的功能吗?

0 个答案:

没有答案