floe-python 0.8.0__py3-none-any.whl

floe/checks/mismatch.py

@@ -0,0 +1,137 @@
+"""Schema mismatch: missing/extra columns and their policies (Rust: `checks/mismatch.rs`)."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import polars as pl
+
+from floe.checks.normalize import normalize_name, resolve_strategy, source_names
+from floe.config.types import EntityConfig
+from floe.errors import FloeConfigError, FloeIoError
+from floe.read.json import json_records
+
+MAX_MISMATCH_COLUMNS = 50
+# nested-path markers: `/` and `@` are XML path syntax, but ordinary characters in a JSON key
+_NESTED = {"json": (".", "["), "xml": (".", "[", "/", "@")}
+
+
+@dataclass(frozen=True)
+class MismatchOutcome:
+    """`report` is the `mismatch` object of the file report; the rest drives the run."""
+
+    report: dict[str, Any]
+    rejected: bool
+    aborted: bool
+    warnings: int
+    errors: int
+    missing: list[str]
+    extra: list[str]
+    fill_missing: bool
+    ignore_extra: bool
+
+
+def declared_names(entity: EntityConfig) -> list[str]:
+    """The declared names input columns are compared with (normalised)."""
+    schema = entity.schema_
+    strategy = resolve_strategy(schema)
+    if entity.source.format not in ("json", "xml"):
+        return source_names(schema.columns, strategy)
+    # Nested selectors are not file columns; only the top-level ones are compared.
+    names = [
+        normalize_name(c.source or c.name, strategy)
+        for c in schema.columns
+        if not any(t in (c.source or c.name) for t in _NESTED[entity.source.format])
+    ]
+    if len(set(names)) != len(names):
+        dup = next(n for n in names if names.count(n) > 1)
+        raise FloeConfigError(f"duplicate top-level column selector: {dup}")
+    return names
+
+
+def plan_mismatch(
+    entity: EntityConfig, declared: list[str], input_names: list[str]
+) -> MismatchOutcome:
+    strategy = resolve_strategy(entity.schema_)
+    inputs = [normalize_name(n, strategy) for n in input_names]
+    missing = sorted(set(declared) - set(inputs))
+    extra = sorted(set(inputs) - set(declared))
+    policy = entity.schema_.mismatch
+    on_missing = (policy and policy.missing_columns) or "fill_nulls"
+    on_extra = (policy and policy.extra_columns) or "ignore"
+    severity = entity.policy.severity
+
+    def requested() -> bool:
+        return (on_missing == "reject_file" and bool(missing)) or (
+            on_extra == "reject_file" and bool(extra)
+        )
+
+    warning = None
+    if requested() and severity == "warn":
+        warning = (
+            f"entity.name={entity.name} schema mismatch requested reject_file "
+            "but policy.severity=warn; continuing"
+        )
+        on_missing, on_extra = "fill_nulls", "ignore"
+    rejected = requested() and severity == "reject"
+    aborted = requested() and severity == "abort"
+    fill = not (rejected or aborted) and on_missing == "fill_nulls" and bool(missing)
+    ignore = not (rejected or aborted) and on_extra == "ignore" and bool(extra)
+    action = (
+        "rejected_file" if rejected
+        else "aborted" if aborted
+        else "filled_nulls" if fill
+        else "ignored_extras" if ignore
+        else "none"
+    )  # fmt: skip
+    report: dict[str, Any] = {
+        "declared_columns_count": len(declared),
+        "input_columns_count": len(inputs),
+        "missing_columns": missing[:MAX_MISMATCH_COLUMNS],
+        "extra_columns": extra[:MAX_MISMATCH_COLUMNS],
+        "mismatch_action": action,
+    }
+    if rejected or aborted:
+        report["error"] = {
+            "rule": "schema_mismatch",
+            "message": f"entity.name={entity.name} schema mismatch: "
+            f"missing={len(missing)} extra={len(extra)}",
+        }
+    if warning:
+        report["warning"] = warning
+    return MismatchOutcome(
+        report, rejected, aborted, int(bool(warning)), int(rejected or aborted),
+        missing, extra, fill, ignore,
+    )  # fmt: skip
+
+
+def apply_mismatch(df: pl.DataFrame, plan: MismatchOutcome) -> pl.DataFrame:
+    """Add the missing columns as nulls and drop the extra ones, as the plan says."""
+    if plan.fill_missing:
+        # a column the reader already produced (an XML `a:id` selector the probe names `id`) stays
+        df = df.with_columns(
+            pl.lit(None, pl.String).alias(n)
+            for n in plan.missing
+            if n not in df.columns
+        )
+    if plan.ignore_extra:
+        df = df.drop(plan.extra, strict=False)
+    return df
+
+
+def json_input_columns(path: Path, mode: str = "array") -> list[str] | None:
+    """Sorted top-level keys of every record, read exactly as `read_json` reads them.
+
+    D40: the union over all records, so a key only a later record carries is still an input
+    column (Rust 0.7 probed the first object only). `None` for a file without records (D41), which
+    differs from records that carry no keys (`[]`): those are still checked for missing columns.
+    """
+    if mode not in ("array", "ndjson"):
+        raise FloeIoError(
+            f"unsupported source.options.json_mode={mode} (allowed: array, ndjson)"
+        )
+    # ponytail: the file is parsed twice (probe, then read); fold the probe into the read if it costs
+    keys: set[str] | None = None
+    for _, record in json_records(path, mode):
+        keys = (keys or set[str]()) | {str(k) for k in record}
+    return None if keys is None else sorted(keys)

floe/checks/normalize.py

@@ -0,0 +1,118 @@
+"""Column-name normalization (Rust: `checks/normalize.rs`). ASCII only, as Rust."""
+
+from collections.abc import Sequence
+from string import ascii_lowercase, ascii_uppercase
+
+import polars as pl
+
+from floe.config.types import ColumnConfig, SchemaConfig
+from floe.errors import FloeConfigError
+
+_STRATEGIES = {"snakecase", "lower", "camelcase", "none"}
+_LOWER = str.maketrans(ascii_uppercase, ascii_lowercase)
+
+
+def _canonical(value: str) -> str:
+    return value.lower().replace("-", "").replace("_", "")
+
+
+def resolve_strategy(schema: SchemaConfig) -> str | None:
+    """The canonical strategy when `normalize_columns.enabled`, else None."""
+    config = schema.normalize_columns
+    if config is None or not config.enabled:
+        return None
+    raw = config.strategy or "snake_case"
+    if (strategy := _canonical(raw)) not in _STRATEGIES:
+        raise FloeConfigError(f"unsupported normalize_columns.strategy: {raw}")
+    return strategy
+
+
+def _words(value: str) -> list[str]:
+    words: list[str] = []
+    current = ""
+    for i, ch in enumerate(value):
+        if not (ch.isascii() and ch.isalnum()):
+            if current:
+                words.append(current)
+            current = ""
+            continue
+        prev = value[i - 1] if i else ""
+        nxt = value[i + 1] if i + 1 < len(value) else ""
+        prev_lower = prev.isascii() and prev.islower()
+        prev_digit = prev.isascii() and prev.isdigit()
+        prev_upper = prev.isascii() and prev.isupper()
+        next_lower = nxt.isascii() and nxt.islower()
+        if (
+            current
+            and ch.isupper()
+            and (prev_lower or prev_digit or (prev_upper and next_lower))
+        ):
+            words.append(current)
+            current = ""
+        current += ch.translate(_LOWER)
+    return [*words, current] if current else words
+
+
+def normalize_name(value: str, strategy: str | None) -> str:
+    match None if strategy is None else _canonical(strategy):
+        case "snakecase":
+            return "_".join(_words(value))
+        case "lower":
+            return value.translate(_LOWER)
+        case "camelcase":
+            return "".join(
+                w if i == 0 else w[:1].upper() + w[1:]
+                for i, w in enumerate(_words(value))
+            )
+        case _:
+            return value
+
+
+def source_names(columns: Sequence[ColumnConfig], strategy: str | None) -> list[str]:
+    """Normalised `source or name` of each column; two columns on one name are an error."""
+    seen: dict[str, str] = {}
+    for column in columns:
+        source = column.source or column.name
+        name = normalize_name(source, strategy)
+        if name in seen:
+            raise FloeConfigError(
+                f"column source collision: {seen[name]} and {column.name} -> {name}"
+            )
+        seen[name] = source
+    return list(seen)
+
+
+def output_names(columns: Sequence[ColumnConfig], strategy: str | None) -> list[str]:
+    """Output name of each column: an explicit `source` keeps `name`, otherwise it is normalised."""
+    return [c.name if c.source else normalize_name(c.name, strategy) for c in columns]
+
+
+def output_column_mapping(
+    columns: Sequence[ColumnConfig], strategy: str | None
+) -> dict[str, str]:
+    """Normalised source name -> output name, for the columns the rename changes."""
+    targets: dict[str, str] = {}
+    mapping: dict[str, str] = {}
+    for column, target in zip(columns, output_names(columns, strategy), strict=True):
+        source = normalize_name(column.source or column.name, strategy)
+        if target in targets:
+            raise FloeConfigError(
+                f"output column name collision: {targets[target]} and {source} -> {target}"
+            )
+        targets[target] = source
+        if source != target:
+            mapping[source] = target
+    return mapping
+
+
+def normalize_columns(df: pl.DataFrame, strategy: str | None) -> pl.DataFrame:
+    """Rename the frame's columns; two columns on one name are an error."""
+    names = {c: normalize_name(c, strategy) for c in df.columns}
+    seen: dict[str, str] = {}
+    for old, new in names.items():
+        if new in seen:
+            raise FloeConfigError(
+                f"normalized input column collision: {seen[new]} and {old} -> {new}"
+            )
+        seen[new] = old
+    return df.rename(names)

