在Spark SQL表中展开（转置？）多个列

Question

我正在使用Spark SQL（我提到它在Spark中，以防对SQL语法有影响 - 我还不太确定），我有一个表格，我试图重新构造它，但我卡在了同时转置多列的步骤。

基本上我的数据看起来像：

userId    someString      varA     varB
   1      "example1"    [0,2,5]   [1,2,9]
   2      "example2"    [1,20,5]  [9,null,6]

我希望同时展开变量varA和varB（长度始终保持一致），以便最终输出如下：

userId    someString      varA     varB
   1      "example1"       0         1
   1      "example1"       2         2
   1      "example1"       5         9
   2      "example2"       1         9
   2      "example2"       20       null
   2      "example2"       5         6

但是我似乎只能在一个命令中使用单个的 explode(var) 语句，如果我尝试链接它们（例如，在第一个 explode 命令之后创建临时表），那么我显然会得到大量重复且不必要的行。

非常感谢！

- anthr

Answer 1

Spark >= 2.4

您可以跳过使用 zip 和 udf，而使用 arrays_zip 函数：

df.withColumn("vars", explode(arrays_zip($"varA", $"varB"))).select(
  $"userId", $"someString",
  $"vars.varA", $"vars.varB").show

Spark < 2.4

如果你没有自定义 UDF，是无法实现你想要的功能的。在 Scala 中，你可以这样做：

val data = sc.parallelize(Seq(
    """{"userId": 1, "someString": "example1",
        "varA": [0, 2, 5], "varB": [1, 2, 9]}""",
    """{"userId": 2, "someString": "example2",
        "varA": [1, 20, 5], "varB": [9, null, 6]}"""
))

val df = spark.read.json(data)

df.printSchema
// root
//  |-- someString: string (nullable = true)
//  |-- userId: long (nullable = true)
//  |-- varA: array (nullable = true)
//  |    |-- element: long (containsNull = true)
//  |-- varB: array (nullable = true)
//  |    |-- element: long (containsNull = true)

现在我们可以定义zip的UDF：

import org.apache.spark.sql.functions.{udf, explode}

val zip = udf((xs: Seq[Long], ys: Seq[Long]) => xs.zip(ys))

df.withColumn("vars", explode(zip($"varA", $"varB"))).select(
   $"userId", $"someString",
   $"vars._1".alias("varA"), $"vars._2".alias("varB")).show

// +------+----------+----+----+
// |userId|someString|varA|varB|
// +------+----------+----+----+
// |     1|  example1|   0|   1|
// |     1|  example1|   2|   2|
// |     1|  example1|   5|   9|
// |     2|  example2|   1|   9|
// |     2|  example2|  20|null|
// |     2|  example2|   5|   6|
// +------+----------+----+----+

使用原生 SQL 语句：

sqlContext.udf.register("zip", (xs: Seq[Long], ys: Seq[Long]) => xs.zip(ys))
df.registerTempTable("df")

sqlContext.sql(
  """SELECT userId, someString, explode(zip(varA, varB)) AS vars FROM df""")

- zero323

---

这个能应用在类型为序列的三列上吗？ - Amit Kumar

@AmitKumar 当然可以，你需要调整签名和正文，但这并不难。 - zero323

我想知道在新的数据集API中，是否可以只使用map和zip将数组组合在一起，而不需要创建UDF，以及这样做是否会更快/更具扩展性/是否会被Catalyst执行引擎优化。当我在控制台时，我会尝试一下。 - Davos

@zero323，你能帮我解决如何在Java中编写超过三列的UDF吗？ - ROOT

@SatishKaruturi，从Spark 2.4.0开始，您不再需要编写自己的UDF。有一个新的arrays_zip函数可以应用于多个列。 - haimco

Answer 2

你也可以尝试

case class Input(
 userId: Integer,
 someString: String,
 varA: Array[Integer],
 varB: Array[Integer])

case class Result(
 userId: Integer,
 someString: String,
 varA: Integer,
 varB: Integer)

def getResult(row : Input) : Iterable[Result] = {
 val user_id = row.user_id
 val someString = row.someString
 val varA = row.varA
 val varB = row.varB
 val seq = for( i <- 0 until varA.size) yield {Result(user_id,someString,varA(i),varB(i))}
 seq
 }

val obj1 = Input(1, "string1", Array(0, 2, 5), Array(1, 2, 9))
val obj2 = Input(2, "string2", Array(1, 3, 6), Array(2, 3, 10))
val input_df = sc.parallelize(Seq(obj1, obj2)).toDS

val res = input_df.flatMap{ row => getResult(row) }
res.show
// +------+----------+----+-----+
// |userId|someString|varA|varB |
// +------+----------+----+-----+
// |     1|  string1 |   0|   1 |
// |     1|  string1 |   2|   2 |
// |     1|  string1 |   5|   9 |
// |     2|  string2 |   1|   2 |
// |     2|  string2 |   3|   3 |
// |     2|  string2 |   6|   10|
// +------+----------+----+-----+

- Alexandru Ivana

Answer 3

即使我们有超过3列，这也可以正常工作

case class Input(user_id: Integer, someString: String, varA: Array[Integer], varB: Array[Integer], varC: Array[String], varD: Array[String])

val obj1 = Input(1, "example1", Array(0,2,5), Array(1,2,9), Array("a","b","c"), Array("red","green","yellow"))
val obj2 = Input(2, "example2", Array(1,20,5), Array(9,null,6), Array("d","e","f"), Array("white","black","cyan"))
val obj3 = Input(3, "example3", Array(10,11,12), Array(5,8,7), Array("g","h","i"), Array("blue","pink","brown"))

val input_df = sc.parallelize(Seq(obj1, obj2, obj3)).toDS
input_df.show()

val zip = udf((a: Seq[String], b: Seq[String], c: Seq[String], d: Seq[String]) => {a.indices.map(i=> (a(i), b(i), c(i), d(i)))})

val output_df = input_df.withColumn("vars", explode(zip($"varA", $"varB", $"varC", $"varD"))).
                         select($"user_id", $"someString", $"vars._1".alias("varA"), $"vars._2".alias("varB"), $"vars._3".alias("varC"), $"vars._4".alias("varD"))
output_df.show()

- Aman Mundra