tabalyst 0.7.3__tar.gz → 0.7.4__tar.gz

{tabalyst-0.7.3/src/tabalyst.egg-info → tabalyst-0.7.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabalyst
-Version: 0.7.3
+Version: 0.7.4
 Summary: Open-source toolkit for understanding and working with unfamiliar data.
 Author: Gregory Borelli
 License-Expression: MPL-2.0

{tabalyst-0.7.3 → tabalyst-0.7.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tabalyst"
-version = "0.7.3"
+version = "0.7.4"
 description = "Open-source toolkit for understanding and working with unfamiliar data."
 readme = "README.md"
 requires-python = ">=3.11"

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst/cli/sample.py

@@ -32,6 +32,8 @@ def _print_result(result: SampleResult) -> None:
     typer.echo(f"Source rows:   {result.source_rows:,}", err=True)
     typer.echo(f"Sample rows:   {result.sample_rows:,}", err=True)
     typer.echo(f"Sample rate:   {result.sample_rate:.2f}%", err=True)
+    if result.per_value is not None:
+        typer.echo(f"Per value:     up to {result.per_value:,} rows", err=True)
     if result.strata_known is not None:
         typer.echo(
             f"Strata:        {result.strata_known:,} known, "
@@ -39,9 +41,23 @@ def _print_result(result: SampleResult) -> None:
             f"{result.strata_uncovered:,} without rows",
             err=True,
         )
+    if result.strata_capped is not None:
+        typer.echo(
+            f"Drawn from:    {result.strata_capped:,} strata with more rows than "
+            f"the quota; the other {result.strata_known - result.strata_capped:,} "
+            "are kept whole",
+            err=True,
+        )
+        if result.source_rows and result.sample_rows * 2 >= result.source_rows:
+            typer.echo(
+                "Note:          the sample holds half of the file or more: the "
+                "strata are many or the quota is large.",
+                err=True,
+            )
     if result.seed is not None and result.method in {
         SampleMethod.RANDOM,
         SampleMethod.STRATIFIED,
+        SampleMethod.STRATIFIED_BALANCED,
     }:
         typer.echo(f"Seed:          {result.seed}", err=True)
     smart = result.smart
@@ -144,18 +160,32 @@ def _reject_unused_options(
     seed: int | None,
     yes: bool,
     no_html: bool,
+    per_value: int | None,
+    rows: int | None,
+    percent: float | None,
 ) -> None:
     """Refuse an option the method does not read: none is ignored without an error."""
     smart = method is SampleMethod.SMART
     for name, given in (("--yes", yes), ("--no-html", no_html)):
         if given and not smart:
             raise typer.BadParameter(f"{name} only goes with the smart method.")
-    stratified = method is SampleMethod.STRATIFIED
+    balanced = method is SampleMethod.STRATIFIED_BALANCED
+    stratified = balanced or method is SampleMethod.STRATIFIED
     for name, given in (("--field", bool(field)), ("--max-strata", max_strata is not None)):
         if given and not stratified:
             raise typer.BadParameter(
-                f"{name} only goes with --sample-method stratified, not {method.value}."
+                f"{name} only goes with --sample-method stratified or "
+                f"stratified-balanced, not {method.value}."
             )
+    if per_value is not None and not balanced:
+        raise typer.BadParameter(
+            f"--per-value only goes with --sample-method stratified-balanced, not {method.value}."
+        )
+    if balanced and (rows is not None or percent is not None):
+        raise typer.BadParameter(
+            "--rows and --percent do not go with stratified-balanced: use --per-value "
+            "for the quota of each stratum."
+        )
     if seed is not None and method in {SampleMethod.FIRST, SampleMethod.LAST}:
         raise typer.BadParameter(
             f"--seed has no effect with {method.value}: it keeps the rows in file order."
@@ -197,10 +227,18 @@ def sample_command(
         list[str] | None,
         typer.Option(
             "--field",
-            help="Field used for stratified sampling; repeat it to stratify "
+            help="Field used for stratified sampling (stratified, stratified-balanced); repeat it to stratify "
             "by the combinations of several fields.",
         ),
     ] = None,
+    per_value: Annotated[
+        int | None,
+        typer.Option(
+            "--per-value",
+            help="With stratified-balanced, the most rows kept from each stratum "
+            "(all of a smaller one). The size of the sample follows from the data.",
+        ),
+    ] = None,
     max_strata: Annotated[
         int | None,
         typer.Option(
@@ -273,6 +311,9 @@ def sample_command(
         seed=seed,
         yes=yes,
         no_html=no_html,
+        per_value=per_value,
+        rows=rows,
+        percent=percent,
     )
     try:
         if method is SampleMethod.SMART:
@@ -305,6 +346,7 @@ def sample_command(
                 encoding=encoding,
                 force=force,
                 max_strata=max_strata,
+                per_value=per_value,
             )
     except TabalystError as exc:
         typer.echo(f"Error: {exc}", err=True)

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst/sampling.py

@@ -21,7 +21,7 @@ from tabalyst.batch import apply_default_file_mode
 from tabalyst.config import CsvConfig
 from tabalyst.csv_records import open_records, output_encoding
 from tabalyst.errors import InputError, ReportError
-from tabalyst.scanner.stratify import allocate_strata
+from tabalyst.scanner.stratify import allocate_strata_with_minimum
 
 # The most strata a stratified sample keeps at once (STRAT-08): one counter and
 # one reservoir per stratum are held in memory while the file streams.
@@ -35,9 +35,13 @@ class SampleMethod(StrEnum):
     LAST = "last"
     RANDOM = "random"
     STRATIFIED = "stratified"
+    STRATIFIED_BALANCED = "stratified-balanced"
     SMART = "smart"
 
 
+_STRATIFIED_METHODS = frozenset({"stratified", "stratified-balanced"})
+
+
 @dataclass(frozen=True)
 class SmartSummary:
     """What a smart sample covers of the catalog of targets of its plan (design
@@ -72,11 +76,15 @@ class SampleResult:
     delimiter: str
     # The columns of a stratified sample (``field`` joins them with a comma) and
     # its balance: the strata the file has, the ones that got rows, and the
-    # ones whose rounded quota was zero (STRAT-02). ``None`` for the other methods.
+    # ones without a row (STRAT-02). ``None`` for the other methods.
     fields: tuple[str, ...] = ()
     strata_known: int | None = None
     strata_selected: int | None = None
     strata_uncovered: int | None = None
+    # ``stratified-balanced`` only: the quota asked per stratum, and the strata
+    # that have more rows than it (so that were drawn from, not kept whole).
+    per_value: int | None = None
+    strata_capped: int | None = None
     # A smart sample writes a report beside its CSV file and sums it up here.
     report: Path | None = None
     page: Path | None = None
@@ -107,6 +115,26 @@ def _validate_paths(source: Path, output: Path, *, force: bool) -> None:
         raise ReportError(f"Output path is not a file: {output}")
 
 
+def _field_indexes(headers: Sequence[str], fields: Sequence[str]) -> tuple[int, ...]:
+    indexes: list[int] = []
+    for field in fields:
+        matches = [index for index, name in enumerate(headers) if name == field]
+        if not matches:
+            raise InputError(f"Stratification field does not exist: {field}")
+        if len(matches) > 1:
+            raise InputError(f"Stratification field is ambiguous: {field}")
+        indexes.append(matches[0])
+    return tuple(indexes)
+
+
+def _too_many_strata(fields: Sequence[str], max_strata: int, record: int) -> InputError:
+    return InputError(
+        f"Stratification by {', '.join(fields)} found more than "
+        f"{max_strata:,} strata at data record {record}. Use "
+        "fewer or less detailed fields, or raise --max-strata."
+    )
+
+
 def _scan(
     source: Path,
     config: CsvConfig,
@@ -116,14 +144,7 @@ def _scan(
     strata: Counter[Hashable] = Counter()
     with open_records(source, config) as records:
         headers = records.header
-        field_indexes: list[int] = []
-        for field in fields:
-            matches = [index for index, name in enumerate(headers) if name == field]
-            if not matches:
-                raise InputError(f"Stratification field does not exist: {field}")
-            if len(matches) > 1:
-                raise InputError(f"Stratification field is ambiguous: {field}")
-            field_indexes.append(matches[0])
+        field_indexes = list(_field_indexes(headers, fields))
 
         row_count = 0
         for row_count, row in enumerate(records, start=1):
@@ -131,11 +152,7 @@ def _scan(
                 key = _stratum_key(row, field_indexes)
                 strata[key] += 1
                 if len(strata) > max_strata:
-                    raise InputError(
-                        f"Stratification by {', '.join(fields)} found more than "
-                        f"{max_strata:,} strata at data record {row_count}. Use "
-                        "fewer or less detailed fields, or raise --max-strata."
-                    )
+                    raise _too_many_strata(fields, max_strata, row_count)
     return _ScanResult(headers, row_count, strata, tuple(field_indexes))
 
 
@@ -174,6 +191,17 @@ def _sample_size(row_count: int, rows: int | None, percent: float | None) -> int
     return size
 
 
+def _check_every_stratum_fits(size: int, strata: int, fields: Sequence[str]) -> None:
+    """A proportional sample keeps one row of every stratum or does not run: it
+    never drops categories silently."""
+    if size < strata:
+        raise InputError(
+            f"A stratified sample of {size:,} rows cannot keep one row of each of the "
+            f"{strata:,} strata of {', '.join(fields)}: use --rows {strata:,} or more, "
+            "or fewer or less detailed fields."
+        )
+
+
 def _reservoir_add(
     reservoir: list[tuple[int, list[str]]],
     item: tuple[int, list[str]],
@@ -230,6 +258,52 @@ def _select_rows(
     return [row for _, row in selected]
 
 
+@dataclass(frozen=True)
+class _BalancedSelection:
+    headers: list[str]
+    row_count: int
+    rows: list[list[str]]
+    strata: int
+    capped: int
+
+
+def _select_balanced(
+    source: Path,
+    config: CsvConfig,
+    fields: Sequence[str],
+    per_value: int,
+    seed: int | None,
+    max_strata: int,
+) -> _BalancedSelection:
+    """Draw up to ``per_value`` rows of each stratum in one reading.
+
+    One reservoir of ``per_value`` rows per stratum, filled independently: the
+    places a small stratum leaves unused are never given to another one. The
+    rows come back in file order.
+    """
+    rng = random.Random(seed)
+    reservoirs: dict[Hashable, list[tuple[int, list[str]]]] = {}
+    seen_by_stratum: Counter[Hashable] = Counter()
+    row_count = 0
+    with open_records(source, config) as reader:
+        headers = reader.header
+        indexes = _field_indexes(headers, fields)
+        for index, row in enumerate(reader):
+            row_count = index + 1
+            key = _stratum_key(row, indexes)
+            if key not in reservoirs:
+                if len(reservoirs) >= max_strata:
+                    raise _too_many_strata(fields, max_strata, row_count)
+                reservoirs[key] = []
+            seen_by_stratum[key] += 1
+            _reservoir_add(reservoirs[key], (index, row), seen_by_stratum[key], per_value, rng)
+    chosen = sorted(item for items in reservoirs.values() for item in items)
+    capped = sum(1 for count in seen_by_stratum.values() if count > per_value)
+    return _BalancedSelection(
+        headers, row_count, [row for _, row in chosen], len(reservoirs), capped
+    )
+
+
 def write_sample_csv(
     output: Path,
     header: Sequence[str],
@@ -291,6 +365,7 @@ def validate_sample_options(
     percent: float | None,
     field: str | Sequence[str] | None,
     max_strata: int = DEFAULT_MAX_STRATA,
+    per_value: int | None = None,
 ) -> SampleMethod:
     """Validate options shared by single-file and batch adapters."""
 
@@ -300,12 +375,28 @@ def validate_sample_options(
         choices = ", ".join(item.value for item in SampleMethod)
         raise InputError(f"Unknown sample method {method!r}; choose from {choices}.") from exc
     fields = stratification_fields(field)
-    if selected_method is SampleMethod.STRATIFIED and not fields:
-        raise InputError("--field is required with stratified sampling.")
-    if selected_method is not SampleMethod.STRATIFIED and fields:
-        raise InputError("--field can only be used with stratified sampling.")
+    stratified = selected_method in _STRATIFIED_METHODS
+    if stratified and not fields:
+        raise InputError(f"--field is required with the {selected_method.value} method.")
+    if not stratified and fields:
+        raise InputError(
+            "--field can only be used with the stratified and stratified-balanced methods."
+        )
     if max_strata < 1:
         raise InputError("--max-strata must be at least 1.")
+    if selected_method is SampleMethod.STRATIFIED_BALANCED:
+        if rows is not None or percent is not None:
+            raise InputError(
+                "--rows and --percent do not go with stratified-balanced: it keeps "
+                "up to --per-value rows of each stratum, so its size follows from the data."
+            )
+        if per_value is None:
+            raise InputError("--per-value is required with the stratified-balanced method.")
+        if isinstance(per_value, bool) or not isinstance(per_value, int) or per_value < 1:
+            raise InputError("--per-value must be a whole number greater than 0.")
+        return selected_method
+    if per_value is not None:
+        raise InputError("--per-value can only be used with the stratified-balanced method.")
     _validate_size_options(rows, percent)
     return selected_method
 
@@ -323,18 +414,28 @@ def sample_csv(
     encoding: str = "utf-8-sig",
     force: bool = False,
     max_strata: int = DEFAULT_MAX_STRATA,
+    per_value: int | None = None,
 ) -> SampleResult:
     """Create a bounded-memory sample of a CSV file.
 
-    A stratified sample allocates its rows to the strata in proportion to their
-    size, with no minimum of one row (a rare stratum can be absent: the result
-    counts them). ``field`` may name several columns; a stratum is then a
-    combination of their values. More than ``max_strata`` strata stop the
-    operation with an ``InputError`` rather than grow without bound.
+    A ``stratified`` sample allocates its ``rows`` to the strata in proportion
+    to their size, with at least one row of each: when ``rows`` is under the
+    number of strata it stops with an ``InputError`` instead of dropping some.
+    A ``stratified-balanced`` sample keeps up to ``per_value`` random rows of
+    each stratum (all of a smaller stratum, nothing moved to another), so its
+    size is the sum of ``min(per_value, rows of the stratum)``. ``field`` may
+    name several columns; a stratum is then a combination of their values. More
+    than ``max_strata`` strata stop the operation with an ``InputError`` rather
+    than grow without bound.
     """
 
     selected_method = validate_sample_options(
-        method, rows=rows, percent=percent, field=field, max_strata=max_strata
+        method,
+        rows=rows,
+        percent=percent,
+        field=field,
+        max_strata=max_strata,
+        per_value=per_value,
     )
     if selected_method is SampleMethod.SMART:
         raise InputError(
@@ -352,10 +453,37 @@ def sample_csv(
         Path(output) if output is not None else default_sample_output(source_path)
     )
     _validate_paths(source_path, output_path, force=force)
+    if selected_method is SampleMethod.STRATIFIED_BALANCED:
+        assert per_value is not None
+        balanced = _select_balanced(
+            source_path, config, fields, per_value, seed, max_strata
+        )
+        write_sample_csv(output_path, balanced.headers, balanced.rows, source_path, config)
+        total = balanced.row_count
+        return SampleResult(
+            source=source_path,
+            output=output_path,
+            method=selected_method,
+            source_rows=total,
+            sample_rows=len(balanced.rows),
+            sample_rate=(len(balanced.rows) / total * 100) if total else 0.0,
+            field=", ".join(fields),
+            seed=seed,
+            encoding=config.encoding,
+            delimiter=config.delimiter,
+            fields=fields,
+            strata_known=balanced.strata,
+            strata_selected=balanced.strata,
+            strata_uncovered=0,
+            per_value=per_value,
+            strata_capped=balanced.capped,
+        )
     scan = _scan(source_path, config, fields, max_strata)
     size = _sample_size(scan.row_count, rows, percent)
     stratified = selected_method is SampleMethod.STRATIFIED
-    allocations = allocate_strata(scan.strata, size) if stratified else None
+    if stratified:
+        _check_every_stratum_fits(size, len(scan.strata), fields)
+    allocations = allocate_strata_with_minimum(scan.strata, size) if stratified else None
     selected = _select_rows(
         source_path, config, scan, selected_method, size, seed, allocations
     )

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst/sampling_service.py

@@ -157,6 +157,7 @@ def generate_samples(
     config_path: ConfigPath | None = None,
     force: bool = False,
     max_strata: int | None = None,
+    per_value: int | None = None,
     confirmed: bool = False,
     html: bool = True,
     on_progress: ProgressCallback | None = None,
@@ -173,7 +174,12 @@ def generate_samples(
     explicit_limit = max_strata is not None
     max_strata = DEFAULT_MAX_STRATA if max_strata is None else max_strata
     selected_method = validate_sample_options(
-        method, rows=rows, percent=percent, field=field, max_strata=max_strata
+        method,
+        rows=rows,
+        percent=percent,
+        field=field,
+        max_strata=max_strata,
+        per_value=per_value,
     )
     if selected_method is SampleMethod.SMART:
         from tabalyst.smart_service import execute_smart_samples, prepare_smart_samples
@@ -201,7 +207,7 @@ def generate_samples(
     plan = build_sample_plan(
         input_specs, output=output, output_dir=output_dir, force=force
     )
-    if selected_method is SampleMethod.STRATIFIED and not explicit_limit:
+    if selected_method in {SampleMethod.STRATIFIED, SampleMethod.STRATIFIED_BALANCED} and not explicit_limit:
         _refuse_excessive_stratification(
             [job.source for job in plan.jobs],
             stratification_fields(field),
@@ -225,6 +231,7 @@ def generate_samples(
                 encoding=config.csv.encoding,
                 force=force,
                 max_strata=max_strata,
+                per_value=per_value,
             )
         except TabalystError as exc:
             failures.append(SampleFailure(job=job, error=exc))

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst/scanner/planner/render.py

@@ -91,13 +91,13 @@ def stratification_lines(advice: Stratification) -> list[str]:
         rare = (
             ""
             if column.smallest_share is None
-            else f", rarest {column.smallest_share:.2%}, every stratum kept from "
+            else f", rarest {column.smallest_share:.2%}, one row of each stratum from "
             f"{column.size_for_every_stratum:,} rows"
         )
         absent = (
             ""
             if column.uncovered_at_size is None
-            else f", {column.uncovered_at_size} without a row at {advice.sample_size:,} rows"
+            else f", {column.uncovered_at_size} strata too many for {advice.sample_size:,} rows"
         )
         lines.append(
             f"    {column.dataset} / {column.display}: "

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst/scanner/planner/stratification.py

@@ -34,11 +34,7 @@ from tabalyst.scanner.planner.stratmodels import (
     Stratification,
     StratificationLimits,
 )
-from tabalyst.scanner.stratify import (
-    allocate_strata,
-    saturating_product,
-    size_for_every_stratum,
-)
+from tabalyst.scanner.stratify import saturating_product
 
 STRATIFICATION_VERSION = 1
 
@@ -157,10 +153,11 @@ def _column(
                 counts.append(field.missing.count)
             total = sum(counts)
             smallest, largest = _round(min(counts) / total), _round(max(counts) / total)
-            size_all = size_for_every_stratum(min(counts), total)
+            size_all = len(counts)
             if sample_size is not None:
-                quotas = allocate_strata(dict(enumerate(counts)), sample_size)
-                uncovered = sum(1 for quota in quotas.values() if quota == 0)
+                # ``stratified`` keeps one row of each stratum or refuses to run,
+                # so a size under the strata is the only way to leave some out.
+                uncovered = max(0, len(counts) - sample_size)
             if smallest <= RARE_SHARE or uncovered:
                 reasons.append("rare_strata")
         else:

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst/scanner/planner/stratmodels.py

@@ -61,9 +61,9 @@ class ColumnAdvice(ScanModel):
     # frequencies (analytical domain, so a little coarser than the original).
     smallest_share: float | None = None
     largest_share: float | None = None
-    # The sample size at which the rarest stratum cannot get zero rows.
+    # The smallest size of a ``stratified`` sample: one row of each stratum.
     size_for_every_stratum: int | None = None
-    # With a sample size: the strata the classic allocation leaves without a row.
+    # With a sample size under the strata: how many it cannot hold (the sample is refused).
     uncovered_at_size: int | None = None
 
     @model_serializer(mode="wrap")

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst/scanner/stratify.py

@@ -67,6 +67,46 @@ def allocate_strata(counts: Mapping[K, int], size: int) -> dict[K, int]:
     return allocated
 
 
+def allocate_strata_with_minimum(counts: Mapping[K, int], size: int) -> dict[K, int]:
+    """Proportional allocation of ``size`` rows with at least one row per stratum.
+
+    A stratum whose exact quota is under one row is held at one, and the rows
+    that costs are taken from the others in proportion: this repeats until no
+    free stratum falls under one. The rest is shared by largest remainder (ties:
+    the key as text). Only integers are compared, so the result is exact and
+    identical on every machine. No stratum gets more rows than it has, and the
+    total is ``size``.
+
+    Raises ``ValueError`` when ``size`` is under the number of strata (the
+    minimum cannot be kept) or over the rows counted (nothing to draw).
+    """
+    total = sum(counts.values())
+    if size < len(counts):
+        raise ValueError("the size is under the number of strata")
+    if size > total:
+        raise ValueError("the size is over the rows counted")
+    allocated: dict[K, int] = {}
+    free = dict(counts)
+    budget = size
+    while free:
+        weight = sum(free.values())
+        held = [value for value, count in free.items() if count * budget < weight]
+        if not held:
+            break
+        for value in held:
+            allocated[value] = 1
+            budget -= 1
+            del free[value]
+    weight = sum(free.values())
+    remainders: dict[K, int] = {}
+    for value, count in free.items():
+        allocated[value], remainders[value] = divmod(count * budget, weight)
+    left = size - sum(allocated.values())
+    for value in sorted(remainders, key=lambda item: (-remainders[item], str(item)))[:left]:
+        allocated[value] += 1
+    return allocated
+
+
 def size_for_every_stratum(smallest: int, total: int) -> int:
     """The smallest sample size at which the exact quota of the rarest stratum
     reaches one row, so that its allocation cannot be zero.

{tabalyst-0.7.3 → tabalyst-0.7.4/src/tabalyst.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabalyst
-Version: 0.7.3
+Version: 0.7.4
 Summary: Open-source toolkit for understanding and working with unfamiliar data.
 Author: Gregory Borelli
 License-Expression: MPL-2.0

{tabalyst-0.7.3 → tabalyst-0.7.4}/src/tabalyst.egg-info/SOURCES.txt

@@ -525,6 +525,7 @@ tests/test_public_api.py
 tests/test_report_contrast.py
 tests/test_reporting.py
 tests/test_sampling.py
+tests/test_sampling_stratified_balanced.py
 tests/test_scan_command.py
 tests/test_scanner_amounts.py
 tests/test_scanner_batches.py