Scala Spark中如何高效生成数组的增量子列表/子数组

问题背景 #

原始DataFrame（list为字符串数组类型，len为对应数组长度）：

+---------------+-------------+
|           list|          len|
+---------------+-------------+
|      [a, b, c]|            3|
|[d, e, f, g, h]|            5|
+---------------+-------------+

需要生成的目标DataFrame：

+---------------+-------------+
|           list|          len|
+---------------+-------------+
|      [a, b, c]|            3|
|         [b, c]|            2|
|[d, e, f, g, h]|            5|
|   [e, f, g, h]|            4|
|      [f, g, h]|            3|
|         [g, h]|            2|
+---------------+-------------+

现有实现及问题 #

当前使用posexplode的实现代码：

val arrayData = Seq((3, List("a", "b", "c")), (5, List("d", "e", "f", "g", "h")))
val df = arrayData.toDF("len", "list")

df.select($"*", posexplode($"list").as(Seq("startIndex", "startValue")))
                .withColumn("newLength", col("len") - col("startIndex"))
                .withColumn("newList", when( col("startIndex") > 0, 
                                            slice($"list", col("startIndex")+1, col("newLength")))
                                      .otherwise(col("list")))

该实现会全量展开数组所有元素（包括仅剩余单个元素的行），生成多余中间列，内存开销较大，且输出结果不符合预期的截断要求。

优化方案 #

使用sequence生成精准索引范围，结合explode和slice实现，避免全量展开数组元素，同时减少临时列生成：

import org.apache.spark.sql.functions._
import spark.implicits._

val arrayData = Seq((3, List("a", "b", "c")), (5, List("d", "e", "f", "g", "h")))
val df = arrayData.toDF("len", "list")

val resultDF = df
  // 生成目标索引序列：从0到len-2，确保只保留到长度为2的行
  .withColumn("indices", sequence(lit(0), col("len") - 2))
  // 展开索引序列生成对应行数
  .withColumn("startIndex", explode(col("indices")))
  // 截取数组并更新长度
  .withColumn("list", slice(col("list"), col("startIndex") + 1, col("len") - col("startIndex")))
  .withColumn("len", col("len") - col("startIndex"))
  // 清理临时列
  .drop("indices", "startIndex")

resultDF.show()

代码说明 #

• sequence(lit(0), col("len") - 2)：针对每行生成精准索引范围，比如len=3时生成[0,1]，对应保留长度3和2的行；len=5时生成[0,1,2,3]，对应长度5到2的行，完全匹配预期输出的行数。
• explode(col("indices"))：仅展开需要的索引次数，相比posexplode大幅减少数据展开量，降低内存开销。
• slice函数直接根据索引截取数组，同步更新len列，最后清理临时列得到目标结构。

输出验证 #

运行后输出与预期完全一致：

+---------------+---+
|           list|len|
+---------------+---+
|      [a, b, c]|  3|
|         [b, c]|  2|
|[d, e, f, g, h]|  5|
|   [e, f, g, h]|  4|
|      [f, g, h]|  3|
|         [g, h]|  2|
+---------------+---+

内容的提问来源于stack exchange，提问作者user3388770