Cohort Analysis

PySpark data engineering interview problem. Difficulty: intermediate. Pattern: Aggregation. About 16 minutes. Part of the Pro drill bank.

Active users by signup cohort and transaction month. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

A cohort is a group of users who started at the same time, for example everyone who signed up in January. Cohort tables show how many users from each group are still active in later months. This is a common product analytics question. You have 2 tables. users has user_id and signup_date. txns has user_id, amount and transaction_date. The cohort of a user is the first 7 characters of their signup_date (the year and month, like 2024-01). The month of a transaction is the first 7 characters of its transaction_date. Return cohort, txn_month and active_users, the number of different users in that cohort who made a purchase in that month. Sort by cohort and txn_month. The editor creates both tables and ends with result = users. Change that last line.

Requirements

• Compute cohort from the first 7 characters of signup_date.
• Compute txn_month from the first 7 characters of transaction_date.
• Count distinct users for each (cohort, txn_month) pair.
• Return cohort, txn_month, active_users, sorted by both, in result.

Constraints

• Count distinct user_id, so a user who buys twice in a month is counted once.
• Users who never bought anything do not appear in the result.
• Use substr(1, 7) on the date text to get YYYY-MM.

Examples

Input: signups + purchases Output: cohort | txn_month | active_users 2024-01 | 2024-02 | 2 2024-01 | 2024-03 | 1 2024-02 | 2024-02 | 1 2024-02 | 2024-03 | 1 January cohort has two actives in February.

Topics: lakebench, pyspark, cohort, countDistinct.

More PySpark interview questions · All interview problems · Learn data engineering