Spark:如何在Scala中从对象的ArrayType中获取单个对象

Question

我有一个包含多个列的DF，其中一个是数组类型的，名为"requests"其中包含2个字段"id";和"responses"

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

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

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

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

但是当我输入"id"时，我不知道如何获得一个请求类型的对象。作为参数?我希望数组中只有一个这样的对象，但如果有不止一个，第一个就足够了。我要将新的匹配对象添加到新列中。

Answer 1

考虑使用高阶函数transform来取消不匹配的请求，然后使用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的答案，您还可以使用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]]]   |[]             |
+-------------------------------------+---------------+

如果您只想要第一个struct对象，您可以将[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         |
+-------------------------------------+-------------+