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

Question

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

但是这里缺少连接操作，我读到我可以像这样进行连接

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

但这不是我想要的，因为我更愿意通过 case class 界面来完成，所以更像这样

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

目前最好的选择似乎是在 case class 旁边创建一个对象，并提供此函数以向我提供正确的列名称作为字符串。所以我会使用第一行代码，但放置一个函数而不是硬编码的列名。但这感觉不够优雅..

有人可以在这里就其他选项给我建议吗？目标是从实际的列名中抽象出来，并且最好通过 case class.

的 getter 来工作

我正在使用 Spark 1.6.1 和 Scala 2.10

Answer 1

观察

Spark SQL 只有在连接条件基于相等运算符时才能优化连接。这意味着我们可以分别考虑等值连接和非等值连接。

等值连接

Equijoin 可以通过将 Datasets 映射到（键，值）元组，基于键执行连接并重塑结果来以类型安全的方式实现：

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 2.0

启用 crossJoin 并使用 joinWith 与平凡相等的谓词：

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)
}

Spark 2.1

使用crossJoin方法：

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)

备注

应该注意的是，这些方法在质量上不同于直接 joinWith 应用程序，并且需要昂贵的 DeserializeToObject / SerializeFromObject 转换（与直接 joinWith 可以对数据进行逻辑运算）。

这类似于中描述的行为。

如果您不限于 Spark SQL API frameless 为 Datasets 提供有趣的类型安全扩展（截至今天仅支持 Spark 2.0):

import frameless.TypedDataset

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

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

Dataset API 在 1.6 中不稳定，所以我认为在那里使用它没有意义。
当然这样的设计和描述性的名字不是必须的。您可以轻松地使用类型 class 将此方法隐式添加到 Dataset 并且与内置签名没有冲突，因此两者都可以调用 joinWith.

Answer 2

此外，类型不安全的 Spark API 的另一个更大的问题是，当您连接两个 Datasets 时，它会给您一个 DataFrame。然后您会丢失原始两个数据集中的类型。

val a: Dataset[A]
val b: Dataset[B]

val joined: Dataframe = a.join(b)
// what would be great is 
val joined: Dataset[C] = a.join(b)(implicit func: (A, B) => C)

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

Perform a typed join in Scala with Spark Datasets

scala

apache-spark

join

apache-spark-sql

apache-spark-dataset

观察

等值连接

非等值连接

Spark 2.0

Spark 2.1

例子

备注