Как создать выражение столбца с помощью подзапроса в Spark scala

Question

Учитывая любой df, я хочу вычислить другой столбец для df с именем has_duplicates, а затем добавить столбец с логическим значением, определяющим, является ли каждая строка уникальной. Пример ввода df:

val df = Seq((1, 2), (2, 5), (1, 7), (1, 2), (2, 5)).toDF("A", "B")

Учитывая ввод columns: Seq[String], я знаю, как получить количество каждой строки:

val countsDf = df.withColumn("count", count("*").over(Window.partitionBy(columns.map(col(_)): _*)))

Но я не уверен, как это использовать для создания выражения столбца для последнего столбца, показывающего, является ли каждая строка уникальной.

Что-то вроде

def getEvaluationExpression(df: DataFrame): Column = {
    when("count > 1", lit("fail").otherwise(lit("pass"))
 }

, но счетчик должен быть вычислен на месте, используя запрос выше.

Answer 1

Попробуйте следующий код.

scala> df.withColumn("has_duplicates", when(count("*").over(Window.partitionBy(df.columns.map(col(_)): _*)) > 1 , lit("fail")).otherwise("pass")).show(false)
+---+---+--------------+
|A  |B  |has_duplicates|
+---+---+--------------+
|1  |7  |pass          |
|1  |2  |fail          |
|1  |2  |fail          |
|2  |5  |fail          |
|2  |5  |fail          |
+---+---+--------------+

Или

scala> df.withColumn("count",count("*").over(Window.partitionBy(df.columns.map(col(_)): _*))).withColumn("has_duplicates", when($"count" > 1 , lit("fail")).otherwise("pass")).show(false)
+---+---+-----+--------------+
|A  |B  |count|has_duplicates|
+---+---+-----+--------------+
|1  |7  |1    |pass          |
|1  |2  |2    |fail          |
|1  |2  |2    |fail          |
|2  |5  |2    |fail          |
|2  |5  |2    |fail          |
+---+---+-----+--------------+