scala - 使用来自另一列的键从 MapType 列查找值

Question

我有一个包含两列的 Spark 流数据框。 Integer id 列和 MapType 列，其中 Integer Id 作为键，JSON 对象作为值。

---------------------------------------------------------------------------------------------------------------
id             objects
---------------------------------------------------------------------------------------------------------------
1     |    (1 -> {"id" : 1, "type": "jpeg"}, 2 -> {"id" : 2, "type": "gif"}, 3 -> {"id" : 3, "type": "png"})
5     |    (1 -> {"id" : 1, "type": "jpeg"}, 2 -> {"id" : 2, "type": "gif"}, 3 -> {"id" : 3, "type": "png"})
2     |    (1 -> {"id" : 1, "type": "jpeg"}, 2 -> {"id" : 2, "type": "gif"}, 3 -> {"id" : 3, "type": "png"})
---------------------------------------------------------------------------------------------------------------

我想构建一个新的数据框，其中有一个列包含 JSON 对象，其键与 id 列匹配。

----------------------------------------------------------------------
objects
----------------------------------------------------------------------
{"id" : 1, "type": "jpeg"}
{"id" : 2, "type": "gif" }
----------------------------------------------------------------------

执行此操作的最佳方法是什么？我实现了一个 Scala udf 来执行查找并返回相应的对象值，但想知道是否可以使用内置的 Spark 函数来完成同样的操作。

我试着做这样的事情:

df.withColumn("obj", $"objects".getItem($"id"))

但是它抛出一个异常:

java.lang.RuntimeException: Unsupported literal type class org.apache.spark.sql.ColumnName

这是有道理的，因为 $"id" 是一种列类型。但是，如果我进行收集，那将导致我希望避免的行为。

Answer 1

不需要 UDF，您可以使用内置功能来执行此操作。但是，您不能使用 getItem，因为参数是另一列，而不是字符串值。

相反，您可以按如下方式从 Map 中获取值:

df.withColumn("value", $"objects"($"id"))

要创建一个新的数据框并删除 id 在 map 中不存在的行，

df.select($"objects"($"id").as("objects")).na.drop

这会给你，

+-------------------------+
|objects                  |
+-------------------------+
|{"id": 1, "type": "jpeg"}|
|{"id": 3, "type": "png"} |
+-------------------------+