floe/checks/rows.py

@@ -0,0 +1,116 @@
+"""not_null, row error columns and the accept/reject split (Rust: `checks/not_null.rs`, `checks/mod.rs`)."""
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+import polars as pl
+
+from floe.checks.cast import json_escape
+
+NOT_NULL_MESSAGE = "required value missing"
+ROW_INDEX = "__floe_row_index"
+ERRORS = "__floe_errors"
+
+Formatter = Literal[
+    "json", "csv", "text"
+]  # `report.formatter`, validated with the config
+
+# (prefix, separator, suffix) around the joined entries. csv and text are JSON strings: the
+# whole value is escaped and quoted, so the csv separator `\n` is written as the two characters.
+_WRAP: dict[Formatter, tuple[str, str, str]] = {
+    "json": ("[", ",", "]"),
+    "csv": ('"', "\n", '"'),
+    "text": ('"', "; ", '"'),
+}
+
+
+@dataclass(frozen=True)
+class RowCheck:
+    """One check over the rows: `mask` is true where the row violates it.
+
+    Pass the checks in Rust's error order: not_null (required columns), cast_error (declared
+    non-string columns), then unique.
+    """
+
+    rule: str
+    column: str
+    message: str
+    mask: pl.Series
+
+
+def not_null_check(df: pl.DataFrame, column: str) -> RowCheck:
+    """A required column. `df` holds the typed values, so a failed cast also violates it."""
+    return RowCheck("not_null", column, NOT_NULL_MESSAGE, df[column].is_null())
+
+
+def error_entry(rule: str, column: str, message: str, source: str | None = None) -> str:
+    """Rust's `RowError::to_json_with_source`."""
+    src = "" if source is None else f'"source":"{json_escape(source)}",'
+    return (
+        f'{{"rule":"{json_escape(rule)}","column":"{json_escape(column)}",{src}'
+        f'"message":"{json_escape(message)}"}}'
+    )
+
+
+def _csv_escape(value: str) -> str:
+    escaped = value.replace('"', '""')
+    return f'"{escaped}"' if any(c in escaped for c in ",\n\r") else escaped
+
+
+def _entry(
+    formatter: Formatter, check: RowCheck, sources: Mapping[str, str] | None
+) -> str:
+    """What one violated check adds to a row's errors, per Rust `RowErrorFormatter`."""
+    source = sources.get(check.column) if sources else None
+    rule, column, message = check.rule, check.column, check.message
+    if formatter == "json":
+        return error_entry(rule, column, message, source)
+    if formatter == "csv":
+        fields = [rule, column, message]
+        if (
+            sources
+        ):  # with a source map every line has the source field, empty if unmapped
+            fields.insert(2, source or "")
+        return ",".join(_csv_escape(f) for f in fields)
+    src = "" if source is None else f" source={source}"
+    return f"{rule}:{column}{src} {message}"
+
+
+def error_column(
+    height: int,
+    checks: Sequence[RowCheck],
+    formatter: Formatter = "json",
+    sources: Mapping[str, str] | None = None,
+) -> pl.Series:
+    """`__floe_errors`: the formatted errors of each row, null for rows that pass every check.
+
+    `sources` maps a runtime column name to its source name (Rust's `source_column_map`).
+    """
+    if not checks:
+        return pl.Series(ERRORS, [None] * height, dtype=pl.String)
+    head, sep, tail = _WRAP[formatter]
+    escape = str if formatter == "json" else json_escape
+    entries = [
+        pl.when(c.mask).then(pl.lit(escape(_entry(formatter, c, sources))))
+        for c in checks
+    ]
+    joined = pl.concat_str(entries, separator=escape(sep), ignore_nulls=True)
+    failed = pl.any_horizontal([c.mask for c in checks])
+    return pl.select(
+        pl.when(failed).then(pl.lit(head) + joined + pl.lit(tail)).alias(ERRORS)
+    )[ERRORS]
+
+
+def split_rows(
+    df: pl.DataFrame,
+    checks: Sequence[RowCheck],
+    formatter: Formatter = "json",
+    sources: Mapping[str, str] | None = None,
+) -> tuple[pl.DataFrame, pl.DataFrame]:
+    """`(accepted, rejected)`; rejected adds `__floe_row_index` (UInt64, row position in `df`) and `__floe_errors`."""
+    errors = error_column(df.height, checks, formatter, sources)
+    rejected = df.with_columns(
+        pl.int_range(pl.len(), dtype=pl.UInt64).alias(ROW_INDEX), errors
+    ).filter(errors.is_not_null())
+    return df.filter(errors.is_null()), rejected

