Spark：使用 Map iterable 将数组转换为具有数组大小的列

Question

我有以下数据

df.show

+---------+--------------------+------------+
|     name| age|         tokens| tokensCount|
+---------+----+---------------+------------+
|    Alice|  29|        [A,B,C]|           3|
|      Bob|  28|      [A,B,C,D]|           4|
|  Charlie|  29|    [A,B,C,D,E]|           5|    
+---------+----+---------------+------------+

我使用以下命令转换数据

val newDF = df.select(($"name") +: (0 until 4).map(i => ($"tokens")(i).alias(s"token$i")): _*).show

+---------+-------+-------+-------+-------+
|     name| token0| token1| token2| token3| 
+---------+-------+-------+-------+-------+
|    Alice|      A|      B|      C|   null|
|      Bob|      A|      B|      C|      D|
|  Charlie|      A|      B|      C|      D|
+---------+-------+-------+-------+-------+

我想在 (0 until 4)

处给出 tokensCount 而不是静态值 4

我尝试了一些方法，例如 $"tokensCount" 和 size($"tokens")，但无法通过。

谁能建议如何根据数组的大小或数组的数量进行循环或映射？

非常感谢

Answer 1

您可以修改代码以找到 tokens 的最大长度，然后使用它来创建必要的列：

val df = Seq(
  ("Alice", 29, Array("A", "B", "C")),
  ("Bob", 28, Array("A", "B", "C", "D")),
  ("Charlie", 29, Array("A", "B", "C", "D", "E")),
).toDF("name", "age", "tokens")

val maxTokenCount = df.withColumn("token_count", size(col("tokens"))).agg(max("token_count") as "mtc")

val newDF = df.select(($"name") +: (0 until maxTokenCount.first.getInt(0)).map(i => ($"tokens")(i).alias(s"token$i")): _*).show

哪个会给你：

+-------+------+------+------+------+------+
|   name|token0|token1|token2|token3|token4|
+-------+------+------+------+------+------+
|  Alice|     A|     B|     C|  null|  null|
|    Bob|     A|     B|     C|     D|  null|
|Charlie|     A|     B|     C|     D|     E|
+-------+------+------+------+------+------+

解释您为什么要进行此转换可能会有用，因为可能有更有效的方法。这有可能创建一个非常稀疏的数据框。想象一下，大多数名字都没有标记，但 Bob 有 100 个标记：突然间你有一百列大部分是空值。

Spark：使用 Map iterable 将数组转换为具有数组大小的列

Spark: Transform array to Column with size of Array using Map iterable

scala

apache-spark

apache-spark-sql