比较 PySpark 中的 3 列

Question

我想比较 PySpark 中的 3 列（百分比总和为 100%）以创建一个新列，其中包含 3 列最大值的列名，或者，如果最大值不唯一，则包含列名具有相同的价值。我在这里看到了一些类似的例子，但是当最大值不唯一时，它们不会处理这种情况。下面是我的蛮力解决方案，但运行需要很长时间才能变得无用：

df\
  .withColumn("MaxName", 
      F.when( (col(A)>col(B)) & (col(A)>col(C)), "A")\
      .when( (col(B)>col(A)) & (col(B)>col(C)), "B")\
      .when( (col(C)>col(A)) & (col(C)>col(B)), "C")\
      .when( (col(A)==col(B)) &\
            (col(A)>col(C)) | (col(B)>col(C)), "AB")\
      .when( (col(C)==col(B)) | (col(C)==col(A)) &\
            (col(C)>col(B)) | (col(C)>col(A)), "CAB")\
      .otherwise("ABC")

任何见解来构建更有效的解决方案？

Answer 1

如果我理解正确，您可以比较greatest并返回列名，然后连接：示例：

输入：

np.random.seed(111)
df = spark.createDataFrame(pd.DataFrame(np.random.randint(0,100,(5,5)),
                                        columns=list('ABCDE')))
df.show()

+---+---+---+---+---+
|  A|  B|  C|  D|  E|
+---+---+---+---+---+
| 84| 84| 84| 86| 19|
| 41| 66| 82| 40| 71|
| 57|  7| 12| 10| 65|
| 88| 28| 14| 34| 21|
| 54| 72| 37| 76| 58|
+---+---+---+---+---+

建议的解决方案：

import pyspark.sql.functions as F

cols = ['A','B','C']
df.withColumn("max_of_ABC",F.concat_ws("",
             *[F.when(F.col(i) == F.greatest(*cols),i) for i in cols])).show()

+---+---+---+---+---+----------+
|  A|  B|  C|  D|  E|max_of_ABC|
+---+---+---+---+---+----------+
| 84| 84| 84| 86| 19|       ABC|
| 41| 66| 82| 40| 71|         C|
| 57|  7| 12| 10| 65|         A|
| 88| 28| 14| 34| 21|         A|
| 54| 72| 37| 76| 58|         B|
+---+---+---+---+---+----------+