dataset-splitter 0.1.0__tar.gz

dataset_splitter-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Pranay Mahendrakar
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

dataset_splitter-0.1.0/PKG-INFO

@@ -0,0 +1,121 @@
+Metadata-Version: 2.5
+Name: dataset-splitter
+Version: 0.1.0
+Summary: Leakage-safe train/validation/test splits in one call: stratified, grouped, time-aware, and checked
+Project-URL: Homepage, https://pypi.org/project/dataset-splitter/
+Project-URL: Author, https://pypi.org/user/pranaymahendrakar/
+Author: Pranay Mahendrakar
+License-Expression: MIT
+License-File: LICENSE
+Keywords: cross-validation,data-leakage,group-split,machine-learning,pandas,stratified,time-series-split,train-test-split
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: numpy>=1.23
+Requires-Dist: pandas>=1.5
+Provides-Extra: dev
+Requires-Dist: pyarrow>=12; extra == 'dev'
+Requires-Dist: pytest>=7; extra == 'dev'
+Provides-Extra: parquet
+Requires-Dist: pyarrow>=12; extra == 'parquet'
+Description-Content-Type: text/markdown
+
+# dataset-splitter
+
+Train/validation/test splits that cannot leak: stratified, grouped, time-aware, and checked, in one call.
+
+## Install
+
+```
+pip install dataset-splitter
+```
+
+## Quickstart
+
+```python
+import pandas as pd
+from dataset_splitter import split
+
+df = pd.DataFrame({"user_id": [i // 3 for i in range(60)], "x": range(60), "label": [i % 2 for i in range(60)]})
+s = split(df, target="label", group="user_id", test_size=0.2, val_size=0.1)
+print(s.train.shape, s.val.shape, s.test.shape)
+print(s.report().summary())
+```
+
+The summary shows the sizes, the class balance of every part, and the leakage checks: no
+`user_id` on two sides, no duplicate row on two sides, rows in exactly one part, `ok: True`.
+
+## What it does
+
+- **Stratified** on `target`: classes are balanced across train/val/test. Numeric targets are
+  stratified on quantile bins. Tiny classes are pooled, and a tiny dataset falls back to a plain
+  random split, both with a warning instead of an error.
+- **Grouped** on `group`: every value of the column (or combination of columns) lands entirely on
+  one side, like scikit-learn's `GroupShuffleSplit` but with row-accurate sizes and stratification.
+  `group="auto"` detects id-like columns (`customer_id`, `userId`, `session_uuid`, `patient_no`, ...)
+  and, when several are found, links rows that share any of them.
+- **Time-aware** on `time`: a chronological split with no shuffling; the oldest rows train, the
+  newest test. Combined with `group`, groups are ordered by their first timestamp and kept whole.
+- **Duplicate-safe**: exact duplicate rows are grouped before splitting (`dedupe=True`), so the
+  same row can never sit in both train and test.
+- **Checked**: `report()` recomputes everything from the output frames: sizes, class balance,
+  group overlap (must be 0), duplicate leakage (must be 0), time ordering, and an exact partition
+  check (no row lost or duplicated). `report().ok` is the single flag to look at.
+- **Predictable**: the original index values are preserved and listed in `.indices`; row order
+  inside each part is the input order; `random_state` makes every split reproducible.
+
+## API
+
+```python
+split(df, *, target=None, group=None, time=None, test_size=0.2, val_size=0.1,
+      random_state=0, dedupe=True) -> Split
+```
+
+`df` is a DataFrame or a path to a `.csv` / `.tsv` / `.parquet` file. `test_size` and `val_size`
+are fractions (floats below 1) or absolute row counts (ints); `val_size=0` disables validation.
+
+`Splitter(target=..., group=..., time=..., test_size=..., val_size=..., random_state=...,
+dedupe=..., stratify="auto", n_bins=10, max_categories=20)` is the class underneath, for
+control over stratification (`stratify=True/False`), the number of quantile bins for numeric
+targets, and the cardinality below which a numeric target counts as categorical. `.split(df)`.
+
+`Split`
+
+- `.train`, `.val`, `.test` - DataFrames (`.val` is `None` when `val_size=0`)
+- `.indices` - `{"train": [...], "val": [...], "test": [...]}` original index values
+- `.strategy` - what was done (method, columns used, number of groups, ...); `.warnings`
+- `.report()` - a `SplitReport`
+- `.save(dir, format="csv" | "parquet")` - writes `train/val/test.<format>` and `report.json`,
+  returns the paths (parquet needs `pip install dataset-splitter[parquet]`)
+
+`SplitReport`
+
+- `.ok` - True when nothing leaks and the parts partition the input exactly
+- `.sizes`, `.fractions`, `.class_balance`, `.class_counts`, `.balance_max_deviation`,
+  `.target_stats` (numeric targets)
+- `.group_overlap`, `.duplicate_leakage`, `.time_ordering`, `.partition`, `.warnings`
+- `.summary()` - human-readable text; `.to_dict()` - JSON-safe dict
+
+`detect_id_columns(df)` returns the columns `group="auto"` would use; `load_table(path)` reads
+a `.csv` / `.tsv` / `.parquet` file.
+
+## CLI
+
+```
+dataset-splitter data.csv --target label --group customer_id
+dataset-splitter data.csv --time timestamp --test-size 0.2 --val-size 0.1 --json
+dataset-splitter data.csv --group auto --output splits/ --format parquet
+```
+
+Prints the report (`--json` for `to_dict()` as JSON), writes the parts with `--output DIR`, and
+exits with status 1 when the report is not ok, 2 on bad input. `dataset-splitter --help` lists
+every option.
+
+## License
+
+MIT

