max_by / min_by

What it does

max_by(page, ts) finds the row with the largest ts in each group and returns that row's page. min_by does the same for the smallest. It is often called argmax.

Step by step

1. 1
   Spark splits the rows into groups by session_id.
2. 2
   For each group it keeps one running pair: the largest ts seen so far and the page from that row.
3. 3
   Each new row either replaces the pair, if its ts is larger, or is ignored.
4. 4
   When the group ends, Spark returns the page half of the pair. No sorting happens at any point.

Run the example

from pyspark.sql import functions as F
result = pageviews.groupBy("session_id").agg(
    F.min_by("page", "ts").alias("entry"),
    F.max_by("page", "ts").alias("exit"),
)

SELECT session_id,
       min_by(page, ts) AS entry,
       max_by(page, ts) AS exit
FROM pageviews
GROUP BY session_id

Without it

The usual workaround ranks rows in a window and keeps the first one: three steps and a sort for a single value.

from pyspark.sql import functions as F
from pyspark.sql.window import Window
w = Window.partitionBy("session_id").orderBy(F.desc("ts"))
(pageviews.withColumn("pos", F.row_number().over(w))
    .filter("pos = 1").select("session_id", "page"))

SELECT session_id, page
FROM (
  SELECT *, ROW_NUMBER() OVER (
    PARTITION BY session_id ORDER BY ts DESC) AS rn
  FROM pageviews)
WHERE rn = 1

Performance

With many rows per group, max_by moves far less data across the network. The window version is still right when you need the top N rows, or the whole winning row with many columns and a defined tie rule.

Several columns from the winning row

Calling max_by once per column is fine when the ordering column is unique. With ties, each call can pick a different row. Pack the columns into one struct so they always come from the same row:

pageviews.groupBy("session_id").agg(
    F.max_by(F.struct("page", "referrer"), "ts").alias("last")
).select("session_id", "last.page", "last.referrer")

SELECT session_id, last.page, last.referrer
FROM (
  SELECT session_id, max_by(struct(page, referrer), ts) AS last
  FROM pageviews GROUP BY session_id)

An older trick does the same without max_by: max(struct(ts, page)). Structs compare field by field, so the maximum struct is the one with the latest ts, and you read .page from it.

Common mistakes