python - 如何在没有UDF的情况下如何计算PySpark数据帧中数组列中的尾随零

Question

我有一个带有固定数量的整数的数组的列的数据框。
如何在df中添加包含数组中尾随零的数量的列？
我想避免使用UDF以获得更好的性能。

例如，输入df:

>>> df.show()
+------------+
|           A|
+------------+
| [1,0,1,0,0]|
| [2,3,4,5,6]|
| [0,0,0,0,0]|
| [1,2,3,4,0]|
+------------+

和通缉的输出:

>>> trailing_zeroes(df).show()
+------------+-----------------+
|           A|   trailingZeroes|
+------------+-----------------+
| [1,0,1,0,0]|                2|
| [2,3,4,5,6]|                0|
| [0,0,0,0,0]|                5|
| [1,2,3,4,0]|                1|
+------------+-----------------+

Answer 1

从Spark 2.4开始，您可以使用高阶函数AGGREGATE来做到这一点:

from pyspark.sql.functions import reverse

(
  df.withColumn("arr_rev", reverse("A"))
  .selectExpr(
    "arr_rev", 
    "AGGREGATE(arr_rev, (1 AS p, CAST(0 AS LONG) AS sum), (buffer, value) -> (if(value != 0, 0, buffer.p), if(value=0, buffer.sum + buffer.p, buffer.sum)), buffer -> buffer.sum) AS result"
  )
)

假设 A是带有数字的数组。这里只是注意数据类型。假设数组中的数字也很长，我正在将初始值转换为 LONG。