如何在Apache Spark中基於時間的兩個數據集的聯合？

Question

給定S和R的兩個數據集都具有時間列（t），如下所述：

//snapshot with id at t
case class S(id: String, t: Int)

//reference data at t
case class R(t: Int, fk: String)

//Example test case
val ss: Dataset[S] = Seq(S("a", 1), S("a", 3), S("b", 5), S("b", 7))
      .toDS

    val rs: Dataset[R] = Seq(R(0, "a"), R(2, "a"), R(6, "b"))
      .toDS

    val srs: Dataset[(S, Option[R])] = ss
      .asOfJoin(rs)

    srs.collect() must contain theSameElementsAs
      Seq((S("a", 1), Some(R(0, "a"))), (S("a", 3), Some(R(2, "a"))), (S("b", 5), None), (S("b", 7), Some(R(6, "b"))))

目標是在R中找到與E的ID匹配的最新行，即R在輸出中可以是可選的。

asOfJoin定義如下：

  implicit class SOps(ss: Dataset[S]) {
    def asOfJoin(rs: Dataset[R])(implicit spark: SparkSession): Dataset[(S, Option[R])] = ???
  }

使用數據集API的一種解決方案如下：

def asOfJoin(rs: Dataset[R])(implicit spark: SparkSession): Dataset[(S, Option[R])] = {
      import spark.implicits._

      ss
        .joinWith(
          rs,
          ss("id") === rs("fk") && ss("t") >= rs("t"),
          "left_outer")
       .map { case (l, r) => (l, Option(r)) }
       .groupByKey { case (s, _) => s }
       .reduceGroups { (x, y) =>
         (x, y) match {
           case ((_, Some(R(tx, _))), (_, Some(R(ty, _)))) => if (tx > ty) x else y
           case _ => x
         }
       }
       .map { case (_, r) => r }
}

Answer 1

我不確定數據集S和數據集R的大小。但是從您的代碼中，我可以看到join（使用不相等的表達式）的效率很差，並且我可以根據不同的特定情況給出一些建議：

數據集R或數據集S都沒有太多數據。

我建議您可以廣播較小的數據集，並借助廣播變量在spark udf中完成業務邏輯。 這樣，您不需要shuffle（join）流程，這可以幫助您節省大量時間和資源。

對於每個唯一ID，count（distinct t）都不大。

我建議您可以通過對id和collect_set（t）進行分組來進行預聚合，如下所示：

select id,collect_set(t) as t_set from S 

這樣，您可以刪除聯接中的不等式（ss（“ t”）> = rs（“ t”））。 並使用數據集S和數據集R中的兩個t_set編寫業務邏輯。

對於其他情況：

我建議您使用相等的聯接和窗口函數來優化代碼。 由於我對SQL更加熟悉，因此我在此處編寫了SQL，可以將其轉換為數據集API：

select
  sid,
  st,
  rt
from
(
    select 
      S.id as sid,
      S.t as st,
      R.t as rt,
      row_number() over (partition by S.id order by (S.t - NVL(R.t, 0)) rn
    from
      S
    left join R on S.id = R.fk) tbl
where tbl.rn = 1

Answer 2

我接受了@bupt_ljy的有關避免theta聯接的評論，並且跟隨似乎可以很好地擴展：

def asOfJoin(rs: Dataset[R])(implicit spark: SparkSession): Dataset[(S, Option[R])] = {
  import spark.implicits._

  ss
    .joinWith(
      rs.sort(rs("fk"), rs("t")),
      ss("id") === rs("fk"),
      "left_outer")
    .map { case (l, r) => (l, Option(r)) }
    .groupByKey { case (s, _) => s }
    .flatMapGroups { (k, vs) =>
      new Iterator[(S, Option[R])] {
        private var didNotStart: Boolean = true

        override def hasNext: Boolean = didNotStart

        override def next(): (S, Option[R]) = {
          didNotStart = false
          vs
            .find { case (l, rOpt) =>
              rOpt match {
                case Some(r) => l.t >= r.t
                case _ => false
              }
            }.getOrElse((k, None))
        }
      }
    }
}

但是，仍然是超級命令性代碼，必須有更好的方法...