比较pyspark中的3列

gudnpqoy  于 2021-05-29  发布在  Spark
关注(0)|答案(1)|浏览(383)

我想比较pyspark中的3列(百分比总和为100%),以创建一个新的列,其中包含3列中最大值的列名,或者,如果最大值不唯一,则包含具有相同值的列的名称。我在这里看到过一些类似的例子,但是它们不能处理max不是唯一的情况。下面是我的暴力解决方案,但它需要这么多时间才能变得毫无用处:

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

有什么见解可以帮助您构建更高效的解决方案吗?

mcvgt66p

mcvgt66p1#

如果我理解正确,你可以和 greatest 返回列名,然后concat:example:
输入:

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

相关问题