如何使用 pyspark 計算連續值？

Question

我正在嘗試計算出現在 Pyspark 列中的連續值。 我的 dataframe 中有“a”列，並希望創建“b”列。

+---+---+
|  a|  b|
+---+---+
|  0|  1|
|  0|  2|
|  0|  3|
|  0|  4|
|  0|  5|
|  1|  1|
|  1|  2|
|  1|  3|
|  1|  4|
|  1|  5|
|  1|  6|
|  2|  1|
|  2|  2|
|  2|  3|
|  2|  4|
|  2|  5|
|  2|  6|
|  3|  1|
|  3|  2|
|  3|  3|
+---+---+

我試圖在一些 window 上創建滯后 function 的列“b”，但沒有成功。

w = Window\
  .partitionBy(df.some_id)\
  .orderBy(df.timestamp_column)

df.withColumn(
  "b",
  f.when(df.a == f.lag(df.a).over(w),
         f.sum(f.lit(1)).over(w)).otherwise(f.lit(0))
)

Answer 1

我可以使用以下代碼解決此問題：

df.withColumn("b",
  f.row_number().over(Window.partitionBy("a").orderBy("timestamp_column"))