Scala中按指定值序列过滤拆分DataFrame生成DataFrame数组

randomSplit是按设定权重随机划分数据集，无法匹配自定义值分组的需求，可按以下步骤实现目标：

步骤1：定义分组规则 #

先把预设的值和分组的对应关系固化为映射结构，后续扩展分组直接修改该映射即可：

// 分组索引 -> 该组包含的取值集合
val groupMapping = Map(
  0 -> Set("vl1", "vl4", "vl5"),
  1 -> Set("vl3", "vl2", "vl6"),
  2 -> Set("vl7", "vl8", "vl9")
)
// 两表关联的字段名，根据实际列名替换
val joinCol = "value"

步骤2：按规则拆分valueDf #

遍历分组规则，通过isin条件过滤得到每个分组对应的DataFrame，按索引顺序组成数组：

import org.apache.spark.sql.functions.col

// 按分组索引排序后逐个过滤，保证数组下标和分组编号一致
val splitDfArray = groupMapping.toSeq
  .sortBy(_._1)
  .map { case (_, valueSet) =>
    valueDf.filter(col(joinCol).isin(valueSet.toSeq: _*))
  }
  .toArray

若存在未纳入预设规则的取值，可额外新增一个兜底分组存放这部分数据，避免数据遗漏。

步骤3：关联全量表并写入目标路径 #

遍历拆分后的DataFrame数组，依次和全量表Df1做left semi join，将结果写入对应路径：

splitDfArray.zipWithIndex.foreach { case (groupDf, index) =>
  val joinedResult = Df1.join(groupDf, Seq(joinCol), "leftsemi")
  val savePath = s"/folder_Df($index)/"
  // 写入模式、存储格式可根据实际需求调整
  joinedResult.write.mode("overwrite").parquet(savePath)
}

性能优化方案 #

如果Df1数据量较大、分组数量多，循环多次join会导致全量表被重复扫描，性能较差，可改为一次打标、分区写入的方式：

• 先给valueDf的所有取值打上对应分组索引的标签
• 一次性和Df1完成join
• 按分组索引分区写入，Spark会自动将不同分组的数据分流到对应目录，仅需简单调整目录命名即可匹配要求

示例代码：

import org.apache.spark.sql.functions.{when, lit, col}

// 给取值打分组标签
val taggedValueDf = groupMapping.toSeq.foldLeft(valueDf.withColumn("group_idx", lit(-1))){
  case (tempDf, (idx, values)) =>
    tempDf.withColumn("group_idx",
      when(col(joinCol).isin(values.toSeq:_*), lit(idx))
      .otherwise(col("group_idx"))
    )
}
// 一次join全量数据
val allJoinedData = Df1.join(taggedValueDf, Seq(joinCol), "leftsemi")
// 按分组分区写入，仅需扫描Df1一次
allJoinedData.write
  .partitionBy("group_idx")
  .mode("overwrite")
  .parquet("/output_base_path/")

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