Transform nested dicts with pandas apply/lambda
Company: Pinterest
Role: Data Scientist
Category: Data Manipulation (SQL/Python)
Difficulty: medium
Interview Round: Onsite
Given a pandas DataFrame df with columns: user_id (int), ts (datetime64[ns]), events (list of dicts), attrs (dict). Example rows (conceptual):
user_id=1, ts=2025-08-08 09:12:00, events=[{"type":"click","ts":"2025-08-08T09:12:00"},{"type":"view","ts":"2025-08-08T09:12:05"}], attrs={"version":"1.2.0","flags":{"beta":true,"dark":false}}
user_id=1, ts=2025-08-09 15:00:00, events=[{"type":"purchase","ts":"2025-08-09T15:00:00","amount":20.0}], attrs={"version":"1.2.0","flags":{"beta":true,"dark":false}}
user_id=2, ts=2025-08-04 12:30:00, events=[{"type":"view","ts":"2025-08-04T12:30:00"}], attrs={"version":"1.1.0","flags":{"beta":false,"dark":true}}
Tasks (write Python/pandas):
1) Explode events into a long table with one row per event: columns [user_id, event_type, event_ts (datetime), amount (nullable), attrs_version, attrs_beta_flag, attrs_dark_flag]. Use Series.explode and apply/lambda to parse dicts; no for-loops over DataFrame rows.
2) Aggregate to a per-user features table with columns: click_count, view_count, purchase_count, last_event_ts, total_purchase_amount, version_most_recent, beta_flag_most_recent, dark_flag_most_recent. Use groupby with named aggregations; where multiple versions exist, take the attrs associated with the most recent event.
3) From the features table, compute a conversion_rate by user segment defined by (version_most_recent, beta_flag_most_recent, dark_flag_most_recent): purchases/users. Return a compact DataFrame with one row per segment and columns [version, beta_flag, dark_flag, users, purchasers, conversion_rate].
4) Implement a robust helper that safely extracts nested keys from attrs using a lambda and dictionary iteration (handle missing keys and None). Explain in comments why your approach avoids SettingWithCopy pitfalls and preserves vectorization.
Overview: This question evaluates proficiency in pandas-based data manipulation, including flattening nested list and dict structures, parsing nested attributes, computing per-user aggregates and segment metrics, and implementing robust, vectorized extraction logic that handles missing keys and avoids SettingWithCopy pitfalls.
Read the full Pinterest Data Scientist interview experience this question came from
Explode JSON event arrays into a flat event table
You are given the table `user_sessions`, which stores **one row per user session**. Each row carries a JSONB array of events (`events`) and a JSONB attributes object (`attrs`).
Write a single PostgreSQL `SELECT` that **explodes the `events` array so the result has one row per individual event**, joined with the session's denormalized attributes.
For each event, produce these output columns:
- `user_id` — the session's user id.
- `event_type` — the event's `type` field (text).
- `event_ts` — the event's `ts` field, cast to `timestamp`.
- `amount` — the event's `amount` field as `numeric`, or `NULL` when the event has no `amount` key.
- `attrs_version` — `attrs.version` (text).
- `attrs_beta_flag` — `attrs.flags.beta` cast to `boolean`.
- `attrs_dark_flag` — `attrs.flags.dark` cast to `boolean`.
Use set-based SQL (a `LATERAL` join over `jsonb_array_elements`) to unnest the array — **do not** use loops or procedural code.
Sort the output by `user_id` ascending, then by `session_ts` ascending, then by the event's position within its array (first event in the array first).
Tables
user_sessions(session_id INT, user_id INT, session_ts TIMESTAMP, events JSONB, attrs JSONB)
Hints
- Use `jsonb_array_elements(events)` in a `CROSS JOIN LATERAL` to emit one row per element of the array.
- `->>` returns NULL when a key is absent — so `(ev->>'amount')::numeric` is automatically NULL for events with no amount.
Aggregate per-user event features from exploded events
You are given a single table, **`user_sessions`**, where each row is one browsing session. Each session stores its events as a JSON array in the `events` column (each event is an object with a `type` of `'click'`, `'view'`, or `'purchase'`, a `ts` timestamp string, and — for purchases — an `amount`), plus session-level metadata in the `attrs` column (an object with a `version` string and a nested `flags` object holding boolean `beta` and `dark` flags).
Write **one** PostgreSQL query that first explodes the `events` array into one row per individual event, then aggregates to produce a per-user feature table with exactly **one row per `user_id`**. The result must contain these columns:
- `user_id`
- `click_count` — number of `'click'` events for that user
- `view_count` — number of `'view'` events for that user
- `purchase_count` — number of `'purchase'` events for that user
- `last_event_ts` — timestamp of the user's most recent event
- `total_purchase_amount` — sum of `amount` over the user's purchase events (treat a missing/NULL amount as 0)
- `version_most_recent` — the `version` from the `attrs` of the session that contains the user's most recent event
- `beta_flag_most_recent` — the `flags.beta` boolean from that same most-recent session's `attrs`
- `dark_flag_most_recent` — the `flags.dark` boolean from that same most-recent session's `attrs`
When a user has events across multiple sessions/versions, the `version_most_recent` and the two flag columns must come from the session attributes attached to the single most recent event (by event timestamp). Use JSON unnesting, `GROUP BY`, and window functions as needed. Order the output by `user_id` ascending.
Tables
user_sessions(session_id INT, user_id INT, session_ts TIMESTAMP, events JSONB, attrs JSONB)
Hints
- Use `CROSS JOIN LATERAL jsonb_array_elements(events)` to turn each session's JSON event array into one row per event, carrying the session's `attrs` along.
- `ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY event_ts DESC)` marks each user's most recent event; combine the aggregate with `FILTER (WHERE rn = 1)` to read that event's session attributes.
Compute conversion rate by user segment from feature table
Suppose you have a precomputed user_features table that contains one row per user with aggregated metrics. Using this table, compute a conversion_rate by user segment defined by the tuple (version_most_recent, beta_flag_most_recent, dark_flag_most_recent). For each segment, return: version, beta_flag, dark_flag, users (number of users in the segment), purchasers (number of users in the segment with purchase_count > 0), and conversion_rate defined as purchasers divided by users as a decimal. Return one row per segment.
Tables
user_features(user_id INT, click_count INT, view_count INT, purchase_count INT, last_event_ts TIMESTAMP, total_purchase_amount DECIMAL(10,2), version_most_recent VARCHAR(20), beta_flag_most_recent BOOLEAN, dark_flag_most_recent BOOLEAN)
Hints
- Group by the three segment-defining columns (version_most_recent, beta_flag_most_recent, dark_flag_most_recent).
- Use COUNT(*) FILTER (WHERE purchase_count > 0) to count purchasers and divide by COUNT(*) as a decimal to get conversion_rate.
Safely extract nested JSON attributes with robust SQL
Using the user_sessions table, write a SELECT that extracts attrs_version, attrs_beta_flag, and attrs_dark_flag from the attrs JSONB column into separate columns. Your query should be robust to missing keys or NULL values in attrs: if attrs.version is missing, return 'unknown'; if attrs.flags.beta or attrs.flags.dark are missing or NULL, default to false. Use JSON operators and COALESCE so the query does not fail when keys are absent, and add brief SQL comments explaining how this approach safely handles missing keys in a set-based way.
Tables
user_sessions(session_id INT, user_id INT, session_ts TIMESTAMP, events JSONB, attrs JSONB)
Hints
- Use the -> and ->> operators to navigate attrs and attrs->'flags', then wrap the extractions in COALESCE to provide default values.
- Remember that casting NULL to a boolean is still NULL, so you can safely cast first and then COALESCE to false.