Impute missing values without leakage

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.

Impute Missing User Event Features in SQL Without Validation Leakage

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.

Table user_events

column	type	constraints
event_id	INT	PRIMARY KEY, NOT NULL
user_id	INT	NOT NULL
event_date	DATE	NOT NULL
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	NOT NULL
is_train	SMALLINT	NOT NULL

Sample data:

age	country	device_type	event_date	event_id	income	is_active_30d	is_train	last_purchase_days_ago	session_length	user_id
25	US	iOS	2025-05-01	1	50000	1	1	3	10	101
NULL	US	iOS	2025-05-05	2	NULL	1	1	NULL	15	101
26	NULL	NULL	2025-05-25	3	52000	0	0	NULL	NULL	101
30	CA	Android	2025-05-02	4	60000	1	1	5	20	102
NULL	CA	NULL	2025-05-10	5	NULL	1	1	NULL	NULL	102
32	CA	Android	2025-05-12	6	NULL	0	0	NULL	25	102
40	US	Web	2025-04-01	7	80000	0	1	1	5	103
NULL	NULL	NULL	2025-05-20	8	NULL	1	0	NULL	NULL	103
22	US	Android	2025-05-15	9	40000	1	1	2	8	104
23	CA	iOS	2025-05-18	10	42000	1	1	NULL	12	104
NULL	NULL	NULL	2025-05-20	11	NULL	0	0	NULL	NULL	104

Expected output

One row per input event, ordered by event_id, with country/device_type/age/income/last_purchase_days_ago/session_length imputed using only training-row statistics and a 14-day-gap-bounded time-ordered forward fill.

age_imputed	country_imputed	device_type_imputed	event_date	event_id	income_imputed	is_active_30d	is_train	last_purchase_days_ago_imputed	session_length_imputed	user_id
25	US	iOS	2025-05-01	1	50000	1	1	3	10	101
25	US	iOS	2025-05-05	2	50000	1	1	3	15	101
26	US	iOS	2025-05-25	3	52000	0	0	3	12	101
30	CA	Android	2025-05-02	4	60000	1	1	5	20	102
26.5	CA	Android	2025-05-10	5	50000	1	1	5	20	102
32	CA	Android	2025-05-12	6	60000	0	0	5	25	102
40	US	Web	2025-04-01	7	80000	0	1	1	5	103
25	US	Web	2025-05-20	8	50000	1	0	3	12	103
22	US	Android	2025-05-15	9	40000	1	1	2	8	104
23	CA	iOS	2025-05-18	10	42000	1	1	2	12	104
25	US	iOS	2025-05-20	11	50000	0	0	2	12	104

Hints

1. Postgres has no DATEDIFF: subtracting two DATE columns (a - b) gives the integer day gap, so compare LAG-based gaps with > 14 directly.
2. 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.
3. Filter every median/mode CTE to is_train = 1 BEFORE aggregating, then LEFT/CROSS JOIN the precomputed stats back to all rows so validation rows never leak into the statistics.