Impute missing values without leakage
Company: Capital One
Role: Data Scientist
Category: Data Manipulation (SQL/Python)
Difficulty: medium
Interview Round: HR Screen
Given a DataFrame df with columns: user_id, event_date (datetime), country (categorical), device_type (categorical), age (numeric), income (numeric), last_purchase_days_ago (numeric), session_length (numeric), is_active_30d (binary label). Implement code to impute missing values for model training with strict no‑leakage. Requirements: 1) Split into train/validation indices; all statistics/models for imputation must be fit on train only and then applied to validation. 2) Numeric: age → median within country (train‑only medians); income → train a ridge regression imputer on train rows using predictors [age_imputed, country, device_type, last_purchase_days_ago, session_length] (one‑hot encoded), then predict income for both train/validation; do not use the label. 3) Time‑ordered within user: for last_purchase_days_ago and session_length, sort by event_date per user_id and forward‑fill gaps up to 14 days; if the gap between consecutive event_date exceeds 14 days, do not propagate; after sequence fills, fill remaining NaNs with the global train median for that feature. 4) Categoricals: device_type and country → per‑user mode; break ties with the global train mode. 5) Deliver: functions fit_imputers(df, train_idx) and transform_impute(df, imputers, idx) where imputers holds all train‑fit objects/statistics; include assertions that no value derived from validation data was used to compute train statistics.
Overview: This question evaluates proficiency in data imputation, strict train/validation leakage prevention, temporal per-user propagation rules, and categorical aggregation for model-ready features.
You are given the table `user_events`, a log of user-level events used to train a predictive model. Each row is one event. The column `is_train` flags whether a row belongs to the training set (`is_train = 1`) or the validation set (`is_train = 0`).
Write a **single PostgreSQL query** that returns one output row per input row with imputed feature values. **Every aggregate (median or mode) used for imputation must be computed using only training rows (`is_train = 1`)** — validation rows may never influence any statistic.
### Table: `user_events`
| column | meaning |
|---|---|
| `event_id` | unique event id |
| `user_id` | user id |
| `event_date` | date of the event |
| `country` | categorical, may be NULL |
| `device_type` | categorical, may be NULL |
| `age` | numeric, may be NULL |
| `income` | numeric, may be NULL |
| `last_purchase_days_ago` | numeric, may be NULL |
| `session_length` | numeric, may be NULL |
| `is_active_30d` | binary label (0/1) |
| `is_train` | 1 = training row, 0 = validation row |
### Imputation rules
**1) Time-ordered forward fill for `last_purchase_days_ago` and `session_length`**
- For each `user_id`, order events by `event_date`.
- Split a user's events into sequences: if the gap between two consecutive `event_date` values is **strictly greater than 14 days**, start a new sequence (values must NOT propagate across such a gap).
- Within each (`user_id`, sequence), forward-fill each feature with the most recent non-NULL value **at or before** the current event.
- If a value is still NULL after the forward fill, replace it with the **global median of the forward-filled values** of that feature over training rows only.
**2) `age`**
- `age_imputed = COALESCE(age, median age for the row's country, global median age)`, where the per-country median uses training rows with non-NULL `age` and `country`, and the global median uses all training rows with non-NULL `age`.
**3) `income`**
- `income_imputed = COALESCE(income, median income for the row's (country, device_type), global median income)`, where the per-(country, device_type) median uses training rows with non-NULL `income`, and the global median uses all training rows with non-NULL `income`.
**4) `country` and `device_type`**
- Using training rows only, compute each user's per-user mode (most frequent non-NULL value) of `country` and of `device_type`, and the global mode of each over all training rows.
- Tie-break a per-user mode by preferring the value equal to the global mode if it is among the tied values; otherwise any tied value is acceptable.
- `country_imputed = COALESCE(country, per-user mode country, global mode country)`
- `device_type_imputed = COALESCE(device_type, per-user mode device_type, global mode device_type)`
### Required output
Return exactly these columns, **one row per input row, ordered by `event_id` ascending**:
`event_id`, `user_id`, `event_date`, `country_imputed`, `device_type_imputed`, `age_imputed`, `income_imputed`, `last_purchase_days_ago_imputed`, `session_length_imputed`, `is_active_30d`, `is_train`.
Tables
user_events(event_id INT, user_id INT, event_date DATE, country VARCHAR(10), device_type VARCHAR(20), age DECIMAL(5,2), income DECIMAL(10,2), last_purchase_days_ago INT, session_length INT, is_active_30d SMALLINT, is_train SMALLINT)
Hints
- Postgres has no DATEDIFF: subtracting two DATE columns (a - b) gives the integer day gap, so compare LAG-based gaps with > 14 directly.
- Postgres LAST_VALUE has no IGNORE NULLS. Emulate forward-fill by a running COUNT() of non-NULLs to form blocks, then MAX() the value within each block.