load_dataset()

Load a bundled dataset by name.

Usage

load_dataset(
    name,
    *,
    backend=None,
)

Greenwood ships several classic survival-analysis datasets from R’s survival package, stored as gzipped CSVs. This function decompresses them on the fly and returns a DataFrame in your preferred backend.

Parameters

name: str: One of available_datasets() (e.g., "lung", "veteran", "ovarian", "pbc", "pbcseq", "colon", "mgus2", "e1684").
backend: str | None = None: "pandas", "polars", or "pyarrow". When left as None (the default), Greenwood picks a backend for you: it prefers Polars if it is installed, otherwise uses Pandas, and raises if neither is available.

Returns

DataFrame: A Polars DataFrame, a Pandas DataFrame, or a PyArrow Table with the dataset contents. Columns that are factors in R (veteran celltype, colon rx, and sex in pbc, pbcseq, and mgus2) are categoricals with R’s levels in R’s order: a Polars Enum, a Pandas Categorical, or an Arrow dictionary column. Models then take the first level as the reference. Curves and tests come in level order for the Polars and Pandas frames. An Arrow dictionary sorts by value, as Arrow itself does, so its groups come in sorted order.

Examples

Load the NCCTG lung cancer dataset as a Polars DataFrame:

import greenwood as gw

# Load the NCCTG lung cancer dataset as a Polars DataFrame
lung = gw.load_dataset("lung", backend="polars")
lung.head()

shape: (5, 10)

inst	time	status	age	sex	ph.ecog	ph.karno	pat.karno	meal.cal	wt.loss
i64	i64	i64	i64	i64	i64	i64	i64	i64	i64
3	306	2	74	1	1	90	100	1175	null
3	455	2	68	1	0	90	90	1225	15
3	1010	1	56	1	0	90	90	null	15
5	210	2	57	1	1	90	60	1150	11
1	883	2	60	1	0	100	90	null	0

The cell type in the veteran dataset keeps R’s levels, with squamous first:

# The cell type is an Enum with R's levels
veteran = gw.load_dataset("veteran", backend="polars")
veteran["celltype"].dtype

Enum(categories=['squamous', 'smallcell', 'adeno', 'large'])

See all available dataset names:

# List all available dataset names
gw.available_datasets()

['colon', 'e1684', 'lung', 'mgus2', 'ovarian', 'pbc', 'pbcseq', 'veteran']