Pyspark：将一列与其他列进行比较，如果相似则进行标记

Question

我想创建一个列，如果在 column2、column3 或 column4 中找到我的 column1 值，它会进行标记。

我可以这样做：

import pyspark.sql.functions as f

df.withColumn("FLAG", f.when((f.col("column1") == f.col("column2")) |
                             (f.col("column1") == f.col("column3")) | 
                             (f.col("column1") == f.col("column4")), 'Y')\
                       .otherwise('N'))

这需要相当长的时间，而且我觉得效率很低。我想知道是否有更好的方法使用 UDF 来编写此代码？我正在尝试弄清楚如何引用匹配的列。

任何帮助都有帮助！谢谢

Answer 1

您可以使用 Spark 内置函数 isin

from pyspark.sql import functions as F

df = (spark
    .sparkContext
    .parallelize([
        ('A', 'A', 'B', 'C'),
        ('A', 'B', 'C', 'A'),
        ('A', 'C', 'A', 'B'),
        ('A', 'X', 'Y', 'Z'),
    ])
    .toDF(['ca', 'cb', 'cc', 'cd'])
)

(df
  .withColumn('flag', F.col('ca').isin(
    F.col('cb'),
    F.col('cc'),
    F.col('cd'),
  ))
  .show()
)
# +---+---+---+---+-----+
# | ca| cb| cc| cd| flag|
# +---+---+---+---+-----+
# |  A|  A|  B|  C| true|
# |  A|  B|  C|  A| true|
# |  A|  C|  A|  B| true|
# |  A|  X|  Y|  Z|false|
# +---+---+---+---+-----+

Pyspark：将一列与其他列进行比较，如果相似则进行标记

Pyspark: compare one column with other columns and flag if similar

python

lambda

pyspark

user-defined-functions