floe/checks/unique.py

@@ -0,0 +1,247 @@
+"""Unique keys within a file and across the files of one run (Rust: `checks/unique.rs`)."""
+# pyarrow ships no type stubs, so strict pyright cannot see its types.
+# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false, reportUnknownArgumentType=false
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from datetime import datetime
+from decimal import Decimal
+
+import polars as pl
+import pyarrow as pa
+
+from floe.checks.rows import RowCheck
+
+SAMPLE_LIMIT = 5
+_SEEN = "__floe_seen"
+_DEBUG = "__floe_debug"
+_COUNT = "__floe_count"
+_BITS = "_bits"  # suffix of the bit-pattern column that keys a float column
+
+
+@dataclass(frozen=True)
+class UniqueSample:
+    values: dict[str, str]
+    count: int
+
+
+@dataclass(frozen=True)
+class UniqueResult:
+    """Rust's `UniqueConstraintResult`; every duplicate is a batch duplicate (no target history)."""
+
+    columns: list[str]
+    duplicates_count: int
+    samples: list[UniqueSample]
+    target_duplicates_count: int = 0  # duplicates of keys seeded from the existing sink
+
+
+def resolve_unique_keys(
+    unique_keys: Sequence[Sequence[str]] | None,
+    unique_columns: Sequence[str],
+    primary_key: Sequence[str] | None,
+) -> list[list[str]]:
+    """Rust's `resolve_schema_unique_keys`: `unique_keys`, else the `unique: true` columns, then the primary key.
+
+    Names are trimmed; empty and repeated keys are dropped.
+    """
+    keys = unique_keys if unique_keys is not None else [[c] for c in unique_columns]
+    resolved: list[list[str]] = []
+    for key in [*keys, *([primary_key] if primary_key else [])]:
+        names = [c.strip() for c in key]
+        if names and names not in resolved:
+            resolved.append(names)
+    return resolved
+
+
+# `str` Debug; ponytail: non-ASCII non-printable characters (Rust `\\u{..}`) are left raw
+_STR_ESCAPES = {
+    "\\": "\\\\",
+    '"': '\\"',
+    "\n": "\\n",
+    "\r": "\\r",
+    "\t": "\\t",
+    "\0": "\\0",
+    **{chr(i): f"\\u{{{i:x}}}" for i in [*range(1, 32), 127] if chr(i) not in "\n\r\t"},
+}
+
+
+def _float_bits(s: pl.Series) -> pl.Series:
+    """Rust's `to_bits()`: Polars cannot reinterpret a float, so view the Arrow buffer."""
+    return pl.Series(s.cast(pl.Float64).to_arrow().view(pa.uint64()))
+
+
+def _debug(name: str, dtype: pl.DataType) -> pl.Expr:
+    """Rust's `{:?}` of the `UniqueKey` variant: the sample tie-break sorts on these strings."""
+    c = pl.col(name)
+    if dtype == pl.String:
+        escaped = c
+        for raw, esc in _STR_ESCAPES.items():
+            escaped = escaped.str.replace_all(raw, esc, literal=True)
+        return pl.format('String("{}")', escaped)
+    if dtype == pl.Boolean:
+        return pl.format("Bool({})", c.cast(pl.String))
+    if dtype.is_signed_integer():
+        return pl.format("I64({})", c.cast(pl.String))
+    if dtype.is_unsigned_integer():
+        return pl.format("U64({})", c.cast(pl.String))
+    if dtype.is_float():
+        return pl.format("F64({})", pl.col(name + _BITS).cast(pl.String))
+    if isinstance(dtype, pl.Datetime):  # chrono prints the fraction only when not zero
+        return pl.format('Other("{}")', c.dt.to_string("%Y-%m-%d %H:%M:%S%.f"))
+    # ponytail: date keys use `Other(display)`, unverified against Rust; ties may order differently
+    return pl.format('Other("{}")', c.cast(pl.String))
+
+
+def _text(value: object) -> str:
+    """Rust's `to_string` of a key value."""
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, float):
+        if not math.isfinite(value):
+            return "NaN" if math.isnan(value) else str(value).replace("infinity", "inf")
+        return format(Decimal(repr(value)).normalize(), "f")
+    if isinstance(value, datetime):
+        if not value.microsecond:
+            spec = "seconds"
+        else:
+            spec = "milliseconds" if value.microsecond % 1000 == 0 else "microseconds"
+        return value.replace(tzinfo=None).isoformat(sep=" ", timespec=spec)
+    return str(value)
+
+
+@dataclass
+class _Constraint:
+    columns: list[str]
+    seen: pl.DataFrame
+    dup_counts: list[pl.DataFrame] = field(default_factory=list[pl.DataFrame])
+    duplicates: int = 0
+    target: int = 0
+    seeded: pl.DataFrame = field(default_factory=pl.DataFrame)
+
+    @property
+    def ids(self) -> list[str]:
+        """Positional names, so key columns named like our temporaries or regexes stay literal."""
+        return [f"k{i}" for i in range(len(self.columns))]
+
+
+class UniqueTracker:
+    """Keeps the keys seen so far, so a key repeated in a later file is a duplicate too."""
+
+    def __init__(self, keys: Sequence[Sequence[str]]) -> None:
+        self._constraints = [_Constraint(list(k), pl.DataFrame()) for k in keys]
+
+    @property
+    def keys(self) -> list[list[str]]:
+        return [list(c.columns) for c in self._constraints]
+
+    def check(self, df: pl.DataFrame) -> list[RowCheck]:
+        """One `RowCheck` per constraint for the rows of the next file, in constraint order.
+
+        Pass every row, rejected ones included. A row with a NULL in any key column is
+        skipped: never a duplicate, never remembered.
+        """
+        return [self._apply(c, df) for c in self._constraints]
+
+    def snapshot(self) -> list[tuple[pl.DataFrame, list[pl.DataFrame], int, int]]:
+        """The state `restore` goes back to: a file that is rejected as unreadable leaves no keys."""
+        return [
+            (c.seen, list(c.dup_counts), c.duplicates, c.target)
+            for c in self._constraints
+        ]
+
+    def restore(
+        self, snapshot: list[tuple[pl.DataFrame, list[pl.DataFrame], int, int]]
+    ) -> None:
+        for c, (seen, counts, dups, target) in zip(
+            self._constraints, snapshot, strict=True
+        ):
+            c.seen, c.dup_counts, c.duplicates, c.target = (
+                seen,
+                list(counts),
+                dups,
+                target,
+            )
+
+    def seed(self, df: pl.DataFrame) -> None:
+        """Remember the keys of rows already in the sink (append mode); a later hit is a target duplicate."""
+        for c in self._constraints:
+            if all(n in df.columns for n in c.columns):
+                keys = self._keys(c, df)
+                valid = keys.select(pl.all_horizontal(pl.all().is_not_null()))[c.ids[0]]
+                fresh = keys.filter(valid).unique()
+                c.seen = pl.concat([c.seen, fresh]).unique() if c.seen.width else fresh
+                # Iceberg seeds chunk by chunk: keys the batch already holds stay batch keys
+                c.seeded = (
+                    pl.concat([c.seeded, fresh]).unique() if c.seeded.width else fresh
+                )
+
+    @staticmethod
+    def _keys(c: _Constraint, df: pl.DataFrame) -> pl.DataFrame:
+        keys = pl.DataFrame(
+            [df[n].alias(k) for n, k in zip(c.columns, c.ids, strict=True)]
+        )
+        # Rust compares `to_bits()`: -0.0 differs from 0.0, NaN payloads differ; key on the bits too
+        return keys.with_columns(
+            pl.col(k).map_batches(_float_bits, return_dtype=pl.UInt64).alias(k + _BITS)
+            for k in c.ids
+            if keys.schema[k].is_float()
+        )
+
+    def _apply(self, c: _Constraint, df: pl.DataFrame) -> RowCheck:
+        keys = self._keys(c, df)
+        valid = keys.select(pl.all_horizontal(pl.all().is_not_null()))[c.ids[0]]
+        dup = pl.Series([False] * df.height)
+        if (
+            valid.any()
+        ):  # skips all-null (Null dtype) frames, which cannot join the seen keys
+            first = keys.select(pl.struct(keys.columns).is_first_distinct())[c.ids[0]]
+            dup = valid & ~first
+            if c.seen.width:
+                hit = keys.join(
+                    c.seen.with_columns(pl.lit(True).alias(_SEEN)),
+                    on=keys.columns,
+                    how="left",
+                    maintain_order="left",
+                )[_SEEN]
+                dup = dup | (valid & hit.fill_null(False))
+            fresh = keys.filter(valid & ~dup)
+            c.seen = pl.concat([c.seen, fresh]) if c.seen.width else fresh
+            dups = keys.filter(dup)
+            c.duplicates += dups.height
+            if c.seeded.width:
+                c.target += dups.join(c.seeded, on=keys.columns, how="semi").height
+            if dups.height:
+                c.dup_counts.append(dups.group_by(keys.columns).len(_COUNT))
+        single = len(c.columns) == 1
+        name = c.columns[0] if single else f"[{','.join(c.columns)}]"
+        message = "duplicate value" if single else "duplicate key"
+        return RowCheck("unique", name, message, dup)
+
+    def results(self) -> list[UniqueResult]:
+        """Per constraint, the top `SAMPLE_LIMIT` keys by count, ties by Rust's Debug string."""
+        return [self._result(c) for c in self._constraints]
+
+    def _result(self, c: _Constraint) -> UniqueResult:
+        if not c.dup_counts:
+            return UniqueResult(c.columns, 0, [])
+        parts = pl.concat(c.dup_counts)
+        counts = parts.group_by(parts.columns[:-1]).agg(pl.col(_COUNT).sum())
+        debug = pl.format(
+            "CompositeKey([{}])",
+            pl.concat_str([_debug(k, counts.schema[k]) for k in c.ids], separator=", "),
+        )
+        top = (
+            counts.with_columns(debug.alias(_DEBUG))
+            .sort([_COUNT, _DEBUG], descending=[True, False])
+            .head(SAMPLE_LIMIT)
+        )
+        samples = [  # at most SAMPLE_LIMIT rows
+            UniqueSample(
+                {n: _text(row[k]) for n, k in zip(c.columns, c.ids, strict=True)},
+                row[_COUNT],
+            )
+            for row in top.iter_rows(named=True)
+        ]
+        return UniqueResult(c.columns, c.duplicates, samples, c.target)