data-sampler 3.2.1__py3-none-any.whl

data_sampler/io.py

@@ -0,0 +1,75 @@
+"""File loading and saving for all supported tabular formats."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pandas as pd
+
+SUPPORTED_EXTENSIONS = (".csv", ".tsv", ".json", ".xlsx", ".xls", ".parquet")
+
+
+def load_file(filepath: str | Path, sheet: str | None = None) -> pd.DataFrame:
+    """Load a data file into a DataFrame.
+
+    Supports CSV, TSV, JSON, Excel (.xlsx/.xls) and Parquet. For Excel files,
+    ``sheet`` selects the sheet by name (default: first sheet).
+    """
+    ext = Path(filepath).suffix.lower()
+    readers = {
+        ".csv": lambda: pd.read_csv(filepath),
+        ".tsv": lambda: pd.read_csv(filepath, sep="\t"),
+        ".json": lambda: pd.read_json(filepath),
+        ".xlsx": lambda: pd.read_excel(filepath, sheet_name=sheet or 0),
+        ".xls": lambda: pd.read_excel(filepath, sheet_name=sheet or 0),
+        ".parquet": lambda: pd.read_parquet(filepath),
+    }
+    if ext not in readers:
+        raise ValueError(
+            f"Unsupported file type: '{ext}'. "
+            f"Supported: {', '.join(readers.keys())}"
+        )
+    return readers[ext]()
+
+
+def list_sheets(filepath: str | Path) -> list[str]:
+    """Return the sheet names of an Excel file (empty list for other formats)."""
+    ext = Path(filepath).suffix.lower()
+    if ext not in (".xlsx", ".xls"):
+        return []
+    with pd.ExcelFile(filepath) as xls:
+        return [str(s) for s in xls.sheet_names]
+
+
+def save_output(
+    df: pd.DataFrame,
+    source_path: str | Path,
+    tag: str,
+    output_folder: str | Path | None = None,
+) -> Path:
+    """Save ``df`` next to ``source_path`` (or into ``output_folder``).
+
+    The output keeps the source file's format and is named
+    ``{stem}_{tag}{ext}`` (e.g. ``data_sample_500.csv``).
+    """
+    p = Path(source_path)
+    out_name = f"{p.stem}_{tag}{p.suffix}"
+    out_dir = Path(output_folder) if output_folder else p.parent
+    out_path = out_dir / out_name
+    ext = p.suffix.lower()
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    if ext == ".csv":
+        df.to_csv(out_path, index=False)
+    elif ext == ".tsv":
+        df.to_csv(out_path, sep="\t", index=False)
+    elif ext in (".xlsx", ".xls"):
+        df.to_excel(out_path, index=False)
+    elif ext == ".json":
+        df.to_json(out_path, orient="records", indent=2)
+    elif ext == ".parquet":
+        df.to_parquet(out_path, index=False)
+    else:
+        raise ValueError(f"Unsupported output file type: '{ext}'")
+
+    return out_path

data_sampler/report.py

@@ -0,0 +1,289 @@
+"""Terminal-friendly stratification reports (returned as strings, not printed)."""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pandas as pd
+
+from .sampling import SampleResult
+from .stats import _classify, _fmt_num
+
+BAR_WIDTH = 30  # max width of distribution bars
+HIST_BAR_WIDTH = 20  # max width of the per-column comparison histogram bars
+
+
+def format_distribution(df: pd.DataFrame, col: str, label: str = "Source") -> str:
+    """Build a bar chart of a column's distribution."""
+    out: list[str] = []
+    counts = df[col].value_counts().sort_index()
+    total = counts.sum()
+    max_count = counts.max()
+    max_label_len = max(len(str(v)) for v in counts.index)
+
+    out.append(f"\n  {label} distribution for '{col}' ({counts.shape[0]} categories):")
+    out.append(f"  {'─' * (max_label_len + BAR_WIDTH + 20)}")
+
+    for value, count in counts.items():
+        pct = count / total * 100
+        bar_len = int(count / max_count * BAR_WIDTH) if max_count > 0 else 0
+        bar = "█" * bar_len + "░" * (BAR_WIDTH - bar_len)
+        out.append(f"  {str(value):>{max_label_len}}  {bar}  {count:>6} ({pct:5.1f}%)")
+
+    out.append(f"  {'─' * (max_label_len + BAR_WIDTH + 20)}")
+    return "\n".join(out)
+
+
+def column_histogram_data(
+    df_source: pd.DataFrame,
+    df_sample: pd.DataFrame,
+    bins: int = 10,
+    top: int = 8,
+) -> list[dict]:
+    """Per-column source-vs-sample distributions, aligned for comparison.
+
+    Returns one dict per column present in both frames::
+
+        {"name", "kind", "labels": [str],
+         "source_counts": [int], "sample_counts": [int],
+         "source_pct": [float], "sample_pct": [float]}
+
+    Numeric columns share bin edges (computed from the source) so the two
+    histograms line up bin-for-bin; other columns use the source's ``top``
+    most-frequent categories, looked up in the sample. Percentages are of each
+    frame's non-null count, so for categorical columns the shown bars may sum
+    to under 100 % when there are more than ``top`` categories.
+
+    Pass the *pre-anonymization* sample (``SampleResult.data``) so the numbers
+    describe how well the sample preserved the source distribution.
+    """
+    out: list[dict] = []
+    for col in df_source.columns:
+        if col not in df_sample.columns:
+            continue
+        src = df_source[col]
+        samp = df_sample[col]
+        kind = _classify(src)
+
+        if kind == "numeric":
+            s_src = pd.to_numeric(src, errors="coerce")
+            s_src = s_src[np.isfinite(s_src)]
+            s_samp = pd.to_numeric(samp, errors="coerce")
+            s_samp = s_samp[np.isfinite(s_samp)]
+            if len(s_src) == 0:
+                continue
+            edges = np.histogram_bin_edges(s_src, bins=bins)
+            source_counts = np.histogram(s_src, bins=edges)[0]
+            sample_counts = (
+                np.histogram(s_samp, bins=edges)[0]
+                if len(s_samp)
+                else np.zeros(len(source_counts), dtype=int)
+            )
+            labels = [
+                f"{_fmt_num(float(edges[i]))} – {_fmt_num(float(edges[i + 1]))}"
+                for i in range(len(source_counts))
+            ]
+            # percentages over the FINITE counts the bins were built from —
+            # ±inf values are excluded from the bars, so including them in the
+            # denominator would silently deflate every percentage
+            src_total = len(s_src) or 1
+            samp_total = len(s_samp) or 1
+        else:
+            src_nonnull = src.dropna()
+            if src_nonnull.empty:
+                continue
+            # a near-unique column (ids, emails, free text) has no meaningful
+            # top-8 distribution — skip it rather than hash-count every value
+            # to draw eight bars of count 1
+            n_unique = src_nonnull.nunique()
+            if n_unique > max(100, 0.5 * len(src_nonnull)):
+                continue
+            src_vc = src_nonnull.astype(str).value_counts()
+            top_labels = list(src_vc.head(top).index)
+            samp_vc = samp.dropna().astype(str).value_counts()
+            source_counts = np.array([int(src_vc.get(l, 0)) for l in top_labels])
+            sample_counts = np.array([int(samp_vc.get(l, 0)) for l in top_labels])
+            labels = [str(l) for l in top_labels]
+            src_total = int(src.notna().sum()) or 1
+            samp_total = int(samp.notna().sum()) or 1
+        out.append(
+            {
+                "name": str(col),
+                "kind": kind,
+                "labels": labels,
+                "source_counts": [int(c) for c in source_counts],
+                "sample_counts": [int(c) for c in sample_counts],
+                "source_pct": [c / src_total * 100 for c in source_counts],
+                "sample_pct": [c / samp_total * 100 for c in sample_counts],
+            }
+        )
+    return out
+
+
+def format_column_histograms(
+    df_source: pd.DataFrame,
+    df_sample: pd.DataFrame,
+    bins: int = 10,
+    top: int = 8,
+) -> str:
+    """Text render of :func:`column_histogram_data`: per-column source-vs-sample
+    bars, so you can eyeball how each column's distribution held up."""
+    data = column_histogram_data(df_source, df_sample, bins=bins, top=top)
+    if not data:
+        return ""
+
+    out: list[str] = []
+    out.append("╔═════════════════════════════════════════════════════════════════════╗")
+    out.append("║                    COLUMN DISTRIBUTIONS (source → sample)           ║")
+    out.append("╚═════════════════════════════════════════════════════════════════════╝")
+
+    for d in data:
+        labels = d["labels"]
+        if not labels:
+            continue
+        label_w = min(20, max(len(l) for l in labels))
+        peak = max([*d["source_pct"], *d["sample_pct"], 1e-9])
+        out.append(f"\n  {d['name']}  ({d['kind']})")
+        for label, s_pct, m_pct in zip(labels, d["source_pct"], d["sample_pct"]):
+            lbl = label if len(label) <= label_w else label[: label_w - 1] + "…"
+            s_len = int(s_pct / peak * HIST_BAR_WIDTH)
+            m_len = int(m_pct / peak * HIST_BAR_WIDTH)
+            s_bar = "█" * s_len + "░" * (HIST_BAR_WIDTH - s_len)
+            m_bar = "█" * m_len + "░" * (HIST_BAR_WIDTH - m_len)
+            out.append(
+                f"  {lbl:>{label_w}}  src {s_bar} {s_pct:5.1f}%   "
+                f"sam {m_bar} {m_pct:5.1f}%"
+            )
+    return "\n".join(out)
+
+
+def format_stratification_report(
+    df_original: pd.DataFrame, result: SampleResult
+) -> str:
+    """Build the side-by-side distribution comparison for a stratified sample.
+
+    For random / all-rows results, returns the run notes only.
+    """
+    if result.method != "stratified":
+        return "\n".join(result.notes)
+
+    df_sample = result.data
+    strat_cols = result.strat_cols
+    group_sizes = result.group_sizes
+    allocations = result.allocations
+    sample_count = result.requested
+
+    out: list[str] = []
+    out.append("╔═════════════════════════════════════════════════════════════════════╗")
+    out.append("║                        STRATIFICATION REPORT                        ║")
+    out.append("╚═════════════════════════════════════════════════════════════════════╝")
+    out.append(f"\n  Columns used: {', '.join(strat_cols)}")
+
+    if len(strat_cols) > 1:
+        out.append(
+            "\n  Method: rows are grouped by the intersection of all stratification\n"
+            "  columns simultaneously (e.g. A=x AND B=y AND C=z). Proportional\n"
+            "  allocation is then computed per intersection group. A category value\n"
+            "  that is small in every intersection group it appears in will receive\n"
+            "  0 allocation — even if it looks sizable in the per-column view below."
+        )
+    else:
+        out.append(
+            f"\n  Method: rows are grouped by '{strat_cols[0]}' and sampled proportionally."
+        )
+
+    for col in strat_cols:
+        orig_counts = df_original[col].value_counts(dropna=False).sort_index(na_position="last")
+        samp_counts = df_sample[col].value_counts(dropna=False).sort_index(na_position="last")
+        orig_total = orig_counts.sum()
+        samp_total = samp_counts.sum()
+
+        # sort with NaN last; use pandas isna to safely detect NaN keys
+        raw_values = orig_counts.index.tolist()
+        all_values = sorted((v for v in raw_values if not pd.isna(v)), key=str) + \
+                     [v for v in raw_values if pd.isna(v)]
+
+        def make_label(v):
+            s = "(missing)" if pd.isna(v) else str(v)
+            return s if len(s) <= 13 else s[:10] + "..."
+
+        all_labels = [make_label(v) for v in all_values]
+        max_label_len = 13
+        half_bar = BAR_WIDTH // 2
+
+        def get_count(counts, v):
+            if pd.isna(v):
+                nan_mask = counts.index.isna()
+                return int(counts[nan_mask].sum()) if nan_mask.any() else 0
+            return int(counts.get(v, 0))
+
+        zero_rep = [v for v in all_values if get_count(samp_counts, v) == 0]
+
+        out.append(f"\n  Column: '{col}' ({len(all_values)} categories)")
+        header = f"  {'Value':>{max_label_len}}  {'Original':^{half_bar + 10}}  {'Sample':^{half_bar + 10}}"
+        out.append(header)
+        out.append(f"  {'─' * len(header)}")
+
+        orig_max = orig_counts.max() if len(orig_counts) > 0 else 1
+        samp_max = samp_counts.max() if len(samp_counts) > 0 else 1
+
+        for value, label in zip(all_values, all_labels):
+            o_count = get_count(orig_counts, value)
+            s_count = get_count(samp_counts, value)
+            o_pct = o_count / orig_total * 100 if orig_total > 0 else 0
+            s_pct = s_count / samp_total * 100 if samp_total > 0 else 0
+
+            o_bar_len = int(o_count / orig_max * half_bar) if orig_max > 0 else 0
+            s_bar_len = int(s_count / samp_max * half_bar) if samp_max > 0 else 0
+
+            o_bar = "█" * o_bar_len + "░" * (half_bar - o_bar_len)
+            s_bar = "█" * s_bar_len + "░" * (half_bar - s_bar_len)
+
+            flag = " ← not represented" if s_count == 0 else ""
+            out.append(
+                f"  {label:>{max_label_len}}  "
+                f"{o_bar} {o_pct:5.1f}%  "
+                f"{s_bar} {s_pct:5.1f}%{flag}"
+            )
+
+        out.append(f"  {'─' * len(header)}")
+        out.append(
+            f"  {'Totals':>{max_label_len}}  {' ' * half_bar} {orig_total:>5}   "
+            f"{' ' * half_bar} {samp_total:>5}"
+        )
+
+        if zero_rep and group_sizes is not None and allocations is not None:
+            total_rows = int(group_sizes.sum())
+            # min group size needed to receive at least 1 sample
+            min_size_for_one = math.ceil(total_rows / sample_count)
+
+            out.append(
+                f"\n  Warning: {len(zero_rep)} of {len(all_values)} categories in '{col}' "
+                f"received 0 samples.\n"
+                f"  (A group needs ≥ {min_size_for_one} rows to receive ≥ 1 sample at this scale.)"
+            )
+
+            if len(strat_cols) > 1:
+                # For each zero-represented value, look up its intersection sub-groups
+                col_idx = strat_cols.index(col)
+                for value in zero_rep:
+                    if isinstance(group_sizes.index, pd.MultiIndex):
+                        mask = group_sizes.index.get_level_values(col_idx) == value
+                        sub_sizes = group_sizes[mask]
+                        sub_allocs = allocations[mask]
+                    else:
+                        sub_sizes = group_sizes[[value]] if value in group_sizes.index else pd.Series(dtype=int)
+                        sub_allocs = allocations[[value]] if value in allocations.index else pd.Series(dtype=int)
+
+                    n_subgroups = len(sub_sizes)
+                    largest = int(sub_sizes.max()) if n_subgroups > 0 else 0
+                    total_alloc = int(sub_allocs.sum())
+
+                    out.append(
+                        f"    {str(value):>{max_label_len}}: {n_subgroups} intersection group(s), "
+                        f"largest has {largest} rows (needs {min_size_for_one}), "
+                        f"total allocated = {total_alloc}"
+                    )
+
+    return "\n".join(out)

