火花数据框根据条件从多列中选择值

Question

数据架构，

root
|-- id: string (nullable = true)
|-- col1: string (nullable = true)
|-- col2: string (nullable = true)


|id|col1         |col2               |
|1 |["x","y","z"]|[123,"null","null"]|

从上面的数据中，我想过滤 x 在 col1 中的退出位置以及 col2 中 x 的相应值。 （col1 和 col2 的值是有序的。如果 col1 中的 x 索引 2 和 col2 中的值索引也为 2）

结果：（需要col1和col2类型的数组类型）

|id |col1 |col2 |
|1  |["x"]|[123]|

如果 col1 中不存在 x，则需要类似

的结果

|id| col1    |col2 |
|1 |["null"] |["null"]|

我试过了，

val df1 = df.withColumn("result",when($"col1".contains("x"),"X").otherwise("null"))

Answer 1

诀窍是将您的数据从愚蠢的string 列转换为更有用的数据结构。一旦col1 和col2 被重新构建为数组（或映射，正如您期望的输出所建议的那样），您可以使用Spark 的内置函数，而不是@baitmbarek 建议的混乱的UDF。

首先，使用trim 和split 将col1 和col2 转换为数组：

scala> val df = Seq(
     |       ("1", """["x","y","z"]""","""[123,"null","null"]"""),
     |         ("2", """["a","y","z"]""","""[123,"null","null"]""")
     |     ).toDF("id","col1","col2")
df: org.apache.spark.sql.DataFrame = [id: string, col1: string ... 1 more field]

scala> val df_array = df.withColumn("col1", split(trim($"col1", "[\"]"), "\"?,\"?"))
                        .withColumn("col2", split(trim($"col2", "[\"]"), "\"?,\"?"))
df_array: org.apache.spark.sql.DataFrame = [id: string, col1: array<string> ... 1 more field]

scala> df_array.show(false)
+---+---------+-----------------+
|id |col1     |col2             |
+---+---------+-----------------+
|1  |[x, y, z]|[123, null, null]|
|2  |[a, y, z]|[123, null, null]|
+---+---------+-----------------+


scala> df_array.printSchema
root
 |-- id: string (nullable = true)
 |-- col1: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- col2: array (nullable = true)
 |    |-- element: string (containsNull = true)

从这里，您应该能够使用array_position 在col1 中找到“x”的索引（如果有）并从col2 中检索匹配数据，从而实现您想要的目标。但是，首先将两个数组转换为映射应该可以更清楚地理解您的代码在做什么：

scala> val df_map = df_array.select(
                        $"id", 
                        map_from_entries(arrays_zip($"col1", $"col2")).as("col_map")
                        )
df_map: org.apache.spark.sql.DataFrame = [id: string, col_map: map<string,string>]

scala> df_map.show(false)
+---+--------------------------------+
|id |col_map                         |
+---+--------------------------------+
|1  |[x -> 123, y -> null, z -> null]|
|2  |[a -> 123, y -> null, z -> null]|
+---+--------------------------------+
scala> val df_final = df_map.select(
                                $"id",
                                when(isnull(element_at($"col_map", "x")), 
                                    array(lit("null")))
                                .otherwise(
                                    array(lit("x")))
                                .as("col1"),  
                                when(isnull(element_at($"col_map", "x")), 
                                    array(lit("null")))
                                .otherwise(
                                    array(element_at($"col_map", "x")))
                                .as("col2")
                                )
df_final: org.apache.spark.sql.DataFrame = [id: string, col1: array<string> ... 1 more field]

scala> df_final.show
+---+------+------+
| id|  col1|  col2|
+---+------+------+
|  1|   [x]| [123]|
|  2|[null]|[null]|
+---+------+------+
scala> df_final.printSchema
root
 |-- id: string (nullable = true)
 |-- col1: array (nullable = false)
 |    |-- element: string (containsNull = false)
 |-- col2: array (nullable = false)
 |    |-- element: string (containsNull = true)