在 PySpark 中添加基于行的操作列答案

【问题标题】：Adding a column on row based operations in PySpark在 PySpark 中添加基于行的操作列
【发布时间】：2023-03-23 14:08:02
【问题描述】：

我有一个 Spark DataFrame，我正在尝试根据以前的列创建一个新列，但对我来说困难的部分是我已经逐行计算了列的值。例如：

col1 |col2 |col3

1 | 2 | 3

4 | 5 | 0

3 | 1 | 1

所以，我想要一个新列，其中包含表达式列的名称
每行最大（col1，col2，col3）。所以，想要的输出：

col1 |col2 |col3 |col4

1 | 2 | 3 | 'col3'

4 | 5 | 0 | 'col2'

3 | 1 | 1 | 'col1'

无论如何都可以在 PySpark 中做？

【问题讨论】：

对于您发布的数据，您希望的输出是什么样的？
我已更新问题以反映所需的输出。
如果出现平局怎么办？如果两个数字相同怎么办？
原始问题陈述处理了这种情况。每行没有两列具有相同的数据。

标签： apache-spark pyspark

【解决方案1】：

这不是一个理想的答案，因为它迫使您回到 RDD。如果我找到一个更好的让你留在DataFrame 宇宙中的，我会更新我的答案。但这现在应该可行。

a = sc.parallelize([[1,2,3],[4,5,0],[3,1,1]])
headers = ["col1", "col2", "col3"]

b = a.map(lambda x: (x[0], x[1], x[2], headers[x.index(max(x))]))

b.toDF(headers.append("max_col")).show()

这基本上允许您通过逐行遍历您的 RDD 来使用 python 中的max 操作。然后它通过索引标题列表找到正确的列。

再说一次，我不确定这是不是最好的方法，我希望找到更好的方法。

【讨论】：