-- Make/get assignments for all customers IDs
WITH assignments AS (
SELECT
customer_id,
create_assignment(customer_id, 'ml-test') AS variant
FROM customer_test_data
),
-- Pull model v1 (our CONTROL variant)
model_v1 AS MODEL customer_churn_model VERSION v1,
-- Pull model v2 (our TREATMENT variant)
model_v2 AS MODEL customer_churn_model VERSION v2
-- If the assigned variant for the customer is 'TREATMENT',
-- use v2, otherwise use v1 (CONTROL)
SELECT
c.customer_id,
a.variant as experiment_variant,
CASE
WHEN a.variant = 'TREATMENT'
THEN model_v2!predict(age, usage)
ELSE model_v1!predict(age, usage)
END AS prediction
FROM customer_test_data c
LEFT JOIN assignments a ON c.customer_id = a.customer_id;