PySpark 列表中元素的差異

Question

我有一個 PySpark dataframe （ df ），其中有一列包含兩個元素的列表。 列表中的兩個元素不按升序或降序排序。

+--------+----------+-------+
| version| timestamp| list  |
+--------+-----+----|-------+
| v1     |2012-01-10| [5,2] |
| v1     |2012-01-11| [2,5] |
| v1     |2012-01-12| [3,2] |
| v2     |2012-01-12| [2,3] |
| v2     |2012-01-11| [1,2] |
| v2     |2012-01-13| [2,1] |
+--------+----------+-------+

我想在列表的第一個和第二個元素之間獲取差異，並將其作為另一列（ diff ）。 這是我想要的 output 的示例。

+--------+----------+-------+-------+
| version| timestamp| list  |  diff | 
+--------+-----+----|-------+-------+
| v1     |2012-01-10| [5,2] |   3   |
| v1     |2012-01-11| [2,5] |  -3   |
| v1     |2012-01-12| [3,2] |   1   |
| v2     |2012-01-12| [2,3] |  -1   |
| v2     |2012-01-11| [1,2] |  -1   |
| v2     |2012-01-13| [2,1] |   1   |
+--------+----------+-------+-------+

如何使用 PySpark 做到這一點？

我嘗試了以下方法：

transform_expr = (
        "transform(diff, x-y ->"
        + "x as list[0], y as list[1])"
    )

df = df.withColumn("diff", F.expr(transform_expr)) 

但是，上述技術沒有給我任何 output。

如果有人需要，我也願意使用 UDF 來獲得我想要的 output。

沒有 UDF 的方法和基於 UDF 的方法都受到歡迎。 謝謝。

Answer 1

有多種方法可以做到這一點，您可以使用element_at （Spark 2.4 或更高版本）、 transform 、 array index[0]或.getItem()中的任何一種來獲得差異。

#sample dataframe
df=spark.createDataFrame([([5,2],),([2,5],)],["list"])

#using element_at
df.withColumn("diff",element_at(col("list"),1) - element_at(col("list"),2)).show()

#using transform 
df.withColumn("diff",concat_ws("",expr("""transform(array(list),x -> x[0] - x[1])"""))).show()

#using array index
df.withColumn("diff",col("list")[0]- col("list")[1]).show()

#using .getItem
df.withColumn("diff",col("list").getItem(0)- col("list").getItem(1)).show()

#+------+----+
#|  list|diff|
#+------+----+
#|[5, 2]|   3|
#|[2, 5]|  -3|
#+------+----+