scala - 有没有办法使用火花数据框根据列条件从地图中获取键？

Question

我有以下数据框：

val df = Seq(
    (Map("a" -> "1", "b" -> "2", "c" -> "3"), Seq("a","b"))
    ).toDF("internalMap","commonList")

df.show()

+------------------------+----------+
|internalMap             |commonList|
+------------------------+----------+
|[a -> 1, b -> 2, c -> 3]|[a, b]    |
+------------------------+----------+

如何获取 interalMap 值，其中键数组等于 valueList 数组？

我曾尝试使用：

val getMapElements = df.select(map_keys(col("internalMap")).as("internalMapKeys"), map_values(col("internalMap")).as("internalMapValues"))
    
getMapElements.show()

+---------------+-----------------+
|internalMapKeys|internalMapValues|
+---------------+-----------------+
|      [a, b, c]|        [1, 2, 3]|
+---------------+-----------------+

    
getMapElements.select("internalMapValues").where(col("commonList") isin col("internalMapKeys")).show()

+-----------------+
|internalMapValues|
+-----------------+
+-----------------+

但它返回一个空数组。预期的是：

+-----------------+
|internalMapValues|
+-----------------+
|           [1, 2]|
+-----------------+

最困难的是，在我的用例中只接受 spark 数据框函数。

预先感谢您的帮助

标签： scalaapache-sparkapache-spark-sql

Answer 1

使用expr,transform函数来获得预期的结果。检查下面的代码。

scala> df
.withColumn(
    "internalMapValues",
    expr("transform(commonList,v -> internalMap[v])") // Pass commonList value to map as key.
)
.show(false)

+------------------------+----------+-----------------+
|internalMap             |commonList|internalMapValues|
+------------------------+----------+-----------------+
|[a -> 1, b -> 2, c -> 3]|[a, b]    |[1, 2]           |
+------------------------+----------+-----------------+