Fix the cohort retention check for a dashboard query

Review the AI's codeEasy15 minSnowflake style

Your task

An agent wrote a 60-line SQL query for monthly cohort retention, and nobody on the team fully follows its window functions. Before executives see the numbers, you check the query against an oracle: a second calculation, small enough to be plainly right, that the query must agree with.

An assistant wrote cohort_retention_counts(signup_lines, activity_lines, cohort_month, max_offset). It is in the editor. It has a bug. Find it and fix it.

Input and output:

  • signup_lines: one string per signup, "user_id|date", with dates like "2026-01-05".
  • activity_lines: one string per activity event, in the same format.
  • cohort_month: the signup month to study, as "YYYY-MM".
  • max_offset: how many months after the cohort month to count, an integer.
  • Return a list of integers.

Rules the tests check:

  • Split each signup or activity line on |. Skip a line that does not have exactly 2 parts, or whose date part is shorter than 7 characters.
  • A line's month is the first 7 characters of its date. The rest of the date is not checked.
  • User ids compare exactly. An empty id is still an id.
  • cohort_month must be 7 characters: 4 digits, -, then 2 digits from 01 to 12. Otherwise return an empty list. A negative max_offset also returns an empty list.
  • The cohort is every distinct user with at least one signup line in cohort_month. A user with several signup lines in that month is one member.
  • Index 0 of the result is the size of the cohort.
  • For each i from 1 to max_offset, index i is the number of distinct cohort members with at least one activity in the month i months after cohort_month. Months roll into the next year: 1 month after "2025-12" is "2026-01".
  • Several activities by one member in one month count once. Activity by users outside the cohort is ignored, and so is activity in any month that is not one of the counted months.
  • The result always has max_offset + 1 entries, with 0 for a month nobody was active.

Examples

cohort_retention_counts(signup_lines: list[str], activity_lines: list[str], cohort_month: str, max_offset: int) -> list[int]
Example 1
Inputcohort_retention_counts(["u1|2026-01-05", "u2|2026-01-20", "u3|2026-02-02", "u4|2026-01-31"], ["u1|2026-02-11", "u1|2026-02-28", "u2|2026-03-01", "u3|2026-02-03"], "2026-01", 2)
Output[3, 1, 1]
Example 2
Inputcohort_retention_counts(["a|2025-12-31", "b|2025-12-01"], ["a|2026-01-15", "a|2026-01-16", "b|2026-02-02", "a|2026-03-05", "b|2026-03-05"], "2025-12", 3)
Output[2, 1, 1, 2]