tidaldb/tidal/src/experiment/mod.rs

420 lines
14 KiB
Rust

//! A/B experiment framework for baseline comparison.
//!
//! Provides deterministic user-to-group assignment and metric aggregation for
//! comparing personalized ranking against a non-personalized baseline feed.
//! The experiment module is read-only with respect to storage -- all metrics
//! are derived from existing `SignalLedger` and `UserSignalIndex` data.
pub mod report;
use std::time::Duration;
use crate::schema::TidalError;
// ── ExperimentGroup ──────────────────────────────────────────────────────────
/// Which arm of an A/B experiment a user is assigned to.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ExperimentGroup {
/// The personalized (non-baseline) arm.
Treatment,
/// The baseline (non-personalized) arm.
Control,
}
// ── ExperimentConfig ─────────────────────────────────────────────────────────
/// Configuration for an A/B experiment comparing two ranking profiles.
#[derive(Debug, Clone)]
pub struct ExperimentConfig {
/// Unique experiment identifier (e.g., "pg1-personalization-v1").
pub experiment_id: String,
/// Fraction of users assigned to treatment (personalized). Must be in `0.0..=1.0`.
/// Control gets the remainder.
pub treatment_fraction: f64,
/// Profile name for the treatment group (e.g., "`for_you`").
pub treatment_profile: String,
/// Profile name for the control group (e.g., "chronological").
pub control_profile: String,
/// Signals that count as "click" for CTR computation.
pub click_signals: Vec<String>,
/// Signals that count as "completion" for completion rate.
pub completion_signals: Vec<String>,
/// Observation window for return rate (e.g., 7 days).
pub return_window: Duration,
}
impl ExperimentConfig {
/// Validate the experiment configuration.
///
/// # Errors
///
/// Returns `TidalError::InvalidInput` if:
/// - `treatment_fraction` is not in `0.0..=1.0`
/// - `click_signals` is empty
/// - `completion_signals` is empty
pub fn validate(&self) -> crate::Result<()> {
if !(0.0..=1.0).contains(&self.treatment_fraction) {
return Err(TidalError::invalid_input(format!(
"treatment_fraction must be in 0.0..=1.0, got {}",
self.treatment_fraction
)));
}
if self.click_signals.is_empty() {
return Err(TidalError::invalid_input("click_signals must not be empty"));
}
if self.completion_signals.is_empty() {
return Err(TidalError::invalid_input(
"completion_signals must not be empty",
));
}
Ok(())
}
}
// ── Assignment ───────────────────────────────────────────────────────────────
/// Assign a user to an experiment group.
///
/// Uses FNV-1a hash over `experiment_id` and `user_id` for deterministic,
/// stable assignment. The same `(user_id, experiment_id)` pair always maps
/// to the same group. This is a pure function with no I/O or state.
///
/// The hash is taken modulo 10,000 to produce a bucket. If the bucket is
/// less than `treatment_fraction * 10_000`, the user is assigned to Treatment.
#[must_use]
pub fn assign_group(user_id: u64, experiment_id: &str, treatment_fraction: f64) -> ExperimentGroup {
let hash = fnv1a_hash(experiment_id.as_bytes(), user_id);
let bucket = hash % 10_000;
#[allow(clippy::cast_sign_loss)]
let threshold = (treatment_fraction * 10_000.0) as u64;
if bucket < threshold {
ExperimentGroup::Treatment
} else {
ExperimentGroup::Control
}
}
/// FNV-1a hash combining `experiment_id` bytes and `user_id`.
///
/// Deterministic across platforms: uses fixed FNV offset basis and prime,
/// processes `experiment_id` bytes followed by `user_id` little-endian bytes.
fn fnv1a_hash(experiment_bytes: &[u8], user_id: u64) -> u64 {
const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
const FNV_PRIME: u64 = 0x0100_0000_01b3;
let mut hash = FNV_OFFSET_BASIS;
for &byte in experiment_bytes {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
for &byte in &user_id.to_le_bytes() {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
hash
}
// ── Report types ─────────────────────────────────────────────────────────────
/// Per-group engagement metrics.
#[derive(Debug, Clone, PartialEq)]
pub struct GroupMetrics {
/// Click-through rate: `total_clicks / total_views`.
pub ctr: f64,
/// Completion rate: `total_completions / total_views`.
pub completion_rate: f64,
/// Return rate: fraction of users who returned within the observation window.
pub return_rate: f64,
/// Total view events across all users in this group.
pub total_views: u64,
/// Total click events across all users in this group.
pub total_clicks: u64,
/// Total completion events across all users in this group.
pub total_completions: u64,
}
impl GroupMetrics {
/// Create zeroed-out group metrics.
#[must_use]
pub const fn zero() -> Self {
Self {
ctr: 0.0,
completion_rate: 0.0,
return_rate: 0.0,
total_views: 0,
total_clicks: 0,
total_completions: 0,
}
}
}
/// Lift of treatment metrics relative to control.
///
/// Each field is `(treatment - control) / control`. When the control value
/// is zero and treatment is also zero, lift is `0.0`. When control is zero
/// and treatment is positive, lift is `f64::INFINITY`.
#[derive(Debug, Clone, PartialEq)]
pub struct MetricLift {
/// CTR lift: `(treatment_ctr - control_ctr) / control_ctr`.
pub ctr_lift: f64,
/// Completion rate lift.
pub completion_lift: f64,
/// Return rate lift.
pub return_lift: f64,
}
impl MetricLift {
/// Compute lift from treatment and control metrics.
#[must_use]
pub fn compute(treatment: &GroupMetrics, control: &GroupMetrics) -> Self {
Self {
ctr_lift: relative_lift(treatment.ctr, control.ctr),
completion_lift: relative_lift(treatment.completion_rate, control.completion_rate),
return_lift: relative_lift(treatment.return_rate, control.return_rate),
}
}
}
/// Compute relative lift: `(treatment - control) / control`.
///
/// Edge cases:
/// - Both zero: 0.0
/// - Control zero, treatment positive: `f64::INFINITY`
/// - Control zero, treatment negative: `f64::NEG_INFINITY`
fn relative_lift(treatment: f64, control: f64) -> f64 {
if control == 0.0 {
if treatment == 0.0 {
0.0
} else if treatment > 0.0 {
f64::INFINITY
} else {
f64::NEG_INFINITY
}
} else {
(treatment - control) / control
}
}
/// Full A/B experiment comparison report.
#[derive(Debug, Clone)]
pub struct ExperimentReport {
/// The experiment identifier.
pub experiment_id: String,
/// Number of users in the treatment group.
pub treatment_users: usize,
/// Number of users in the control group.
pub control_users: usize,
/// Aggregated metrics for treatment.
pub treatment_metrics: GroupMetrics,
/// Aggregated metrics for control.
pub control_metrics: GroupMetrics,
/// Lift of treatment over control.
pub lift: MetricLift,
}
// ── Tests ────────────────────────────────────────────────────────────────────
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::float_cmp)]
mod tests {
use super::*;
#[test]
fn assign_group_is_deterministic() {
let group = assign_group(42, "exp-1", 0.5);
for _ in 0..1000 {
assert_eq!(assign_group(42, "exp-1", 0.5), group);
}
}
#[test]
fn different_experiment_id_can_change_group() {
// With enough experiment IDs, at least one should differ for user 42.
let base_group = assign_group(42, "exp-a", 0.5);
let found_different = (0..100)
.map(|i| assign_group(42, &format!("exp-{i}"), 0.5))
.any(|g| g != base_group);
assert!(
found_different,
"different experiment IDs should sometimes produce different groups"
);
}
#[test]
fn treatment_fraction_1_0_assigns_all_treatment() {
for user_id in 0..10_000u64 {
assert_eq!(
assign_group(user_id, "all-treatment", 1.0),
ExperimentGroup::Treatment
);
}
}
#[test]
fn treatment_fraction_0_0_assigns_all_control() {
for user_id in 0..10_000u64 {
assert_eq!(
assign_group(user_id, "all-control", 0.0),
ExperimentGroup::Control
);
}
}
#[test]
fn treatment_fraction_0_5_produces_balanced_split() {
let treatment_count = (0..10_000u64)
.filter(|&uid| assign_group(uid, "balanced-test", 0.5) == ExperimentGroup::Treatment)
.count();
assert!(
(4750..=5250).contains(&treatment_count),
"expected ~5000 treatment users, got {treatment_count}"
);
}
#[test]
fn treatment_fraction_0_1_produces_10_percent_split() {
let treatment_count = (0..10_000u64)
.filter(|&uid| assign_group(uid, "ten-percent", 0.1) == ExperimentGroup::Treatment)
.count();
assert!(
(800..=1200).contains(&treatment_count),
"expected ~1000 treatment users, got {treatment_count}"
);
}
#[test]
fn validate_rejects_out_of_range_fraction() {
let config = ExperimentConfig {
experiment_id: "test".into(),
treatment_fraction: 1.5,
treatment_profile: "for_you".into(),
control_profile: "chronological".into(),
click_signals: vec!["click".into()],
completion_signals: vec!["complete".into()],
return_window: Duration::from_secs(7 * 86400),
};
assert!(config.validate().is_err());
}
#[test]
fn validate_rejects_negative_fraction() {
let config = ExperimentConfig {
experiment_id: "test".into(),
treatment_fraction: -0.1,
treatment_profile: "for_you".into(),
control_profile: "chronological".into(),
click_signals: vec!["click".into()],
completion_signals: vec!["complete".into()],
return_window: Duration::from_secs(7 * 86400),
};
assert!(config.validate().is_err());
}
#[test]
fn validate_rejects_empty_click_signals() {
let config = ExperimentConfig {
experiment_id: "test".into(),
treatment_fraction: 0.5,
treatment_profile: "for_you".into(),
control_profile: "chronological".into(),
click_signals: vec![],
completion_signals: vec!["complete".into()],
return_window: Duration::from_secs(7 * 86400),
};
assert!(config.validate().is_err());
}
#[test]
fn validate_rejects_empty_completion_signals() {
let config = ExperimentConfig {
experiment_id: "test".into(),
treatment_fraction: 0.5,
treatment_profile: "for_you".into(),
control_profile: "chronological".into(),
click_signals: vec!["click".into()],
completion_signals: vec![],
return_window: Duration::from_secs(7 * 86400),
};
assert!(config.validate().is_err());
}
#[test]
fn validate_accepts_valid_config() {
let config = ExperimentConfig {
experiment_id: "test".into(),
treatment_fraction: 0.5,
treatment_profile: "for_you".into(),
control_profile: "chronological".into(),
click_signals: vec!["click".into()],
completion_signals: vec!["complete".into()],
return_window: Duration::from_secs(7 * 86400),
};
assert!(config.validate().is_ok());
}
#[test]
fn validate_accepts_boundary_fractions() {
let mut config = ExperimentConfig {
experiment_id: "test".into(),
treatment_fraction: 0.0,
treatment_profile: "for_you".into(),
control_profile: "chronological".into(),
click_signals: vec!["click".into()],
completion_signals: vec!["complete".into()],
return_window: Duration::from_secs(7 * 86400),
};
assert!(config.validate().is_ok());
config.treatment_fraction = 1.0;
assert!(config.validate().is_ok());
}
#[test]
fn relative_lift_both_zero_is_zero() {
assert_eq!(relative_lift(0.0, 0.0), 0.0);
}
#[test]
fn relative_lift_control_zero_treatment_positive_is_infinity() {
assert_eq!(relative_lift(0.5, 0.0), f64::INFINITY);
}
#[test]
fn relative_lift_normal_case() {
// treatment 0.3, control 0.2 => lift = (0.3 - 0.2) / 0.2 = 0.5
let lift = relative_lift(0.3, 0.2);
assert!((lift - 0.5).abs() < 1e-10, "expected 0.5, got {lift}");
}
#[test]
fn relative_lift_treatment_worse() {
// treatment 0.1, control 0.2 => lift = -0.5
let lift = relative_lift(0.1, 0.2);
assert!((lift - (-0.5)).abs() < 1e-10, "expected -0.5, got {lift}");
}
#[test]
fn metric_lift_compute_all_fields() {
let treatment = GroupMetrics {
ctr: 0.3,
completion_rate: 0.6,
return_rate: 0.8,
total_views: 100,
total_clicks: 30,
total_completions: 60,
};
let control = GroupMetrics {
ctr: 0.2,
completion_rate: 0.4,
return_rate: 0.5,
total_views: 100,
total_clicks: 20,
total_completions: 40,
};
let lift = MetricLift::compute(&treatment, &control);
assert!((lift.ctr_lift - 0.5).abs() < 1e-10);
assert!((lift.completion_lift - 0.5).abs() < 1e-10);
assert!((lift.return_lift - 0.6).abs() < 1e-10);
}
}