如何为Spark DataFrame新增列，用带冒号的行值作为前缀填充后续行？

这是一个典型的分组填充+组内聚合场景，其实不需要用UDF就能优雅解决（UDF在Spark中往往不如内置函数高效，还可能带来序列化开销）。下面是具体的思路和代码实现：

核心思路 #

1. 划分逻辑分组：给每个带冒号的行（如Total:、Male:）分配唯一分组ID，后续行继承该ID直到下一个冒号行出现，用窗口函数就能实现。
2. 提取组内关键值：对每个分组，提取两个核心信息：
   • 分组的前缀（即组内带冒号的name值）
   • 组内第一个非冒号行的name值（如Under 5 years）
3. 拼接生成目标列：关联原始数据和分组信息，根据当前行是否为冒号行，生成对应的new_name：冒号行直接复用原name，其他行用前缀+第一个子项拼接。

Scala代码实现 #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window

// 构建示例原始DataFrame
val df = spark.createDataFrame(Seq(
  (1, "Total:"),
  (2, "Male:"),
  (3, "Under 5 years"),
  (4, "5 to 9 years"),
  (5, "10 to 14 years"),
  (6, "Female:"),
  (7, "Under 5 years"),
  (8, "5 to 9 years"),
  (9, "10 to 14 years")
)).toDF("id", "name")

// 步骤1：生成分组ID，区分每个冒号行及其后续子项
val groupWindow = Window.orderBy("id")
val dfWithGroup = df.withColumn(
  "group_id",
  sum(when(col("name").contains(":"), 1).otherwise(0)).over(groupWindow)
)

// 步骤2：聚合每个分组的前缀和第一个子项
val groupInfo = dfWithGroup.groupBy("group_id")
  .agg(
    // 提取组内第一个带冒号的行作为前缀
    first(when(col("name").contains(":"), col("name"))).alias("prefix"),
    // 提取组内第一个非冒号的行作为子项模板
    first(when(!col("name").contains(":"), col("name"))).alias("first_subitem")
  )

// 步骤3：关联分组信息，生成new_name列
val resultDF = dfWithGroup.join(groupInfo, "group_id")
  .withColumn(
    "new_name",
    // 冒号行直接用原name，其他行拼接前缀和子项
    when(col("name").contains(":"), col("name"))
      .otherwise(concat_ws(" ", col("prefix"), col("first_subitem")))
  )
  .select("id", "name", "new_name")

// 查看结果
resultDF.show(false)

代码说明 #

• 分组ID生成：sum(when(...)).over(groupWindow)会在遇到冒号行时累加1，后续行的group_id会和当前冒号行保持一致，直到下一个冒号行触发新的分组。
• 组内聚合：first(when(...))确保我们拿到的是分组内第一个符合条件的值，避免了复杂的排序逻辑。
• 目标列生成：用when分支处理两种情况，concat_ws负责安全拼接字符串（自动处理空值）。

运行代码后，输出结果完全符合你期望的格式。

内容的提问来源于stack exchange，提问作者J.Hammond