Scala / Spark DataFrame：連接行以在其內部獲得一對

Question

我正在嘗試編寫一個將df-> df2轉換為如下的函數：

// input dataframe df
+-----+-----+
|  T  |  S  |
+-----+-----+
|    A|   4|
|    B|   8|
|    C|   8|
|    D|   2|
+-----+-----+

我需要一個將df作為輸入並返回df2作為輸出的函數。

// output dataframe df2
+-----+-----+-----+
| T1 | T2 | S=T1+T2 |
+-----+-----+-----+
|    A|    B|  12|
|    A|    C|  12|
|    A|    D|  6|
|    B|    C|  16|
|    B|    D|  10|
|    C|    D|  10|
+-----+-----+-----+

編輯我想出了這個解決方案。 任何改進都將受到歡迎。

val sumOf = udf((left_score: Float, right_score: Float) => left_score + right_score)

val left = df.select("T", "S").withColumnRenamed("T", "T1").withColumnRenamed("S", "S1")
val right= df.select("T", "S").withColumnRenamed("T", "T2").withColumnRenamed("S", "S2")

val joinDF = left.join(right, left.col("T1") !== right.col("T2"))
val outDF = joinDF.withColumn("S", sumOf($"S1", $"S2")).select("T1", "T2", "S")

Answer 1

val df = sc.parallelize(Seq("A" -> 4, "B" -> 8, "C" -> 8, "D" -> 2))
           .toDF("T", "S")

val df1 = df.withColumnRenamed("T", "T1")
            .withColumnRenamed("S", "S1")

val df2 = df.withColumnRenamed("T", "T2")
            .withColumnRenamed("S", "S2")

df1.join(df2, df1("T1") < df2("T2"))
   .withColumn("S", 'S1 + 'S2)
   .drop("S1", "S2")
   .show

+---+---+---+
| T1| T2|  S|
+---+---+---+
|  A|  B| 12|
|  A|  C| 12|
|  A|  D|  6|
|  B|  C| 16|
|  B|  D| 10|
|  C|  D| 10|
+---+---+---+

基本上，您不需要完整的笛卡爾積。 T2> T1時只有所有可能性。 這就是連接條件在代碼中的含義。 請注意，笛卡爾積會生成n²個記錄。 在這里，您將生成n（n-1）/ 2條記錄。 該值小於n²，但仍為O（n²），因此應盡可能避免使用...

Answer 2

撇開性能（提示：不可能使Spark在大型笛卡爾產品上表現良好），可以使用Spark 2.x中引入的交叉聯接。

import sc.implicits._

val df = sc.parallelize(Seq("A" -> 4, "B" -> 8, "C" -> 8, "D" -> 2))
         .toDF("T", "S")

df.as("df1")
  .crossJoin(df.as("df2"))
    .filter($"df1.T" =!= $"df2.T")
      .select($"df1.T".as("T1"), $"df2.T".as("T2"))
      .withColumn("S", $"df1.S"+$"df2.S") // you can use udf here as well

內部聯接可以實現相同的結果，從而使其與Spark 1.6.x兼容

import sc.implicits._

val df = sc.parallelize(Seq("A" -> 4, "B" -> 8, "C" -> 8, "D" -> 2))
         .toDF("T", "S")

df.as("df1")
  .join(df.as("df2"), Seq("T"), "inner") // this line is different
    .filter($"df1.T" =!= $"df2.T")
     .select($"df1.T".as("T1"), $"df2.T".as("T2"))
     .withColumn("S", $"df1.S"+$"df2.S") // you can use udf here as well

Answer 3

我建議的解決方案根本不需要您使用join 。 但是該解決方案也很昂貴，因為所有數據都將被累積到一個執行器中進行處理 。

我的解決方案是將內置函數（例如array ， collect_list和explode與window函數結合在一起，如下所示

import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions._

def windowFunction = Window.orderBy("T").rowsBetween(1, Long.MaxValue)

df.withColumn("array", collect_list(array($"T", $"S")).over(windowFunction))
    .withColumn("array", explode($"array"))
    .select($"T".as("T1"), $"array"(0).as("T2"), ($"array"(1)+$"S").as("S=T1+T2"))
  .show(false)

這應該給你你想要的輸出為

+---+---+-------+
|T1 |T2 |S=T1+T2|
+---+---+-------+
|A  |B  |12.0   |
|A  |C  |12.0   |
|A  |D  |6.0    |
|B  |C  |16.0   |
|B  |D  |10.0   |
|C  |D  |10.0   |
+---+---+-------+