比较 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|

+---+---+---+---+---+----------+