data_sampler/sampling.py

@@ -0,0 +1,193 @@
+"""Stratified and random sampling engine.
+
+Pure logic — nothing here prints. Callers (CLI, TUI) render the
+:class:`SampleResult` via :mod:`data_sampler.report`.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Iterable
+
+import numpy as np
+import pandas as pd
+
+from ._logging import get_logger
+from .stats import is_stratifiable
+
+log = get_logger(__name__)
+
+
+@dataclass
+class SampleResult:
+    """Outcome of a sampling run, with everything needed to build a report."""
+
+    data: pd.DataFrame
+    method: str  # "stratified" | "random" | "all"
+    requested: int
+    strat_cols: list[str] = field(default_factory=list)
+    group_sizes: pd.Series | None = None
+    allocations: pd.Series | None = None
+    notes: list[str] = field(default_factory=list)
+
+
+def find_stratification_columns(
+    df: pd.DataFrame,
+    sample_count: int,
+    exclude: Iterable[str] = (),
+) -> list[str]:
+    """Pick columns suitable for stratification.
+
+    Candidates are categorical / low-cardinality columns with 2–100 unique
+    values; long text and ID-like numeric columns are skipped. Columns listed
+    in ``exclude`` are never considered (the user's "skip" selections). The
+    final set is pruned so the intersection-group count fits ``sample_count``.
+    """
+    exclude = set(exclude)
+    candidates = []
+    n_rows = len(df)
+
+    for col in df.columns:
+        if col in exclude:
+            log.debug("stratification: column %r excluded by user", col)
+            continue
+        series = df[col]
+        n_unique = series.nunique()
+        if not is_stratifiable(series, n_rows, n_unique=n_unique):
+            continue
+        candidates.append((col, n_unique))
+
+    # sort by fewest categories first (easiest to represent)
+    candidates.sort(key=lambda x: x[1])
+
+    # prune: only keep columns whose combined group count fits the sample size
+    selected = []
+    combo_count = 1
+    for col, n_unique in candidates:
+        new_combo = combo_count * n_unique
+        if new_combo > sample_count:
+            break
+        combo_count = new_combo
+        selected.append(col)
+
+    log.debug("stratification columns selected: %s", selected)
+    return selected
+
+
+def stratified_sample(
+    df: pd.DataFrame,
+    count: int,
+    strat_cols: list[str],
+    random_state: int | np.random.Generator | None = None,
+) -> tuple[pd.DataFrame, pd.Series, pd.Series]:
+    """Sample ``count`` rows preserving the joint distribution of ``strat_cols``.
+
+    Returns ``(sampled_df, group_sizes, allocations)``.
+    """
+    rng = np.random.default_rng(random_state) if not isinstance(
+        random_state, np.random.Generator
+    ) else random_state
+
+    grouped = df.groupby(strat_cols, observed=True, dropna=False)
+    group_sizes = grouped.size()
+    total = group_sizes.sum()
+
+    # compute proportional allocation
+    allocations = (group_sizes / total * count).apply(math.floor)
+
+    # distribute remaining slots to the largest remainders
+    remainders = (group_sizes / total * count) - allocations
+    shortfall = count - int(allocations.sum())
+    if shortfall > 0:
+        top_positions = remainders.values.argsort()[-shortfall:]
+        for pos in top_positions:
+            allocations.iloc[pos] += 1
+
+    # sample from each group
+    samples = []
+    for i, (_, group_df) in enumerate(grouped):
+        n = int(allocations.iloc[i])
+        if n == 0:
+            continue
+        n = min(n, len(group_df))
+        samples.append(group_df.sample(n=n, random_state=rng))
+
+    result = pd.concat(samples).sample(frac=1, random_state=rng)  # shuffle
+
+    # adjust if rounding left us short/over
+    if len(result) < count:
+        remaining = df.drop(result.index)
+        extra = remaining.sample(
+            n=min(count - len(result), len(remaining)), random_state=rng
+        )
+        result = pd.concat([result, extra])
+    elif len(result) > count:
+        result = result.head(count)
+
+    return result, group_sizes, allocations
+
+
+def sample(
+    df: pd.DataFrame,
+    count: int,
+    use_random: bool = False,
+    exclude_columns: Iterable[str] = (),
+    random_state: int | np.random.Generator | None = None,
+) -> SampleResult:
+    """Create a representative sample of ``count`` rows.
+
+    Stratifies automatically on suitable columns unless ``use_random`` is
+    set. Columns in ``exclude_columns`` are never used for stratification
+    (they still appear in the output). Returns a :class:`SampleResult`.
+    """
+    log.info("sampling %d of %d rows (random=%s)", count, len(df), use_random)
+
+    if count >= len(df):
+        return SampleResult(
+            data=df,
+            method="all",
+            requested=count,
+            notes=[
+                f"Requested {count} samples but file only has {len(df)} rows. "
+                "Returning all rows."
+            ],
+        )
+
+    rng = np.random.default_rng(random_state) if not isinstance(
+        random_state, np.random.Generator
+    ) else random_state
+
+    if use_random:
+        return SampleResult(
+            data=df.sample(n=count, random_state=rng),
+            method="random",
+            requested=count,
+            notes=["Mode: pure random sampling"],
+        )
+
+    strat_cols = find_stratification_columns(df, count, exclude=exclude_columns)
+
+    if not strat_cols:
+        return SampleResult(
+            data=df.sample(n=count, random_state=rng),
+            method="random",
+            requested=count,
+            notes=[
+                "No suitable columns for stratification found. "
+                "Using pure random sampling."
+            ],
+        )
+
+    result, group_sizes, allocations = stratified_sample(
+        df, count, strat_cols, random_state=rng
+    )
+    return SampleResult(
+        data=result,
+        method="stratified",
+        requested=count,
+        strat_cols=strat_cols,
+        group_sizes=group_sizes,
+        allocations=allocations,
+        notes=[f"Stratifying on columns: {strat_cols}"],
+    )