Spark：使用地图逻辑添加列而不使用 UDF

Question

基本上，我想在数据框的每一行上应用我的函数 countSimilarColumns 并将结果放在一个新列中。

我的代码如下

def main(args: Array[String]) = {
    val customerID           = "customer-1" //args(0)
    val rawData              = readFromResource("json", "/spark-test-data-copy.json")
    val flattenData          = rawData.select(flattenSchema(rawData.schema): _*)
    val referenceCustomerRow = flattenData.transform(getCustomer(customerID)).first
  }

def getCustomer(customerID: String)(dataFrame: DataFrame) = {
    dataFrame.filter($"customer" === customerID)
  }

def countSimilarColumns(first: Row, second: Row): Int = {
    if (!(first.getAs[String]("customer").equals(second.getAs[String]("customer"))))
      first.toSeq.zip(second.toSeq).count { case (x, y) => x == y }
    else
      -1
  }

我想做如下的事情。但我不知道该怎么做。

flattenData
  .withColumn(
    "similarity_score",
    flattenData.map(row => countSimilarColumns(row, referenceCustomerRow))
  )
  .show()

扁平化示例数据：

{"customer":"customer-1","att-a":"7","att-b":"3","att-c":"10","att-d":"10"}
{"customer":"customer-2","att-a":"9","att-b":"7","att-c":"12","att-d":"4"}
{"customer":"customer-3","att-a":"7","att-b":"3","att-c":"1","att-d":"10"}
{"customer":"customer-4","att-a":"9","att-b":"14","att-c":"10","att-d":"4"}

想要的输出：

+--------------------+-----------+
| customer   | similarity_score |
+--------------------+-----------+
|customer-1  |  -1    | 
|customer-2  |  0    |
|customer-3  |  3    |
|customer-4  |  1    |

UDF 是唯一的方法吗？如果是，那么我想保持我的函数 countSimilarColumns 原样，这样它就可以测试了。怎么可能？ 我是 Spark/Scala 的新手。

Answer 1

您需要充分了解 DF 上的所有标准非聚合操作一次只能看到一行。因此，您必须考虑 1 行或聚合。在单行上下文中对整个 DF 调用函数 - 您的 withColumn 调用非常不寻常。当您想将行与其他行进行比较/组合时，您通常将表连接到另一个表或自身，然后使用组合表。

你想要的是类似的东西

// create DF with similarity scores AND customer, which is treated as row id here
val similarityDF = flattenData.map(row => row.getAs[String]("customer") -> countSimilarColumns(row, referenceCustomerRow)).
                   toDF("customer","similarity_score")

// join your original DF with similarityDF based on "customer"
flattenData.join(simlarityDF, usingColumn = "customer").show()

这允许您匹配两个 DF 的行。

UDF 在这里并不真正适用，因为与内置函数类似，UDF 应用于一个（或多个）特定列以生成一个值（可以使用withColumn 作为列添加），您想要的查看整个 Row。

Answer 2

flattenData 是 DataFrame 类型，在 flattenData 上应用 map 函数将得到 Dataset 的结果。

您将flattenData.map(row => countSimilarColumns(row, referenceCustomerRow)) 的结果传递给withColumn 但withColumn 只能获取org.apache.spark.sql.Column 类型的数据

因此，如果您想在没有UDF 的情况下将上述结果添加到列中，您必须使用collect 函数，然后将其传递给lit

请检查以下代码。

flattenData
.withColumn("similarity_score",lit(
        flattenData
        .map(row => countSimilarColumns(row, referenceCustomerRow))
        .collect
        .map(_.toInt)
   )
) 

根据样本数据，添加以下逻辑。

scala> df.show(false)
+-----+-----+-----+-----+----------+
|att-a|att-b|att-c|att-d|customer  |
+-----+-----+-----+-----+----------+
|7    |3    |10   |10   |customer-1|
|9    |7    |12   |4    |customer-2|
|7    |3    |1    |10   |customer-3|
|9    |14   |10   |4    |customer-4|
+-----+-----+-----+-----+----------+


scala> val conditions = df.columns.filterNot(_ == "customer").map(c => (when(count(col(c)).over(Window.partitionBy(col(c)).orderBy(col(c).asc)) =!= 1,lit(1)).otherwise(0))).reduce(_ + _) // if row_number is 1 then adding 1 else 0 ..
conditions: org.apache.spark.sql.Column = (((CASE WHEN (NOT (count(att-a) OVER (PARTITION BY att-a ORDER BY att-a ASC NULLS FIRST unspecifiedframe$()) = 1)) THEN 1 ELSE 0 END + CASE WHEN (NOT (count(att-b) OVER (PARTITION BY att-b ORDER BY att-b ASC NULLS FIRST unspecifiedframe$()) = 1)) THEN 1 ELSE 0 END) + CASE WHEN (NOT (count(att-c) OVER (PARTITION BY att-c ORDER BY att-c ASC NULLS FIRST unspecifiedframe$()) = 1)) THEN 1 ELSE 0 END) + CASE WHEN (NOT (count(att-d) OVER (PARTITION BY att-d ORDER BY att-d ASC NULLS FIRST unspecifiedframe$()) = 1)) THEN 1 ELSE 0 END)

最终结果

scala> df.withColumn("similarity_score",conditions).show(false)
+-----+-----+-----+-----+----------+----------------+
|att-a|att-b|att-c|att-d|customer  |similarity_score|
+-----+-----+-----+-----+----------+----------------+
|9    |7    |12   |4    |customer-2|2               |
|7    |3    |1    |10   |customer-3|3               |
|7    |3    |10   |10   |customer-1|4               |
|9    |14   |10   |4    |customer-4|3               |
+-----+-----+-----+-----+----------+----------------+