基于条件向Spark DataFrame追加新行的技术咨询

需求说明 #

给定初始Spark DataFrame，需为每个item_nbr下所有edit_flag=false的(A,B)组合添加新行：新行除rcvr_nbr和rcv_qty取自同item_nbr下edit_flag=true的行外，其余字段完全复制自原edit_flag=false的行，最终合并原数据与新行得到目标结果。

初始数据 #

+---+---+--------+---------+-------+--------+---------+
|  A|  B|rcvr_nbr|order_qty|rcv_qty|item_nbr|edit_flag|
+---+---+--------+---------+-------+--------+---------+
|123|  1|     500|       10|      2|    1001|    false|
|123|  1|     501|       10|      2|    1001|    false|
|123|  4|     502|       60|      5|    1001|     true|
|123|  2|     504|       40|     30|    1003|    false|
|123|  5|     510|       10|     10|    1003|     true|
+---+---+--------+---------+-------+--------+---------+

目标结果 #

+---+---+--------+---------+-------+--------+---------+
|  A|  B|rcvr_nbr|order_qty|rcv_qty|item_nbr|edit_flag|
+---+---+--------+---------+-------+--------+---------+
|123|  1|     500|       10|      2|    1001|    false|
|123|  1|     501|       10|      2|    1001|    false|
|123|  1|     502|       10|      5|    1001|    false|
|123|  4|     502|       60|      5|    1001|     true|
|123|  2|     504|       40|     30|    1003|    false|
|123|  2|     510|       40|     10|    1003|    false|
|123|  5|     510|       10|     10|    1003|     true|
+---+---+--------+---------+-------+--------+---------+

实现思路 #

1. 分离edit_flag=true和edit_flag=false的数据集，分别提取所需字段
2. 按item_nbr关联两类数据集，生成需要新增的行
3. 合并原数据集与新增行，按指定字段排序得到最终结果

Python 代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# 初始化SparkSession
spark = SparkSession.builder.appName("ComplexDFTransform").getOrCreate()

# 构造初始DataFrame
data = [
    (123, 1, 500, 10, 2, 1001, False),
    (123, 1, 501, 10, 2, 1001, False),
    (123, 4, 502, 60, 5, 1001, True),
    (123, 2, 504, 40, 30, 1003, False),
    (123, 5, 510, 10, 10, 1003, True)
]
schema = ["A", "B", "rcvr_nbr", "order_qty", "rcv_qty", "item_nbr", "edit_flag"]
df = spark.createDataFrame(data, schema)

# 提取edit_flag=true的行，保留关键字段
true_df = df.filter(col("edit_flag") == True).select("item_nbr", "rcvr_nbr", "rcv_qty")

# 提取edit_flag=false的行，保留需复用的字段
false_df = df.filter(col("edit_flag") == False).select("A", "B", "order_qty", "item_nbr", "edit_flag")

# 关联生成新增行
new_rows = false_df.join(true_df, on="item_nbr", how="inner")

# 合并原数据与新增行并排序
final_df = df.union(new_rows).orderBy("item_nbr", "A", "B")

# 展示结果
final_df.show()

Scala 代码实现 #

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

object ComplexDFTransform {
  def main(args: Array[String]): Unit = {
    // 初始化SparkSession
    val spark = SparkSession.builder.appName("ComplexDFTransform").getOrCreate()
    import spark.implicits._

    // 构造初始DataFrame
    val data = Seq(
      (123, 1, 500, 10, 2, 1001, false),
      (123, 1, 501, 10, 2, 1001, false),
      (123, 4, 502, 60, 5, 1001, true),
      (123, 2, 504, 40, 30, 1003, false),
      (123, 5, 510, 10, 10, 1003, true)
    )
    val df = data.toDF("A", "B", "rcvr_nbr", "order_qty", "rcv_qty", "item_nbr", "edit_flag")

    // 提取edit_flag=true的行
    val trueDF = df.filter(col("edit_flag") === true).select("item_nbr", "rcvr_nbr", "rcv_qty")

    // 提取edit_flag=false的行
    val falseDF = df.filter(col("edit_flag") === false).select("A", "B", "order_qty", "item_nbr", "edit_flag")

    // 关联生成新增行
    val newRows = falseDF.join(trueDF, Seq("item_nbr"), "inner")

    // 合并并排序
    val finalDF = df.union(newRows).orderBy("item_nbr", "A", "B")

    // 展示结果
    finalDF.show()
  }
}

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