在 pyspark 数据框中找到第一个非零元素

Question

我正在使用 pyspark 数据帧并尝试查看是否有一种方法可以提取我在 spark 数据帧中的第一个非零元素的索引。我自己添加了索引列，因为 pyspark 不支持它，而不是 pandas。

Answer 1

假设您的数据框如下所示：

df.show()
+---+-----+                                                                     
|idx|value|
+---+-----+
|  0|    0|
|  1|    0|
|  2|    1|  # <-- We want this one
|  3|    2|
|  4|    3|
|  5|    4|
+---+-----+

您可以使用min 轻松实现此目的：

from pyspark.sql import functions as F

df.where(F.col("value") != 0).select(F.min("idx")).show()

或row_number

from pyspark.sql import functions as F, Window

df.where(F.col("value") != 0).withColumn(
    "rwnb", F.row_number().over(Window.orderBy("idx"))
).where(F.col("rwnb") == 1).select("idx").show()
+---+
|idx|
+---+
|  2|
+---+