train.build_dataset¶

def build_training_dataset(
    features_df: pd.DataFrame,
    label_spans: Sequence[LabelSpan],
    *,
    output_dir: Path,
    train_ratio: float = 0.70,
    val_ratio: float = 0.15,
    holdout_user_fraction: float = 0.0,
    bucket_seconds: int = DEFAULT_BUCKET_SECONDS,
) -> DatasetManifest:
    """Join features with labels, apply exclusions, split, and write artifacts.

    Label projection uses strict block-to-window containment rules from
    ``time_spec.md`` Section 6 (full window must fall inside a single
    block; conflicting multi-block overlaps are dropped).

    Outputs:
        ``output_dir/X.parquet`` -- feature matrix with ID columns and
        ``schema_version``.
        ``output_dir/y.parquet`` -- labels keyed by ``user_id`` and
        ``bucket_start_ts``.
        ``output_dir/splits.json`` -- train/val/test index lists and
        metadata.

    Args:
        features_df: Feature DataFrame conforming to ``FeatureSchemaV1``.
        label_spans: Label spans to project onto feature windows.
        output_dir: Directory to write artifacts into (created if needed).
        train_ratio: Fraction of each user's data for training.
        val_ratio: Fraction for validation.
        holdout_user_fraction: Fraction of users held out entirely for
            the test set (cold-start evaluation).
        bucket_seconds: Window width in seconds.

    Returns:
        A :class:`DatasetManifest` with paths and summary statistics.
    """
    features_df = features_df.copy()
    if "user_id" not in features_df.columns:
        features_df["user_id"] = "default-user"

    schema_version = _resolve_schema_version(features_df)
    feature_columns = get_feature_columns(schema_version)

    labeled = project_blocks_to_windows(
        features_df, label_spans, bucket_seconds=bucket_seconds
    )
    pre_exclusion = len(labeled)

    labeled = _exclude_short_sessions(labeled, bucket_seconds=bucket_seconds)
    labeled = _exclude_missing_critical(labeled)
    labeled = labeled.sort_values("bucket_start_ts").reset_index(drop=True)

    excluded = pre_exclusion - len(labeled)

    splits = split_by_time(
        labeled,
        train_ratio=train_ratio,
        val_ratio=val_ratio,
        holdout_user_fraction=holdout_user_fraction,
    )

    id_and_meta = [c for c in _ID_COLUMNS if c in labeled.columns] + ["schema_version"]
    seen = set(id_and_meta)
    x_cols = id_and_meta + [
        c for c in feature_columns if c in labeled.columns and c not in seen
    ]
    x_df = labeled[x_cols]

    provenance_col = "provenance"
    y_cols = [c for c in ("user_id", "bucket_start_ts") if c in labeled.columns]
    y_cols.append("label")
    if provenance_col in labeled.columns:
        y_cols.append(provenance_col)
    y_df = labeled[y_cols]

    output_dir = Path(output_dir)
    x_path = output_dir / "X.parquet"
    y_path = output_dir / "y.parquet"
    splits_path = output_dir / "splits.json"

    write_parquet(x_df, x_path)
    write_parquet(y_df, y_path)

    class_dist = labeled["label"].value_counts().to_dict()

    splits_payload: dict[str, Any] = {
        "train": splits["train"],
        "val": splits["val"],
        "test": splits["test"],
        "holdout_users": splits["holdout_users"],
        "metadata": {
            "feature_schema_version": schema_version,
            "label_schema_version": "labels_v1",
            "total_rows": len(labeled),
            "excluded_rows": excluded,
            "user_count": labeled["user_id"].nunique(),
            "class_distribution": {str(k): int(v) for k, v in class_dist.items()},
            "train_ratio": train_ratio,
            "val_ratio": val_ratio,
            "holdout_user_fraction": holdout_user_fraction,
        },
    }
    splits_path.parent.mkdir(parents=True, exist_ok=True)
    splits_path.write_text(json.dumps(splits_payload, indent=2, default=str))

    return DatasetManifest(
        x_path=str(x_path),
        y_path=str(y_path),
        splits_path=str(splits_path),
        total_rows=len(labeled),
        train_rows=len(splits["train"]),
        val_rows=len(splits["val"]),
        test_rows=len(splits["test"]),
        excluded_rows=excluded,
        holdout_users=splits["holdout_users"],
        class_distribution={str(k): int(v) for k, v in class_dist.items()},
    )