比较 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 Fcols = ['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|+---+---+---+---+---+----------+