TopolearnDeduplication, sessionization and 'latest record per key' are all window-function problems in Spark.
A window function computes a value for each row from a group of related rows, without collapsing them. Define the window with Window from pyspark.sql.window:
w = Window.partitionBy("user").orderBy("ts")then apply a function .over(w):
- Ranking:
F.row_number(),F.rank(),F.dense_rank(),F.percent_rank(),F.ntile(4) - Offsets:
F.lag("amount", 1),F.lead("amount") - Aggregates:
F.sum("amount").over(w)gives a running total, because an ordered window defaults to a frame from the partition start to the current row - Without
orderBy, an aggregate overWindow.partitionBy("user")covers the whole partition, handy for shares of a total
Frames can be set explicitly: .rowsBetween(-2, Window.currentRow) for a 3-row moving window, or .rangeBetween(...) for value-based ranges like the last 7 days in seconds.
The classic latest record per key pattern: add F.row_number().over(Window.partitionBy("id").orderBy(F.desc("updated_at"))) and keep rows where it equals 1.
A window without partitionBy moves all data into a single partition, which is very slow on big data. Spark warns about this.
Example
from pyspark.sql import SparkSession, functions as F
spark = SparkSession.builder.master("local[2]").getOrCreate()
from pyspark.sql.window import Window
df = spark.createDataFrame(
[("u1", 1, 10), ("u1", 2, 30), ("u1", 3, 20), ("u2", 1, 5), ("u2", 2, 15)],
"user STRING, ts INT, amount INT",
)
by_time = Window.partitionBy("user").orderBy("ts")
whole = Window.partitionBy("user")
out = df.select(
"user", "ts", "amount",
F.sum("amount").over(by_time).alias("running"),
F.lag("amount").over(by_time).alias("prev"),
F.round(F.col("amount") / F.sum("amount").over(whole), 2).alias("share"),
)
print([tuple(r) for r in out.orderBy("user", "ts").collect()])
# [('u1', 1, 10, 10, None, 0.17), ('u1', 2, 30, 40, 10, 0.5), ('u1', 3, 20, 60, 30, 0.33),
# ('u2', 1, 5, 5, None, 0.25), ('u2', 2, 15, 20, 5, 0.75)]
latest = Window.partitionBy("user").orderBy(F.desc("ts"))
last = df.withColumn("rn", F.row_number().over(latest)).filter("rn = 1").drop("rn")
print(sorted(tuple(r) for r in last.collect())) # [('u1', 3, 20), ('u2', 2, 15)]Remember
Window.partitionBy(...).orderBy(...)+func.over(w)- Ordered windows give running aggregates; unordered ones give group totals
row_number() == 1per key = latest-record deduplication
Watch out for
- Defining a window without partitionBy on large data, forcing everything into one partition.