Разница элементов в списке в PySpark

Question

У меня есть фрейм данных PySpark (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) , Вот пример вывода, который я хочу.

+--------+----------+-------+-------+
| 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)) 

Но вышеприведенный метод не дал мне никакого вывода.

Я также открыт для использования UDF для получения желаемого результата в случае необходимости.

Приветствуются подходы без 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|
#+------+----+