mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/rowset.py

@@ -0,0 +1,283 @@
+"""The ONE cross-runtime value domain for a cell of a sealed dataset.
+
+Two callers need to turn sealed Parquet bytes into a canonical JSON payload: the hosted
+worker's table preview and the serving surface. They face the same two hard facts.
+``canonical.canonical_json_bytes`` refuses every float, so a numeric cell has to reach the
+encoder as text; and ``canonical.MAX_CANONICAL_MEMBERS`` is exactly the pipeline's
+``MAX_ROWS``, so a whole-dataset payload cannot be canonicalized at all and has to be fitted
+to a byte budget.
+
+Both problems are solved here once. A second renderer would be a second value domain, and
+two restatements of a value domain stay mutually consistent right up until one of them is
+edited -- which is the failure ``docs/CONTRACTS.md`` records for restated format tables: the
+copy keeps agreeing with itself while the real gate rots.
+
+This module is deliberately a LEAF. Its only in-package import is ``canonical``; the bounds
+it enforces arrive as arguments rather than being read from ``pipeline``. A caller that must
+prove a small call graph -- that no request path can reach an origin fetch or an unsealed
+draft -- can then start from a small named set instead of from the pipeline's whole import
+closure.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from collections.abc import Callable, Sequence
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from mostlyright.data_harness.canonical import CanonicalJSONError
+
+__all__ = [
+    "COLUMN_NAME",
+    "RowsetError",
+    "fit_canonical_payload",
+    "inspect_sealed_parquet",
+    "open_sealed_parquet",
+    "render_cell",
+    "render_filtered_window",
+    "render_rows",
+]
+
+# The one refusal code this module raises. A caller that needs a different error identity
+# translates at its own boundary rather than teaching this module about that identity.
+ROWS_INVALID = "DATASET_ROWS_INVALID"
+
+# The frozen column-name shape a sealed dataset may present to either reader.
+COLUMN_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.:-]{0,127}$")
+
+
+class RowsetError(RuntimeError):
+    """A deterministic rowset gate refused a sealed dataset read.
+
+    Shaped like ``pipeline.BuildError``: a typed refusal carrying a greppable ``code``,
+    whose ``str()`` is exactly the message, so a caller can re-raise it under its own error
+    type without the message text changing.
+    """
+
+    def __init__(self, code: str, message: str) -> None:
+        super().__init__(message)
+        self.code = code
+
+
+def render_cell(scalar: pa.Scalar) -> bool | str | None:
+    """Render one Arrow scalar into the cross-runtime value domain.
+
+    The domain is exactly ``bool | str | None``. Every number leaves as text because the
+    canonical encoder refuses non-integer JSON numbers outright, and rendering an integer as
+    text too keeps one rule for all numerics instead of a width-dependent one.
+    """
+
+    if not scalar.is_valid:
+        return None
+    value = scalar.as_py()
+    data_type = scalar.type
+    if pa.types.is_boolean(data_type):
+        if type(value) is not bool:  # pragma: no cover - Arrow type invariant
+            raise RowsetError(ROWS_INVALID, "boolean scalar is invalid")
+        return value
+    if pa.types.is_string(data_type) or pa.types.is_large_string(data_type):
+        if not isinstance(value, str):  # pragma: no cover - Arrow type invariant
+            raise RowsetError(ROWS_INVALID, "string scalar is invalid")
+        return value
+    if pa.types.is_integer(data_type):
+        return str(value)
+    if pa.types.is_floating(data_type):
+        if not isinstance(value, float) or not math.isfinite(value):
+            raise RowsetError(
+                ROWS_INVALID,
+                "floating preview values must be finite",
+            )
+        return format(value, ".17g")
+    if pa.types.is_decimal(data_type):
+        return format(value, "f")
+    if pa.types.is_date(data_type):
+        return value.isoformat()
+    if pa.types.is_timestamp(data_type):
+        rendered = value.isoformat(timespec="microseconds")
+        return rendered.removesuffix("+00:00") + "Z" if rendered.endswith("+00:00") else rendered
+    raise RowsetError(
+        ROWS_INVALID,
+        f"unsupported preview column type: {data_type}",
+    )
+
+
+def render_rows(
+    source: pa.RecordBatch | pa.Table,
+    columns: Sequence[str],
+) -> list[dict[str, bool | str | None]]:
+    """Render every row of one batch or table as an ordered mapping of rendered cells.
+
+    Columns are addressed positionally, in the order the caller read them off the sealed
+    schema, so a duplicate or reordered name in the mapping cannot silently pick a different
+    Arrow column than the one the guards inspected.
+    """
+
+    return [
+        {
+            column: render_cell(source.column(column_index)[row_index])
+            for column_index, column in enumerate(columns)
+        }
+        for row_index in range(source.num_rows)
+    ]
+
+
+def render_filtered_window(
+    parquet: pq.ParquetFile,
+    columns: Sequence[str],
+    *,
+    select: Callable[[pa.RecordBatch], pa.RecordBatch | pa.Table],
+    offset: int,
+    limit: int,
+    batch_size: int = 1_024,
+    exact_match_count: int | None = None,
+) -> tuple[int, list[dict[str, bool | str | None]]]:
+    """Retain one bounded window and count only when no exact metadata count exists.
+
+    Sealed file order remains the order: batches are visited in Parquet order and rows are
+    appended in their order inside each batch. Unfiltered serving supplies the exact Parquet
+    metadata count, so scanning stops as soon as the requested window is full. A filtered read has
+    no trustworthy metadata count and continues only to keep ``matched_row_count`` exact. Live
+    Arrow state is one fixed-size batch plus at most ``limit`` rendered rows, never the table.
+    """
+
+    if exact_match_count is not None and (
+        type(exact_match_count) is not int or exact_match_count < 0
+    ):
+        raise RowsetError(ROWS_INVALID, "exact match count is invalid")
+    matched_row_count = 0
+    rendered: list[dict[str, bool | str | None]] = []
+    for row_group in range(parquet.metadata.num_row_groups):
+        row_group_rows = parquet.metadata.row_group(row_group).num_rows
+        if exact_match_count is not None and matched_row_count + row_group_rows <= offset:
+            matched_row_count += row_group_rows
+            continue
+        for batch in parquet.iter_batches(batch_size=batch_size, row_groups=[row_group]):
+            selected = select(batch)
+            selected_count = selected.num_rows
+            start = max(offset - matched_row_count, 0)
+            remaining = limit - len(rendered)
+            if remaining > 0 and start < selected_count:
+                retained = selected.slice(start, min(remaining, selected_count - start))
+                rendered.extend(render_rows(retained, columns))
+            matched_row_count += selected_count
+            if exact_match_count is not None and len(rendered) == limit:
+                return exact_match_count, rendered
+    return exact_match_count if exact_match_count is not None else matched_row_count, rendered
+
+
+def inspect_sealed_parquet(
+    parquet: pq.ParquetFile,
+    *,
+    max_columns: int,
+    max_rows: int,
+    max_uncompressed_bytes: int,
+) -> tuple[list[str], int]:
+    """Validate one already-open sealed Parquet file entirely from its metadata."""
+
+    columns = list(parquet.schema_arrow.names)
+    total_row_count = parquet.metadata.num_rows
+    if (
+        not 1 <= len(columns) <= max_columns
+        or len(columns) != len(set(columns))
+        or any(not COLUMN_NAME.fullmatch(column) for column in columns)
+    ):
+        raise RowsetError(
+            ROWS_INVALID,
+            "preview columns violate the frozen table-preview contract",
+        )
+    if type(total_row_count) is not int or not 0 <= total_row_count <= max_rows:
+        raise RowsetError(
+            ROWS_INVALID,
+            "preview row count exceeds the candidate row budget",
+        )
+    declared_uncompressed_bytes = sum(
+        parquet.metadata.row_group(index).total_byte_size
+        for index in range(parquet.metadata.num_row_groups)
+    )
+    if (
+        type(declared_uncompressed_bytes) is not int
+        or not 0 <= declared_uncompressed_bytes <= max_uncompressed_bytes
+    ):
+        raise RowsetError(
+            ROWS_INVALID,
+            "preview Parquet exceeds the candidate decompression budget",
+        )
+    return columns, total_row_count
+
+
+def open_sealed_parquet(
+    parquet_raw: bytes,
+    *,
+    max_columns: int,
+    max_rows: int,
+    max_uncompressed_bytes: int,
+) -> tuple[pq.ParquetFile, list[str], int]:
+    """Open sealed Parquet bytes and refuse anything outside the declared bounds.
+
+    Every guard reads only Parquet metadata, so an oversized or malformed file is refused
+    before a single row group is decompressed. The bounds arrive as arguments: this module
+    does not decide what a dataset's limits are, it only enforces the ones it is handed.
+    """
+
+    try:
+        parquet = pq.ParquetFile(pa.BufferReader(parquet_raw))
+        columns, total_row_count = inspect_sealed_parquet(
+            parquet,
+            max_columns=max_columns,
+            max_rows=max_rows,
+            max_uncompressed_bytes=max_uncompressed_bytes,
+        )
+    except RowsetError:
+        raise
+    except Exception as error:
+        raise RowsetError(
+            ROWS_INVALID,
+            f"sealed Parquet preview derivation failed: {type(error).__name__}",
+        ) from error
+    return parquet, columns, total_row_count
+
+
+def fit_canonical_payload(
+    encode: Callable[[int], bytes],
+    row_count: int,
+    *,
+    max_bytes: int,
+) -> bytes:
+    """Return the encoding of the longest row prefix that fits ``max_bytes``.
+
+    The loop searches rather than estimates because the canonical encoding of a row is not a
+    function of its Python size: escaping, key ordering and the shared envelope all move the
+    encoded length, and ``MAX_CANONICAL_MEMBERS`` bounds a member COUNT rather than a byte
+    count. The only honest bound on a candidate prefix is the length of that prefix once
+    encoded, so the prefix length is what gets bisected. A ``CanonicalJSONError`` from the
+    encoder is read as "this prefix is too large" and shrinks the search, which is what keeps
+    the member ceiling a bound on the payload rather than an error propagated to the caller.
+
+    A payload is always shrunk by whole rows from the end. A cell value is never truncated,
+    so every row that is served is the row that was sealed.
+    """
+
+    low = 0
+    high = row_count
+    selected = b""
+    while low <= high:
+        middle = (low + high) // 2
+        try:
+            encoded = encode(middle)
+        except CanonicalJSONError:
+            high = middle - 1
+            continue
+        if len(encoded) <= max_bytes:
+            selected = encoded
+            low = middle + 1
+        else:
+            high = middle - 1
+    if not selected:
+        raise RowsetError(
+            ROWS_INVALID,
+            "preview metadata alone exceeds the artifact byte budget",
+        )
+    return selected