dataset_splitter-0.1.0/README.md

@@ -0,0 +1,94 @@
+# dataset-splitter
+
+Train/validation/test splits that cannot leak: stratified, grouped, time-aware, and checked, in one call.
+
+## Install
+
+```
+pip install dataset-splitter
+```
+
+## Quickstart
+
+```python
+import pandas as pd
+from dataset_splitter import split
+
+df = pd.DataFrame({"user_id": [i // 3 for i in range(60)], "x": range(60), "label": [i % 2 for i in range(60)]})
+s = split(df, target="label", group="user_id", test_size=0.2, val_size=0.1)
+print(s.train.shape, s.val.shape, s.test.shape)
+print(s.report().summary())
+```
+
+The summary shows the sizes, the class balance of every part, and the leakage checks: no
+`user_id` on two sides, no duplicate row on two sides, rows in exactly one part, `ok: True`.
+
+## What it does
+
+- **Stratified** on `target`: classes are balanced across train/val/test. Numeric targets are
+  stratified on quantile bins. Tiny classes are pooled, and a tiny dataset falls back to a plain
+  random split, both with a warning instead of an error.
+- **Grouped** on `group`: every value of the column (or combination of columns) lands entirely on
+  one side, like scikit-learn's `GroupShuffleSplit` but with row-accurate sizes and stratification.
+  `group="auto"` detects id-like columns (`customer_id`, `userId`, `session_uuid`, `patient_no`, ...)
+  and, when several are found, links rows that share any of them.
+- **Time-aware** on `time`: a chronological split with no shuffling; the oldest rows train, the
+  newest test. Combined with `group`, groups are ordered by their first timestamp and kept whole.
+- **Duplicate-safe**: exact duplicate rows are grouped before splitting (`dedupe=True`), so the
+  same row can never sit in both train and test.
+- **Checked**: `report()` recomputes everything from the output frames: sizes, class balance,
+  group overlap (must be 0), duplicate leakage (must be 0), time ordering, and an exact partition
+  check (no row lost or duplicated). `report().ok` is the single flag to look at.
+- **Predictable**: the original index values are preserved and listed in `.indices`; row order
+  inside each part is the input order; `random_state` makes every split reproducible.
+
+## API
+
+```python
+split(df, *, target=None, group=None, time=None, test_size=0.2, val_size=0.1,
+      random_state=0, dedupe=True) -> Split
+```
+
+`df` is a DataFrame or a path to a `.csv` / `.tsv` / `.parquet` file. `test_size` and `val_size`
+are fractions (floats below 1) or absolute row counts (ints); `val_size=0` disables validation.
+
+`Splitter(target=..., group=..., time=..., test_size=..., val_size=..., random_state=...,
+dedupe=..., stratify="auto", n_bins=10, max_categories=20)` is the class underneath, for
+control over stratification (`stratify=True/False`), the number of quantile bins for numeric
+targets, and the cardinality below which a numeric target counts as categorical. `.split(df)`.
+
+`Split`
+
+- `.train`, `.val`, `.test` - DataFrames (`.val` is `None` when `val_size=0`)
+- `.indices` - `{"train": [...], "val": [...], "test": [...]}` original index values
+- `.strategy` - what was done (method, columns used, number of groups, ...); `.warnings`
+- `.report()` - a `SplitReport`
+- `.save(dir, format="csv" | "parquet")` - writes `train/val/test.<format>` and `report.json`,
+  returns the paths (parquet needs `pip install dataset-splitter[parquet]`)
+
+`SplitReport`
+
+- `.ok` - True when nothing leaks and the parts partition the input exactly
+- `.sizes`, `.fractions`, `.class_balance`, `.class_counts`, `.balance_max_deviation`,
+  `.target_stats` (numeric targets)
+- `.group_overlap`, `.duplicate_leakage`, `.time_ordering`, `.partition`, `.warnings`
+- `.summary()` - human-readable text; `.to_dict()` - JSON-safe dict
+
+`detect_id_columns(df)` returns the columns `group="auto"` would use; `load_table(path)` reads
+a `.csv` / `.tsv` / `.parquet` file.
+
+## CLI
+
+```
+dataset-splitter data.csv --target label --group customer_id
+dataset-splitter data.csv --time timestamp --test-size 0.2 --val-size 0.1 --json
+dataset-splitter data.csv --group auto --output splits/ --format parquet
+```
+
+Prints the report (`--json` for `to_dict()` as JSON), writes the parts with `--output DIR`, and
+exits with status 1 when the report is not ok, 2 on bad input. `dataset-splitter --help` lists
+every option.
+
+## License
+
+MIT

