在Scala中使用Spark数据集执行类型化联接

Question

我喜欢Spark数据集，因为它们在编译时会给我带来分析错误和语法错误，还允许我使用getter而不是硬编码的名称/数字。大多数计算都可以通过DataSet的高级API完成。例如，通过访问Dataset类型对象的执行agg、select、sum、avg、map、filter或groupBy操作要比使用RDD行的数据字段简单得多。

但是其中缺少join操作，我读到我可以像这样执行join操作

ds1.joinWith(ds2, ds1.toDF().col("key") === ds2.toDF().col("key"), "inner")

ds1.joinWith(ds2, ds1.key === ds2.key, "inner")

我使用的是Spark 1.6.1和Scala 2.10

Answer 1

Spark SQL只能在联接条件基于相等运算符的情况下优化联接。这意味着我们可以分别考虑等价连接和非等价连接。

通过将数据集映射到（key,value）元组，执行基于key的联接，并重新塑造结果，可以以类型安全的方式实现Equijoin：

import org.apache.spark.sql.Encoder
import org.apache.spark.sql.Dataset

def safeEquiJoin[T, U, K](ds1: Dataset[T], ds2: Dataset[U])
    (f: T => K, g: U => K)
    (implicit e1: Encoder[(K, T)], e2: Encoder[(K, U)], e3: Encoder[(T, U)]) = {
  val ds1_ = ds1.map(x => (f(x), x))
  val ds2_ = ds2.map(x => (g(x), x))
  ds1_.joinWith(ds2_, ds1_("_1") === ds2_("_1")).map(x => (x._1._2, x._2._2))
}

可以用关系代数运算符表示为R？S=？（R？S）并直接转换为代码。

spark.conf.set("spark.sql.crossJoin.enabled", true)

def safeNonEquiJoin[T, U](ds1: Dataset[T], ds2: Dataset[U])
                         (p: (T, U) => Boolean) = {
  ds1.joinWith(ds2, lit(true)).filter(p.tupled)
}

def safeNonEquiJoin[T, U](ds1: Dataset[T], ds2: Dataset[U])
    (p: (T, U) => Boolean)
    (implicit e1: Encoder[Tuple1[T]], e2: Encoder[Tuple1[U]], e3: Encoder[(T, U)]) = {
  ds1.map(Tuple1(_)).crossJoin(ds2.map(Tuple1(_))).as[(T, U)].filter(p.tupled)
}

case class LabeledPoint(label: String, x: Double, y: Double)
case class Category(id: Long, name: String)

val points1 = Seq(LabeledPoint("foo", 1.0, 2.0)).toDS
val points2 = Seq(
  LabeledPoint("bar", 3.0, 5.6), LabeledPoint("foo", -1.0, 3.0)
).toDS
val categories = Seq(Category(1, "foo"), Category(2, "bar")).toDS

safeEquiJoin(points1, categories)(_.label, _.name)
safeNonEquiJoin(points1, points2)(_.x > _.x)

如果您不局限于Spark SQL APIframeless为数据集提供了有趣的类型安全扩展（到目前为止，它只支持Spark2.0):

import frameless.TypedDataset

val typedPoints1 = TypedDataset.create(points1)
val typedPoints2 = TypedDataset.create(points2)

typedPoints1.join(typedPoints2, typedPoints1('x), typedPoints2('x))

datasetAPI在1.6中不稳定，所以我认为在1.6中使用它没有意义。

当然，这种设计和描述性的名称是不必要的。您可以轻松地使用type class将此方法隐式添加到DataSet中，并且与内置签名没有冲突，因此两者都可以称为JoinWith。