Spark DataFrame数据反规范化：每4行合并为1行的实现需求

To solve this problem, we'll use window functions and pivoting to group every 4 rows of the same Key into a single row with columns item_desc1 through item_desc4. Here's a step-by-step implementation in both Scala and PySpark:

Approach Overview #

1. Assign Row Numbers: Track the position of each row within its Key group to maintain order.
2. Create Chunk Groups: Group rows into chunks of 4 using integer division on the row number.
3. Pivot to Wide Format: Convert each chunk into a single row with separate columns for each item_desc in the chunk.

Scala Implementation #

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

// Sample source DataFrame
val sourceDF = Seq(
  (1, "desc1"),
  (1, "desc2"),
  (1, "desc3"),
  (1, "desc4"),
  (1, "desc5"),
  (1, "desc6"),
  (2, "descA"),
  (2, "descB")
).toDF("Key", "item_desc")

// Step 1: Assign row numbers within each Key group (order by item_desc to preserve sequence)
val withRowNum = sourceDF.withColumn(
  "row_num",
  row_number().over(Window.partitionBy("Key").orderBy("item_desc"))
)

// Step 2: Create group IDs for chunks of 4 rows
val withGroupId = withRowNum.withColumn(
  "group_id",
  (col("row_num") - 1) / 4 // Integer division to group every 4 rows
).withColumn(
  "pos_in_group",
  (col("row_num") - 1) % 4 + 1 // Position within the chunk (1-4)
)

// Step3: Pivot to create item_desc1-item_desc4 columns
val targetDF = withGroupId
  .groupBy("Key", "group_id")
  .pivot("pos_in_group", Seq(1, 2, 3, 4)) // Explicit sequence ensures column order
  .agg(first("item_desc"))
  .withColumnRenamed("1", "item_desc1")
  .withColumnRenamed("2", "item_desc2")
  .withColumnRenamed("3", "item_desc3")
  .withColumnRenamed("4", "item_desc4")
  .drop("group_id") // Remove group_id if not needed

// Show the result
targetDF.show()

PySpark Implementation #

from pyspark.sql import SparkSession
from pyspark.sql.functions import row_number, first
from pyspark.sql.window import Window

# Initialize Spark session
spark = SparkSession.builder.appName("DenormalizeDF").getOrCreate()

# Sample source DataFrame
sourceDF = spark.createDataFrame([
  (1, "desc1"),
  (1, "desc2"),
  (1, "desc3"),
  (1, "desc4"),
  (1, "desc5"),
  (1, "desc6"),
  (2, "descA"),
  (2, "descB")
], ["Key", "item_desc"])

# Step1: Assign row numbers within each Key group
window_spec = Window.partitionBy("Key").orderBy("item_desc")
withRowNum = sourceDF.withColumn("row_num", row_number().over(window_spec))

# Step2: Create group IDs and position within group
withGroupId = withRowNum.withColumn(
  "group_id",
  (withRowNum["row_num"] - 1) // 4
).withColumn(
  "pos_in_group",
  (withRowNum["row_num"] - 1) % 4 + 1
)

# Step3: Pivot to wide format
targetDF = withGroupId.groupBy("Key", "group_id") \
  .pivot("pos_in_group", [1,2,3,4]) \
  .agg(first("item_desc")) \
  .withColumnRenamed("1", "item_desc1") \
  .withColumnRenamed("2", "item_desc2") \
  .withColumnRenamed("3", "item_desc3") \
  .withColumnRenamed("4", "item_desc4") \
  .drop("group_id")

# Show result
targetDF.show()

Key Notes #

• Order Preservation: The orderBy("item_desc") in the window function ensures rows are processed in a consistent order. If your source data has a natural sequence (like a timestamp or ID), replace item_desc with that column to maintain the original order.
• Handling Partial Chunks: If a Key has fewer than 4 rows (or leftover rows after grouping into 4s), the remaining item_desc columns will be null (as seen in the sample result for Key 2 and the second row of Key 1).
• Performance: For large datasets, ensure your Key column is properly partitioned to optimize window function performance.

Sample Output #

+---+----------+----------+----------+----------+
|Key|item_desc1|item_desc2|item_desc3|item_desc4|
+---+----------+----------+----------+----------+
|  1|     desc1|     desc2|     desc3|     desc4|
|  1|     desc5|     desc6|      null|      null|
|  2|     descA|     descB|      null|      null|
+---+----------+----------+----------+----------+

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