Data API

uplift_bench.data.base

Abstract dataset loader.

A DatasetLoader is responsible for one job: take a path on disk and produce a validated UpliftDataset. Anything fancier (downloading, caching, schema inference) belongs in subclasses or utility modules — keep the contract small so swapping a dataset is a one-class change.

DatasetLoader

Bases: ABC

Loader contract.

Subclasses implement _raw_path, _read, and schema. The base class handles the common bits — file existence, hashing, validation.

download()

Fetch the source file; default is no-op (assumes file already there).

load()

Resolve, read, validate.

uplift_bench.data.validation

Schema-level validation for uplift datasets.

We validate at the boundary between disk and the rest of the pipeline: once a UpliftDataset is in memory we trust its invariants. This avoids re-checking inside every metric and model.

The pydantic schema is intentionally narrow — it catches the dumb mistakes (wrong dtype, missing column, T outside {0, 1}) before they corrupt downstream computations. Distributional checks (overlap, class balance) live under robustness/ where they belong.

DatasetSchema

Bases: BaseModel

Declarative schema for a binary-treatment uplift dataset.

Used by every loader (real and synthetic) to declare which columns are treatment / outcome / features. Loaders pass the raw DataFrame through validate_dataframe and get back a normalized one.

UpliftDataset dataclass

Validated, ready-to-use dataset.

df keeps the original columns; X, t, y are convenience views. We avoid copying because for Criteo (~10M rows) the memory matters.

validate_dataframe(df, schema)

Check the DataFrame matches the schema; coerce where safe.

Coercions performed: - treatment column is cast to int8 (memory; downstream code assumes int). - outcome column is cast to int8 if allowed_outcome_values is set, otherwise float64.

Failures (raised as ValueError): - missing required column, - treatment values outside allowed_treatment_values, - outcome values outside allowed_outcome_values (when set), - any null in treatment, outcome, or feature columns.

uplift_bench.data.splits

Train / val / test splitting for uplift datasets.

Two non-obvious choices documented here so future-me doesn't undo them:

1. We stratify on the joint (T, Y) instead of T alone. I tried T-only stratification first. On Criteo it produced a test fold where the conversion rate was 8% off the train fold purely by chance, which made bootstrap CIs confusingly wide. (T, Y) stratification keeps the marginal P(Y=1|T=t) stable across folds.

2. We force a fixed permutation per seed instead of using sklearn's StratifiedShuffleSplit random state. Reason: when a single seed has to reproduce results across pandas / numpy / sklearn versions (which we pin only loosely), explicit indexing is the only thing that's actually stable. Sklearn changed its shuffle algorithm twice in the 1.x series.

make_splits(dataset, train_frac=0.7, val_frac=0.15, *, seed=42)

Stratified train/val/test split for an uplift dataset.

Stratification is on the joint (treatment, outcome) when the outcome is discrete; on treatment alone otherwise. Indices are positional integers into dataset.df.

Parameters:

Name	Type	Description	Default
dataset	UpliftDataset	Validated UpliftDataset.	required
train_frac	float	Fractions in (0, 1). test_frac = 1 - train - val.	0.7
val_frac	float	Fractions in (0, 1). test_frac = 1 - train - val.	0.7
seed	int	Reproducibility.	42

Returns:

Type	Description
SplitIndices	Three disjoint, sorted arrays of positional integer indices.

uplift_bench.data.factory

Build a DatasetLoader from a name + kwargs.

Mirrors uplift_bench.models.factory. Used by the Hydra entry point.

uplift_bench.data.hillstrom

Hillstrom (MineThatData) email-marketing dataset.

Kevin Hillstrom's classic 64k-row dataset. Three treatment arms: "No E-Mail", "Mens E-Mail", "Womens E-Mail". Outcomes: visit, conversion, spend (we use visit because it has decent base rate; conversion is extremely rare and the bench would need a much bigger n).

We binarize treatment to "Womens E-Mail" vs "No E-Mail" by default — that contrast has the strongest measured ATE in published analyses, which makes it a useful sanity-check setup. The other contrasts are reachable via treatment_arm.

uplift_bench.data.criteo

Criteo Uplift v2 dataset.

The original dataset is hosted by Criteo Research: https://ailab.criteo.com/criteo-uplift-prediction-dataset/

The 12-feature CSV is ~300 MB compressed (~1.5 GB uncompressed) and has roughly 13.9M rows. We materialise the validated DataFrame as parquet on first load — pandas takes 90+ seconds to re-parse the CSV, parquet does it in 4.

For local-machine sanity we also support subsample: a fixed-seed random sample of the full set. The bench reports both "full" and "subsample" runs in results/ and labels them clearly.

uplift_bench.data.retailhero

X5 RetailHero Uplift Modeling Contest dataset.

The data is hosted on Ods.ai behind a free account login: https://ods.ai/competitions/x5-retailhero-uplift-modeling/data

That makes a programmatic download impossible without scraping a session cookie, which we won't do. The loader expects two files placed by the user:

{data_dir}/retailhero/uplift_train.csv
{data_dir}/retailhero/clients.csv

uplift_train.csv has columns: client_id, treatment_flg, target. clients.csv carries the customer features. We join, drop the id column, one-hot the small categoricals, and use median-imputation for missing numeric features (RetailHero ships some genuinely-missing values).

For tests and the smoke-config a tiny synthetic stand-in lives at data/sample/retailhero/. It has the same schema but only a few thousand rows so CI runs in seconds.

uplift_bench.data.megafon

MegaFon Uplift Competition dataset.

Source: https://ods.ai/competitions/megafon-uplift-competition/data (also requires Ods.ai login). Same shape rationale as RetailHero — manual placement, sample stand-in for tests.

The dataset has ~600k rows, ~50 anonymised numeric features, binary treatment_group, and a binary conversion target.

uplift_bench.data.download

HTTP downloader with resume + sha verification.

Used by individual loaders that can self-download (Hillstrom, Criteo). Datasets that gate behind login (RetailHero, MegaFon) raise a clear error in their loader instead of going through here.

download_file(url, dest, *, expected_sha256=None, overwrite=False, timeout=60.0)

Download url to dest. Skips if file exists and hash matches.

Parameters:

Name	Type	Description	Default
url	str	Direct HTTP(S) URL. We don't follow login redirects on purpose — if a host needs auth, the caller should fail loud, not silently.	required
dest	Path	Target path. Parent created if missing.	required
expected_sha256	str | None	If given, verifies the downloaded file. On mismatch the file is deleted and a ValueError is raised — better than silently using corrupted data.	None
overwrite	bool	Force re-download even if dest exists.	False
timeout	float	Per-request connect/read timeout, in seconds.	60.0

Returns:

Type	Description
Path	dest itself, for chaining.

fetch(name, data_dir)

CLI dispatch: download one (or all) auto-downloadable dataset(s).