dataset-splitter 0.1.0__py3-none-any.whl

dataset_splitter/__init__.py

@@ -0,0 +1,16 @@
+"""dataset-splitter: leakage-safe train/validation/test splits in one call."""
+from ._io import load_table
+from .groups import detect_id_columns
+from .report import SplitReport
+from .splitter import Split, Splitter, split
+
+__version__ = "0.1.0"
+__all__ = [
+    "split",
+    "Splitter",
+    "Split",
+    "SplitReport",
+    "detect_id_columns",
+    "load_table",
+    "__version__",
+]

dataset_splitter/_io.py

@@ -0,0 +1,32 @@
+"""Loading tabular input: a DataFrame is passed through, a path is read."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+
+def load_table(data: Any) -> pd.DataFrame:
+    """Return `data` as a DataFrame; accepts a DataFrame or a path to .csv/.tsv/.parquet."""
+    if isinstance(data, pd.DataFrame):
+        return data
+    if isinstance(data, (str, os.PathLike)):
+        path = Path(data)
+        suffix = path.suffix.lower()
+        if suffix == ".csv":
+            return pd.read_csv(path)
+        if suffix == ".tsv":
+            return pd.read_csv(path, sep="\t")
+        if suffix in (".parquet", ".pq"):
+            try:
+                return pd.read_parquet(path)
+            except ImportError as exc:
+                raise ImportError(
+                    "reading parquet needs pyarrow: pip install 'dataset-splitter[parquet]'"
+                ) from exc
+        raise ValueError(f"unsupported file type {suffix!r}; expected .csv, .tsv or .parquet")
+    raise TypeError(
+        f"expected a pandas DataFrame or a path to a .csv/.parquet file, got {type(data).__name__}"
+    )

dataset_splitter/_util.py

@@ -0,0 +1,141 @@
+"""Small shared helpers: JSON conversion, time parsing, target kind and strata."""
+from __future__ import annotations
+
+import datetime as _dt
+import math
+import warnings
+from typing import Any, List, Tuple
+
+import numpy as np
+import pandas as pd
+from pandas.api.types import (
+    is_bool_dtype,
+    is_datetime64_any_dtype,
+    is_numeric_dtype,
+    is_object_dtype,
+    is_string_dtype,
+)
+
+
+def jsonable(obj: Any) -> Any:
+    """Recursively turn numpy/pandas scalars and containers into JSON-safe Python values."""
+    if obj is None or isinstance(obj, (bool, int, str)):
+        return obj
+    if isinstance(obj, float):
+        return obj if math.isfinite(obj) else None
+    if isinstance(obj, np.bool_):
+        return bool(obj)
+    if isinstance(obj, np.integer):
+        return int(obj)
+    if isinstance(obj, np.floating):
+        value = float(obj)
+        return value if math.isfinite(value) else None
+    if isinstance(obj, dict):
+        return {str(k): jsonable(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple, set, frozenset, np.ndarray, pd.Index, pd.Series)):
+        return [jsonable(v) for v in list(obj)]
+    if obj is pd.NaT:
+        return None
+    if isinstance(obj, np.datetime64):
+        return None if np.isnat(obj) else pd.Timestamp(obj).isoformat()
+    if isinstance(obj, (pd.Timestamp, _dt.datetime, _dt.date)):
+        return obj.isoformat()
+    if isinstance(obj, (pd.Timedelta, _dt.timedelta, pd.Interval)):
+        return str(obj)
+    try:
+        if pd.isna(obj):
+            return None
+    except (TypeError, ValueError):
+        pass
+    return str(obj)
+
+
+def parse_time(series: pd.Series, name: Any) -> pd.Series:
+    """Return a time column as datetime64 or float64 values.
+
+    Raises ValueError when values are missing or cannot be interpreted as dates or numbers.
+    """
+    if is_datetime64_any_dtype(series):
+        parsed = series
+    elif is_bool_dtype(series):
+        raise ValueError(f"time column {name!r} is boolean; expected dates or numbers")
+    elif is_numeric_dtype(series):
+        parsed = pd.to_numeric(series, errors="coerce").astype("float64")
+    else:
+        parsed = _to_datetime(series, name)
+    n_missing = int(parsed.isna().sum())
+    if n_missing:
+        raise ValueError(
+            f"time column {name!r} has {n_missing} missing value(s); fill or drop them before splitting"
+        )
+    if is_datetime64_any_dtype(parsed) and getattr(parsed.dt, "tz", None) is not None:
+        parsed = parsed.dt.tz_convert("UTC").dt.tz_localize(None)
+    return parsed
+
+
+def _to_datetime(series: pd.Series, name: Any) -> pd.Series:
+    values = series.astype(object) if not is_object_dtype(series) else series
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            return pd.to_datetime(values, errors="raise")
+        except (ValueError, TypeError, OverflowError):
+            pass
+        try:  # pandas >= 2.0 can parse per-element formats
+            return pd.to_datetime(values, errors="raise", format="mixed")
+        except (ValueError, TypeError, OverflowError) as exc:
+            raise ValueError(
+                f"time column {name!r} could not be parsed as dates or numbers"
+            ) from exc
+
+
+def time_key(parsed: pd.Series) -> np.ndarray:
+    """A sortable numeric array for a column returned by parse_time()."""
+    if is_datetime64_any_dtype(parsed):
+        return parsed.to_numpy(dtype="datetime64[ns]").astype("int64")
+    return parsed.to_numpy(dtype="float64")
+
+
+def target_kind(y: pd.Series, max_categories: int = 20) -> str:
+    """'categorical' for non-numeric or low-cardinality targets, otherwise 'numeric'."""
+    if (
+        isinstance(y.dtype, pd.CategoricalDtype)
+        or is_bool_dtype(y)
+        or is_object_dtype(y)
+        or is_string_dtype(y)
+    ):
+        return "categorical"
+    if is_datetime64_any_dtype(y):
+        return "numeric"
+    if is_numeric_dtype(y):
+        return "categorical" if y.nunique(dropna=True) <= max_categories else "numeric"
+    return "categorical"
+
+
+def strata(y: pd.Series, kind: str, n_bins: int) -> Tuple[np.ndarray, List[str]]:
+    """Integer stratum code per row plus the label of each code.
+
+    Categorical targets use their classes; numeric targets use quantile bins.
+    Missing target values become a stratum of their own, labelled '<missing>'.
+    """
+    if kind == "categorical":
+        codes, uniques = pd.factorize(y)
+        labels = [str(u) for u in uniques]
+    else:
+        categories = None
+        try:
+            binned = pd.qcut(y, q=max(1, int(n_bins)), duplicates="drop")
+            categories = list(binned.cat.categories)
+        except (ValueError, TypeError):
+            binned = None
+        if binned is None or not categories:
+            codes = np.zeros(len(y), dtype=np.int64)
+            labels = ["all"]
+        else:
+            codes = binned.cat.codes.to_numpy()
+            labels = [str(iv) for iv in categories]
+    codes = np.asarray(codes, dtype=np.int64)
+    if (codes < 0).any():
+        codes = np.where(codes < 0, len(labels), codes)
+        labels = labels + ["<missing>"]
+    return codes, labels

dataset_splitter/cli.py

@@ -0,0 +1,104 @@
+"""Command line entry point: ``dataset-splitter data.csv --target y --group customer_id``."""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from typing import List, Optional, Union
+
+from . import __version__
+from .splitter import Splitter
+
+
+def _size(text: str) -> Union[int, float]:
+    """'0.2' -> fraction, '100' -> absolute row count."""
+    if "." in text or "e" in text.lower():
+        return float(text)
+    return int(text)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="dataset-splitter",
+        description=(
+            "Leakage-safe train/validation/test splits: stratified, grouped, time-aware, "
+            "and checked. Prints the split report; exit status 1 when the report is not ok."
+        ),
+    )
+    parser.add_argument("path", help="input table (.csv, .tsv or .parquet)")
+    parser.add_argument(
+        "--target", metavar="COL", help="column to stratify on and report class balance for"
+    )
+    parser.add_argument(
+        "--group",
+        metavar="COL",
+        nargs="+",
+        help="column(s) whose rows must stay on one side, or 'auto' to detect id-like columns",
+    )
+    parser.add_argument(
+        "--time", metavar="COL", help="chronological split on this column (oldest train, newest test)"
+    )
+    parser.add_argument(
+        "--test-size", type=_size, default=0.2, help="fraction or row count (default 0.2)"
+    )
+    parser.add_argument(
+        "--val-size", type=_size, default=0.1, help="fraction or row count (default 0.1)"
+    )
+    parser.add_argument("--random-state", type=int, default=0, help="seed for shuffling (default 0)")
+    parser.add_argument(
+        "--no-dedupe", action="store_true", help="do not keep exact duplicate rows on the same side"
+    )
+    parser.add_argument("--json", action="store_true", help="print the report as JSON")
+    parser.add_argument(
+        "--output", metavar="DIR", help="write train/val/test files and report.json here"
+    )
+    parser.add_argument(
+        "--format", choices=("csv", "parquet"), default="csv", help="output file format"
+    )
+    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    return parser
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    """Parse arguments, split the table, print the report. Returns the process exit status."""
+    for stream in (sys.stdout, sys.stderr):
+        if hasattr(stream, "reconfigure"):
+            try:
+                stream.reconfigure(encoding="utf-8", errors="replace")
+            except (ValueError, OSError):  # a stream that cannot be reconfigured
+                pass
+    args = build_parser().parse_args(argv)
+    group = None
+    if args.group:
+        if args.group == ["auto"]:
+            group = "auto"
+        elif len(args.group) == 1:
+            group = args.group[0]
+        else:
+            group = list(args.group)
+    try:
+        splitter = Splitter(
+            target=args.target,
+            group=group,
+            time=args.time,
+            test_size=args.test_size,
+            val_size=args.val_size,
+            random_state=args.random_state,
+            dedupe=not args.no_dedupe,
+        )
+        result = splitter.split(args.path)
+        report = result.report()
+        if args.output:
+            result.save(args.output, format=args.format)
+    except (ValueError, TypeError, ImportError, FileNotFoundError) as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+    if args.json:
+        print(json.dumps(report.to_dict(), indent=2, ensure_ascii=False))
+    else:
+        print(report.summary())
+    return 0 if report.ok else 1
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())

dataset_splitter/groups.py

@@ -0,0 +1,169 @@
+"""Row grouping: id-like column detection, composite keys, linked ids, exact duplicates.
+
+Every function here returns a compact integer id per row (0..k-1). Rows that share an id
+must land on the same side of a split.
+"""
+from __future__ import annotations
+
+import re
+from typing import Any, Iterable, List, Optional, Sequence
+
+import numpy as np
+import pandas as pd
+from pandas.api.types import is_float_dtype
+
+_ID_SUFFIXES = {"id", "ids", "uuid", "guid"}
+_ID_ANYWHERE = {"uuid", "guid"}
+_ENTITY_WORDS = {
+    "customer", "user", "session", "subject", "patient", "account", "device", "household",
+    "client", "member", "visitor", "participant", "person", "student", "employee", "player",
+    "entity", "case", "tenant", "organisation", "organization", "org", "company", "merchant",
+    "driver", "vehicle", "site", "store", "clinic", "hospital", "school", "family",
+}
+_NUMBER_WORDS = {"no", "num", "number", "nbr", "code", "key", "ref"}
+_COMPOUNDS = {w + "id" for w in _ENTITY_WORDS} | {
+    "orderid", "itemid", "productid", "docid", "fileid", "groupid", "recordid", "txid",
+    "transactionid", "eventid", "requestid", "traceid",
+}
+_SPLIT_RE = re.compile(r"[^0-9a-zA-Z]+")
+_CAMEL_RE = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")
+
+
+def _tokens(name: Any) -> List[str]:
+    text = _CAMEL_RE.sub("_", str(name))
+    return [t for t in _SPLIT_RE.split(text.lower()) if t]
+
+
+def is_id_like(name: Any) -> bool:
+    """True when a column name looks like an entity identifier (customer_id, userId, uuid...)."""
+    tokens = _tokens(name)
+    if not tokens:
+        return False
+    if tokens[-1] in _ID_SUFFIXES:
+        return True
+    if any(t in _ID_ANYWHERE for t in tokens):
+        return True
+    if "".join(tokens) in _COMPOUNDS:
+        return True
+    if len(tokens) == 1 and tokens[0] in _ENTITY_WORDS:
+        return True
+    if tokens[-1] in _NUMBER_WORDS and any(t in _ENTITY_WORDS for t in tokens[:-1]):
+        return True
+    return False
+
+
+def detect_id_columns(df: pd.DataFrame, exclude: Optional[Iterable[Any]] = None) -> List[Any]:
+    """Columns that look like ids by name and whose values repeat (so grouping matters).
+
+    Constant columns and columns unique per row are skipped, as are `exclude`d columns.
+    """
+    skip = set(exclude or ())
+    n = len(df)
+    found: List[Any] = []
+    for col in df.columns:
+        if col in skip or not is_id_like(col):
+            continue
+        series = df[col]
+        if isinstance(series, pd.DataFrame):  # duplicated column name
+            continue
+        if is_float_dtype(series):
+            values = series.dropna()
+            if len(values) and not np.all(np.mod(values.to_numpy(dtype="float64"), 1) == 0):
+                continue  # fractional floats are measurements, not ids
+        try:
+            n_unique = int(series.nunique(dropna=True))
+        except TypeError:
+            continue
+        if n_unique < 2 or n_unique >= n:
+            continue
+        found.append(col)
+    return found
+
+
+def _codes(df: pd.DataFrame, col: Any) -> np.ndarray:
+    series = df[col]
+    if isinstance(series, pd.DataFrame):
+        raise ValueError(f"column {col!r} appears more than once in the frame")
+    try:
+        codes, _ = pd.factorize(series)
+    except TypeError as exc:
+        raise ValueError(f"group column {col!r} contains unhashable values") from exc
+    return np.asarray(codes, dtype=np.int64)
+
+
+def _isolate_missing(ids: np.ndarray, missing: np.ndarray) -> np.ndarray:
+    """Give rows with a missing key an id of their own, then renumber compactly."""
+    ids = ids.copy()
+    if missing.any():
+        start = int(ids[~missing].max()) + 1 if (~missing).any() else 0
+        ids[missing] = start + np.arange(int(missing.sum()))
+    return np.asarray(pd.factorize(ids)[0], dtype=np.int64)
+
+
+def composite_ids(df: pd.DataFrame, cols: Sequence[Any]) -> np.ndarray:
+    """One id per distinct combination of `cols`; rows missing any key stand alone."""
+    n = len(df)
+    if n == 0:
+        return np.zeros(0, dtype=np.int64)
+    code_list = [_codes(df, c) for c in cols]
+    missing = np.zeros(n, dtype=bool)
+    for codes in code_list:
+        missing |= codes < 0
+    if len(code_list) == 1:
+        ids = code_list[0]
+    else:
+        ids = np.zeros(n, dtype=np.int64)
+        for codes in code_list:
+            width = int(codes.max()) + 2
+            ids = np.asarray(pd.factorize(ids * width + (codes + 1))[0], dtype=np.int64)
+    return _isolate_missing(ids, missing)
+
+
+def link_codes(code_list: Sequence[np.ndarray]) -> np.ndarray:
+    """Connected components over rows: two rows are linked when they share a value in any array.
+
+    Negative codes mean "no value" and never link anything.
+    """
+    n = len(code_list[0])
+    labels = np.arange(n, dtype=np.int64)
+    while True:
+        changed = False
+        for codes in code_list:
+            valid = codes >= 0
+            if not valid.any():
+                continue
+            current = labels[valid]
+            lowest = (
+                pd.Series(current).groupby(codes[valid], sort=False).transform("min").to_numpy()
+            )
+            if np.any(lowest != current):
+                labels[valid] = lowest
+                changed = True
+        if not changed:
+            break
+    return np.asarray(pd.factorize(labels)[0], dtype=np.int64)
+
+
+def linked_ids(df: pd.DataFrame, cols: Sequence[Any]) -> np.ndarray:
+    """One id per connected component: rows sharing a value in any of `cols` stay together."""
+    if len(df) == 0:
+        return np.zeros(0, dtype=np.int64)
+    return link_codes([_codes(df, c) for c in cols])
+
+
+def duplicate_ids(df: pd.DataFrame) -> np.ndarray:
+    """One id per distinct row (all columns compared); exact duplicates share an id."""
+    n = len(df)
+    if n == 0 or df.shape[1] == 0:
+        return np.arange(n, dtype=np.int64)
+    ids = np.zeros(n, dtype=np.int64)
+    for i in range(df.shape[1]):
+        column = df.iloc[:, i]
+        try:
+            codes = pd.factorize(column)[0]
+        except TypeError:  # unhashable cells (lists, dicts): compare their text form
+            codes = pd.factorize(column.astype(str))[0]
+        codes = np.asarray(codes, dtype=np.int64) + 1  # missing (-1) becomes a value of its own
+        width = int(codes.max()) + 1
+        ids = np.asarray(pd.factorize(ids * width + codes)[0], dtype=np.int64)
+    return ids