dataset_splitter-0.1.0/pyproject.toml

@@ -0,0 +1,54 @@
+[build-system]
+requires = ["hatchling>=1.27"]
+build-backend = "hatchling.build"
+
+[project]
+name = "dataset-splitter"
+version = "0.1.0"
+description = "Leakage-safe train/validation/test splits in one call: stratified, grouped, time-aware, and checked"
+readme = "README.md"
+requires-python = ">=3.9"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Pranay Mahendrakar" }]
+keywords = [
+  "train-test-split",
+  "data-leakage",
+  "stratified",
+  "group-split",
+  "time-series-split",
+  "cross-validation",
+  "pandas",
+  "machine-learning",
+]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Operating System :: OS Independent",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+  "pandas>=1.5",
+  "numpy>=1.23",
+]
+
+[project.optional-dependencies]
+# heavy or niche deps go here, never in `dependencies`
+parquet = ["pyarrow>=12"]
+dev = ["pytest>=7", "pyarrow>=12"]
+
+[project.scripts]
+dataset-splitter = "dataset_splitter.cli:main"
+
+[project.urls]
+Homepage = "https://pypi.org/project/dataset-splitter/"
+Author = "https://pypi.org/user/pranaymahendrakar/"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/dataset_splitter"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

dataset_splitter-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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())