Scala DataFrame基于ID父子关系的Window Lag函数实现需求

Got it, let's work through this problem. The core issue here is that your original code only calculates lag within the same exact id group, but we need child IDs (like A-1) to inherit the lag values from their parent IDs (like A) instead of their own group. Here's a step-by-step solution:

Step 1: Extract Parent ID from Each Record #

First, we need to map every ID to its parent. Using a regex is robust here—it works even if child IDs have multiple hyphens (like A-1-2).

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

// Add a parent_id column to map child IDs to their parent
val dfWithParent = DF.withColumn(
  "parent_id",
  regexp_extract(col("id"), "^([^-]+)", 1) // Grabs everything before the first hyphen
)

Step 2: Calculate Lag Values for Parent IDs Only #

Next, we isolate parent ID records (those without hyphens) and calculate their lag values using your original window logic—since parents should still get their own previous col1 values.

// Define window for parent ID grouping and time sorting
val parentWindow = Window.partitionBy("parent_id").orderBy("time")

// Process parent records to get their lag values
val parentWithLag = dfWithParent.filter(!col("id").contains("-"))
  .withColumn("col1_lag", lag(col("col1"), 1).over(parentWindow))
  .select("parent_id", "time", "col1_lag")

Step 3: Associate Parent Lag Values to All Records #

Now we need to link every record (parent or child) to the correct parent lag value. For child IDs, we'll find the most recent parent record that's earlier than or equal to the child's timestamp and use its lag value.

// Join the original data with parent lag data, matching parent_id and timestamps
val resultDF = dfWithParent.join(
  parentWithLag,
  dfWithParent("parent_id") === parentWithLag("parent_id") && parentWithLag("time") <= dfWithParent("time"),
  "left_outer"
)
// For each record, keep only the latest matching parent lag value
.withColumn("rank", row_number().over(Window.partitionBy(dfWithParent("id"), dfWithParent("time")).orderBy(parentWithLag("time").desc)))
.filter(col("rank") === 1)
// Drop temporary columns and keep original columns plus the lag feature
.select(DF.columns.map(col): _*, col("col1_lag"))

How This Works #

• Parent Records: They'll match their own lag values exactly like your original code did. For example, parent A's 08:00 record gets a lag of 3 (from A's 07:00 record).
• Child Records: They'll inherit the lag value from the most recent parent record before their timestamp. For example, all A-1 records (10:00, 11:00, 12:00) will get a lag of 2 (from parent A's 09:00 record, whose lag is the 08:00 value of 2).

Example Result Snippet #

For the sample data you provided, part of the result will look like this:

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