Streaming Interview Pattern #7: Retention Cohort

Retention is one of those metrics that every product manager obsesses over, and every data engineer has to calculate at some point. This usually gets framed as a cohort problem in interviews.

You have to track weekly retention. A cohort is defined as:

• Users who signed up in week N,
• And came back and did something (any activity) in week N+1.

The events look like this:

(user_id, signup_date, activity_date)

Question:

• How would you compute weekly retention rate?
• And what changes if the data arrives in batch vs streaming?”

Questions for the Interviewer

1. How do we define “retention”? Is it any activity in week N+1? Or only specific event types (e.g., login, purchase)?
2. What is the exact time window?Do we bucket weeks using calendar weeks (Mon–Sun) or rolling 7-day windows? Do we use signup_date as cohort anchor or first_activity_date?
3. How large is the data? Millions per week? Hundreds of millions? This affects the join strategy.
4. Data arrival pattern: Is it mostly batch (daily load) or streaming (continuous events)?
5. How late can events arrive?

Hints

• A retention metric is a self-join on user activity across time windows.
• In batch: it is a cohort table + join on next week’s activity.
• In streaming: you need stateful processing with event-time windows.

Key trick: group users by signup week, then check if they show up again the next week.

Please try by yourself first before diving to solution

Approach

1. Batch Case:

• Assign each user a signup_week.
• Extract all activity_week for that user.
• Join the two: Did activity occur in signup_week+1?
• Aggregate per cohort to compute retention %.

2. Streaming Case:

• Maintain state keyed by user_id → signup_week.
• For each activity event, check if it falls in signup_week+1.
• Emit a “retained” flag when matched.
• Use watermarks to handle late activity events.

Solutions

Batch Mode (SQL Example)

WITH signup AS (
    SELECT user_id, DATE_TRUNC('week', signup_date) AS signup_week
    FROM users
),
activity AS (
    SELECT user_id, DATE_TRUNC('week', activity_date) AS activity_week
    FROM activities
)
SELECT s.signup_week,
       COUNT(DISTINCT s.user_id) AS cohort_size,
       COUNT(DISTINCT CASE WHEN a.activity_week = s.signup_week + INTERVAL '1 week' THEN s.user_id END) AS retained_users,
       COUNT(DISTINCT CASE WHEN a.activity_week = s.signup_week + INTERVAL '1 week' THEN s.user_id END) * 1.0 /
       COUNT(DISTINCT s.user_id) AS retention_rate
FROM signup s
LEFT JOIN activity a
  ON s.user_id = a.user_id
GROUP BY s.signup_week
ORDER BY s.signup_week;

Streaming Mode (Spark Structured Streaming Example)

signups = spark.readStream \
    .format("kafka") \
    .option("subscribe", "signup_events") \
    .load() \
    .withColumn("signup_week", weekofyear("signup_date"))

activities = spark.readStream \
    .format("kafka") \
    .option("subscribe", "activity_events") \
    .load() \
    .withColumn("activity_week", weekofyear("activity_date"))

# Add watermark to handle late activity events
activities = activities.withWatermark("activity_date", "7 days")

# Join on user_id
joined = signups.join(
    activities,
    (signups.user_id == activities.user_id) &
    (activities.activity_week == signups.signup_week + 1),
    "leftOuter"
)

retention = joined.groupBy("signup_week") \
    .agg(
        countDistinct("signups.user_id").alias("cohort_size"),
        countDistinct("activities.user_id").alias("retained_users")
    ) \
    .withColumn("retention_rate", col("retained_users") / col("cohort_size"))

Lessons Learned from Production Systems(common gotchas):

Late arrivals are the rule, not the exception. Always budget for a few days of lateness, but not infinite.

Partition by signup week if the scale is huge. That way, you’re only ever joining manageable chunks.

Watch for data skew. If there’s a viral campaign, that signup week will dwarf others and could cause uneven load. Usecase for salting and bucketting

Retention without signup data is nonsense. Build pipelines so signup always lands first.

Define retention with the business. Sometimes “open app once” is fine, sometimes only a purchase counts.

Edge Cases

1. Users who sign up multiple times. Which signup do we count?
2. People who never come back: Counted in denominator, but not numerator.
3. Timezone differences: Weeks should be defined in UTC for consistency.
4. Streaming cutoff: If activity happens 8 days later, do we still count it?
5. Partial weeks: How do we handle users who sign up mid-week? e.g. Someone who signs up on a Sunday is still part of that week’s cohort