Spark：如何从 Scala 中的 ArrayType 对象中获取单个 object

Question

我有一个包含多个列的 DF，其中一个列在 ArrayType 中，称为“requests”，其中包含 2 个字段“id”和“responses”：

(requests,ArrayType(StructType(StructField(id,IntegerType,true),StructField(responses,ArrayType(IntegerType,false),true))))

从“请求”数组中，我想获取一个“id”与特定值匹配的请求，并将其添加到新列中。

到目前为止，我添加了一个 boolean 值来表示该值是否存在于列表中：

dF
  .withColumn("idPresent", array_contains(col("requests.id"), 55))
  .show()

但是当我将“id”作为参数时，我无法弄清楚如何获得一个 object 类型的请求？ 我希望只有一个这样的 object 出现在阵列中，但如果有多个，第一个就足够了。 我想将新匹配的 object 添加到新列中。

Answer 1

考虑使用高阶 function transform来取消与“id”不匹配的请求，然后使用array_except删除那些 null 元素：

case class Request(id: Int, responses: Seq[Int])

val df = Seq(
  Seq(Request(1, Seq(11)), Request(2, Seq(21, 22)), Request(3, Seq(31))),
  Seq(Request(4, Seq(41, 42, 43)), Request(5, Seq(51, 52)))
).toDF("requests")

df.
  withColumn("request_id2", array_except(
      expr("transform(requests, r -> case when r.id = 2 then r end)"),
      array(lit(null))
    )
  ).show(false)
// +-------------------------------------+---------------+
// |requests                             |request_id2    |
// +-------------------------------------+---------------+
// |[[1, [11]], [2, [21, 22]], [3, [31]]]|[[2, [21, 22]]]|
// |[[4, [41, 42, 43]], [5, [51, 52]]]   |[]             |
// +-------------------------------------+---------------+

Answer 2

扩展@Leo C's answer ，您还可以使用filter根据它们的 id 过滤数组元素：

df.withColumn("request_id2", expr("filter(requests, r -> r.id = 2)")).show(false)
+-------------------------------------+---------------+
|requests                             |request_id2    |
+-------------------------------------+---------------+
|[[1, [11]], [2, [21, 22]], [3, [31]]]|[[2, [21, 22]]]|
|[[4, [41, 42, 43]], [5, [51, 52]]]   |[]             |
+-------------------------------------+---------------+

如果您只想要第一个结构 object，您可以将[0]添加到expr ：

df.withColumn("request_id2", expr("filter(requests, r -> r.id = 2)[0]")).show(false)
+-------------------------------------+-------------+
|requests                             |request_id2  |
+-------------------------------------+-------------+
|[[1, [11]], [2, [21, 22]], [3, [31]]]|[2, [21, 22]]|
|[[4, [41, 42, 43]], [5, [51, 52]]]   |null         |
+-------------------------------------+-------------+