如何使用 SQL 语言有条件地替换 Spark SQL 数组值？

Question

我的表中有此列：

myColumn
[red, green]
[green, green, red]

我需要修改它，以便我可以用 1 替换红色，用 2 替换绿色：

myColumn
[1, 2]
[2, 2, 1]

简而言之，有没有办法对数组中的每个元素逐行应用 case 子句？

到目前为止我最接近的：

select replace(replace(to_json(myColumn), 'red', 1), 'green', 2)

另一方面，如果我们有一列字符串，我可以简单地使用：

select (
  case
    when myColumn='red' then 1
    when myColumn='green' then 2
  end
) from myTable;

Answer 1

让我们创建一些示例数据和包含替换的地图：你想制作

val df = Seq((1, Seq("red", "green")),
             (2, Seq("green", "green", "red")))
         .toDF("id", "myColumn")
val values = Map("red" -> "1", "green" -> "2")

最直接的方法是定义一个完全符合您要求的 UDF：

val replace = udf((x : Array[String]) =>
    x.map(value => values.getOrElse(value, value)))
df.withColumn("myColumn", replace('myColumn)).show

+---+---------+
| id| myColumn|
+---+---------+
|  1|   [1, 2]|
|  2|[2, 2, 1]|
+---+---------+

如果没有 UDF，您可以使用 concat_ws 使用数组中没有的分隔符将数组转换为字符串。然后我们可以使用字符串函数进行编辑：

val sep = ","
val replace = values
    .foldLeft(col("myColumn")){ case (column, (key, value)) =>
        regexp_replace(column, sep + key + sep, sep + value + sep) 
    }
df.withColumn("myColumn", concat(lit(sep), concat_ws(sep+sep, 'myColumn), lit(sep)))
  .withColumn("myColumn", regexp_replace(replace, "(^,)|(,$)", ""))
  .withColumn("myColumn", split('myColumn, sep+sep))
  .show

Answer 2

假设dataframe已经注册了一个名为tmp的临时视图，使用下面的SQL语句得到结果。

sql = """
    select
        collect_list(
            case col
                when 'red' then 1
                when 'green' then 2
            end)
        myColumn
    from
        (select mid,explode(myColumn) col
        from
            (select monotonically_increasing_id() mid,myColumn
            from tmp)
        )
    group by mid
"""
df = spark.sql(sql)
df.show(truncate=False)

Answer 3

在纯 Spark SQL 中，您可以使用 concat_ws 将数组转换为字符串，使用 regexp_replace 进行替换，然后使用 split 重新创建数组。

select split(
    regexp_replace(
        regexp_replace(
            concat_ws(',', myColumn)
        , 'red', '1')
    , 'green', '2')
, ',') myColumn from df

Answer 4

我可以执行简单的转换（Spark 3 及更高版本）

select transform(myColumn, value ->
  case value
    when 'red' then 1
    when 'green' then 2
  end
from myTable

如何使用 SQL 语言有条件地替换 Spark SQL 数组值？

How to conditionally replace Spark SQL array values using SQL language?

apache-spark

apache-spark-sql