p6flow 0.1.2__py3-none-any.whl

p6flow/inference.py

@@ -0,0 +1,204 @@
+"""Type resolution for XER columns.
+
+Three-tier fallback:
+
+    1. pmSchema.xml (via p6flow.schema.COLUMN_TYPES). If known, use that.
+    2. Naming heuristic + value sniffing must agree. Raise on mismatch.
+    3. If neither has an opinion, fall back to VARCHAR.
+
+Tier 2 only runs for columns absent from pmSchema.xml. The heuristic alone
+is fast but wrong for the long tail; the sniffer alone has small-sample
+ambiguity (an all-zeros column could be int or double). Requiring both to
+agree catches confidently-wrong inferences and surfaces them as errors.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+
+from .schema import COLUMN_TYPES, ColumnSpec
+
+# Order matters: longer/more specific suffixes first.
+_SUFFIX_BUCKETS: tuple[tuple[str, str], ...] = (
+    # Compound suffixes that override their tail
+    ("_pct_type", "string"),
+    ("_hr_cnt", "double"),
+    # Integers
+    ("_id", "integer"),
+    ("_cnt", "integer"),
+    ("_num", "integer"),
+    ("_seq", "integer"),
+    # Doubles
+    ("_pct", "double"),
+    ("_qty", "double"),
+    ("_rate", "double"),
+    ("_wt", "double"),
+    ("_value", "double"),
+    ("_cost", "double"),
+    ("_drtn", "double"),  # duration
+    # Dates
+    ("_date", "date"),
+    # Flags / single-char codes
+    ("_flag", "flag"),
+    # Strings
+    ("_type", "string"),
+    ("_code", "string"),
+    ("_name", "string"),
+    ("_user", "string"),
+    ("_short_name", "string"),
+    ("_path", "string"),
+    ("_state", "string"),
+    ("_guid", "string"),
+)
+
+_DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}( \d{2}:\d{2}(:\d{2})?)?$")
+
+_BUCKET_TO_DUCKDB = {
+    "integer": "BIGINT",
+    "double": "DOUBLE",
+    "date": "TIMESTAMP",
+    "flag": "VARCHAR(1)",
+    "string": "VARCHAR",
+}
+_BUCKET_TO_SNOWFLAKE = {
+    "integer": "NUMBER(19)",
+    "double": "FLOAT",
+    "date": "TIMESTAMP_NTZ",
+    "flag": "VARCHAR(1)",
+    "string": "VARCHAR",
+}
+
+
+class TypeMismatch(Exception):
+    """Naming heuristic and value sniffing disagree on a column's type."""
+
+
+def heuristic_bucket(field: str) -> str | None:
+    """Return a type bucket from the column name suffix, or None."""
+    for suffix, bucket in _SUFFIX_BUCKETS:
+        if field == suffix.lstrip("_") or field.endswith(suffix):
+            return bucket
+    return None
+
+
+def sniff_bucket(values: Iterable[str], max_samples: int = 1000) -> str | None:
+    """Return a type bucket from observed values, or None if all empty.
+
+    Samples up to max_samples non-empty values to avoid O(n) scans of huge tables.
+    """
+    non_empty: list[str] = []
+    for v in values:
+        if v != "" and v is not None:
+            non_empty.append(v)
+            if len(non_empty) >= max_samples:
+                break
+    if not non_empty:
+        return None
+
+    # Date check first; "2025-01-01" parses as nothing else useful.
+    if all(_DATE_RE.match(v) for v in non_empty):
+        return "date"
+    # Y/N flags
+    if all(v in ("Y", "N") for v in non_empty):
+        return "flag"
+    # Integers (no decimal point, parseable)
+    if all(_is_int(v) for v in non_empty):
+        return "integer"
+    # Floats (parseable as float, possibly with decimal)
+    if all(_is_float(v) for v in non_empty):
+        return "double"
+    return "string"
+
+
+def _is_int(v: str) -> bool:
+    try:
+        int(v)
+        return True
+    except ValueError:
+        return False
+
+
+def _is_float(v: str) -> bool:
+    try:
+        float(v)
+        return True
+    except ValueError:
+        return False
+
+
+# What each predicted bucket can cleanly hold. Asymmetric: a column predicted
+# `double` accepts integer-valued data (ints fit in doubles), but a column
+# predicted `integer` does NOT accept floating-point data (would lose precision).
+_ABSORBS: dict[str, frozenset[str]] = {
+    "string": frozenset({"string", "flag", "integer", "double", "date"}),
+    "double": frozenset({"double", "integer"}),
+    "integer": frozenset({"integer"}),
+    "flag": frozenset({"flag"}),
+    "date": frozenset({"date"}),
+}
+
+
+def infer_unknown(table: str, field: str, values: Iterable[str]) -> ColumnSpec:
+    """Resolve type for a column NOT in COLUMN_TYPES.
+
+    Raises TypeMismatch if the naming heuristic and value sniffing disagree
+    in a way that would lose information. Compatible disagreements (naming
+    says double but data is integer) resolve to the wider type. If neither
+    has an opinion (no suffix match, all values empty), default to string.
+    """
+    h = heuristic_bucket(field)
+    s = sniff_bucket(values)
+    if h is None and s is None:
+        bucket = "string"
+    elif h is None:
+        bucket = s
+    elif s is None:
+        bucket = h
+    elif s in _ABSORBS.get(h, frozenset()):
+        bucket = h
+    else:
+        raise TypeMismatch(
+            f"{table}.{field}: naming says {h!r}, sniffing says {s!r}"
+        )
+    return ColumnSpec(
+        duckdb_type=_BUCKET_TO_DUCKDB[bucket],
+        snowflake_type=_BUCKET_TO_SNOWFLAKE[bucket],
+        nullable=True,
+        description="",
+    )
+
+
+def resolve_overrides(
+    blocks: Iterable[tuple[str, list[str], Iterable[list[str]]]],
+) -> dict[tuple[str, str], ColumnSpec]:
+    """Build per-XER type overrides for every column not in COLUMN_TYPES.
+
+    blocks: iterable of (table_name, field_names, row_iter) where each row is
+    a list of string cell values aligned with field_names.
+
+    Returns: {(table, field): ColumnSpec} for every unknown column. Raises
+    TypeMismatch on the first disagreement and stops; this is intentional
+    fail-fast behavior so the caller learns about new ambiguous columns
+    instead of silently mis-typing them.
+    """
+    overrides: dict[tuple[str, str], ColumnSpec] = {}
+    for table, fields, rows in blocks:
+        unknown_idx = [
+            (i, f)
+            for i, f in enumerate(fields)
+            if f not in COLUMN_TYPES.get(table, {})
+        ]
+        if not unknown_idx:
+            continue
+        # Materialize columns we care about.
+        cols: dict[int, list[str]] = {i: [] for i, _ in unknown_idx}
+        for row in rows:
+            for i in cols:
+                if i < len(row):
+                    cols[i].append(row[i])
+        for i, f in unknown_idx:
+            if (table, f) in overrides:
+                continue
+            overrides[(table, f)] = infer_unknown(table, f, cols[i])
+    return overrides

p6flow/loader.py

@@ -0,0 +1,189 @@
+"""Load parsed XER tables into a DuckDB connection.
+
+Pipeline:
+    1. Materialize all tables (rows in memory, since type inference and
+       Arrow registration both need the full row set).
+    2. Resolve type overrides for unknown columns via p6flow.inference.
+       Fail fast if naming heuristic and value sniffing disagree.
+    3. Topo-sort tables so FK parents are inserted before children.
+    4. CREATE OR REPLACE TABLE for each one with constraints from schema.py.
+    5. INSERT data with per-column CAST(NULLIF(col, '') AS <type>) so empty
+       XER cells become SQL NULL.
+
+Returns a per-table row count map for downstream validation.
+"""
+
+from __future__ import annotations
+
+import duckdb
+import pyarrow as pa
+
+from .ddl import duckdb_ddl, topo_sort_tables
+from .inference import resolve_overrides
+from .schema import COLUMN_TYPES
+from .tokenizer import Table  # parse_tables imported lazily by callers
+
+
+class RowArityMismatch(ValueError):
+    """A `%R` row had MORE fields than `%F` declared. Format-level violation."""
+
+
+def _normalize_rows(
+    table: str,
+    fields: list[str],
+    rows: list[list[str]],
+) -> list[list[str]]:
+    """Force every row to len(fields). Pad short rows with '', drop
+    trailing empty cells before checking arity, error on real overflow.
+
+    Short rows are common in P6 exports (trailing optional fields dropped).
+    Some exporters also emit a trailing tab on every row, producing an
+    extra empty cell that's not a real field. We tolerate that quietly
+    and only raise when extra cells carry actual data.
+    """
+    n = len(fields)
+    out: list[list[str]] = []
+    for i, row in enumerate(rows):
+        # Drop trailing empties before arity check; they're an exporter
+        # artifact, not real data.
+        while len(row) > n and row[-1] == "":
+            row = row[:-1]
+        if len(row) > n:
+            raise RowArityMismatch(
+                f"{table} row {i}: {len(row)} fields but %F declared {n}; "
+                f"extra cells: {row[n:]!r}"
+            )
+        if len(row) < n:
+            row = row + [""] * (n - len(row))
+        out.append(row)
+    return out
+
+
+def _materialize(
+    tables: list[Table],
+) -> list[tuple[str, list[str], list[list[str]], list[str]]]:
+    """Snapshot each parsed table into (name, fields, rows, raws) tuples
+    used by downstream type inference, Arrow registration, and the DW
+    column path. Trailing empty cells are dropped from rows but raws is
+    left untouched (its job is byte-faithful capture)."""
+    out = []
+    for t in tables:
+        rows = _normalize_rows(t.name, t.fields, list(t.rows))
+        out.append((t.name, t.fields, rows, list(t.raws)))
+    return out
+
+
+def _cast_select(table: str, fields: list[str], overrides) -> str:
+    """Build a SELECT list that casts each string column to its target type.
+
+    Uses NULLIF(col, '') to convert XER's empty-string convention to SQL NULL.
+    The cast type comes from per-XER overrides first, then schema.py, with
+    VARCHAR as final fallback.
+    """
+    parts: list[str] = []
+    for f in fields:
+        spec = None
+        if overrides and (table, f) in overrides:
+            spec = overrides[(table, f)]
+        if spec is None:
+            spec = COLUMN_TYPES.get(table, {}).get(f)
+        ddb_type = spec.duckdb_type if spec else "VARCHAR"
+        if ddb_type == "VARCHAR" or ddb_type.startswith("VARCHAR("):
+            # No cast needed; just NULLIF empty strings. Don't NULLIF
+            # 'null' because a VARCHAR column might legitimately hold
+            # the literal string "null" (e.g. notes/descriptions).
+            parts.append(f'NULLIF("{f}", \'\') AS "{f}"')
+        elif ddb_type == "BLOB":
+            # CAST(varchar AS BLOB) interprets the string as hex-escaped
+            # bytes and rejects raw non-ASCII (e.g. UTF-8 BOM survivors,
+            # curly quotes in description fields). encode() takes the
+            # string's UTF-8 byte representation directly.
+            parts.append(f'encode(NULLIF("{f}", \'\')) AS "{f}"')
+        else:
+            # Typed casts (BIGINT/DOUBLE/TIMESTAMP/etc.) must treat both
+            # '' and the literal string 'null'/'NULL' as SQL NULL. Some
+            # P6 exporters emit the string "null" in lieu of empty cells
+            # for missing numerics; without this, the cast fails.
+            parts.append(
+                f"CAST(NULLIF(NULLIF(LOWER(\"{f}\"), 'null'), '') "
+                f"AS {ddb_type}) AS \"{f}\""
+            )
+    return ", ".join(parts)
+
+
+def load_xer(
+    con: duckdb.DuckDBPyConnection,
+    tables: list[Table],
+    add_dw_columns: bool = False,
+) -> dict[str, int]:
+    """Load every parsed table into the DuckDB connection. Return row counts per table.
+
+    When add_dw_columns=True, every loaded table gains an `_raw` VARCHAR
+    column containing the byte-faithful `%R\\t...` line captured by the
+    tokenizer (newline-stripped). Used by the warehouse path so
+    consumers can recover the exact source line of any parsed row for
+    debugging. Other DW columns (source_xer, source_sha256,
+    flow_published_at, exported_at) are constants per file and get
+    added at COPY time in output.write_parquet.
+    """
+    materialized = _materialize(tables)
+    rows_by_name = {t: (fs, rs, raws) for t, fs, rs, raws in materialized}
+    available = {t for t, _, _, _ in materialized}
+
+    # Type inference for unknown columns. Fails fast on conflicting evidence.
+    # resolve_overrides only needs (name, fields, row_iter); raws are unused.
+    overrides = resolve_overrides(
+        [(t, fs, rs) for t, fs, rs, _ in materialized]
+    )
+
+    # Topo sort by FK dependencies (parents first; DuckDB enforces FKs at INSERT).
+    ordered_pairs = topo_sort_tables([(t, fs) for t, fs, _, _ in materialized])
+
+    counts: dict[str, int] = {}
+    for table, fields in ordered_pairs:
+        _, rows, raws = rows_by_name[table]
+
+        # An empty %T block (zero columns, zero rows) is a real XER
+        # quirk; treat it as "not present in this export" rather than
+        # creating a column-less table that DuckDB rejects.
+        if not fields:
+            continue
+
+        emit_fields = list(fields)
+        emit_rows = rows
+        if add_dw_columns:
+            # Append _raw alongside the parsed cells using the raw line
+            # captured at tokenization time. VARCHAR by default since
+            # it isn't in COLUMN_TYPES.
+            #
+            # Defensive: if for some reason raws is shorter than rows
+            # (shouldn't happen in well-formed XERs, but guard against
+            # future tokenizer changes), fall back to empty string.
+            emit_fields = [*fields, "_raw"]
+            emit_rows = [
+                [*r, raws[i] if i < len(raws) else ""]
+                for i, r in enumerate(rows)
+            ]
+
+        # CREATE OR REPLACE the staging table; only emit FKs to tables we'll create.
+        con.execute(
+            duckdb_ddl(table, emit_fields, overrides, available_tables=available)
+        )
+
+        if not emit_rows:
+            counts[table] = 0
+            continue
+
+        # Pivot rows into columns, build an Arrow table of all-string cells,
+        # register it, INSERT with casts.
+        cols = list(zip(*emit_rows, strict=True))
+        arrow_table = pa.table({f: list(cols[i]) for i, f in enumerate(emit_fields)})
+        con.register("xer_raw", arrow_table)
+        try:
+            select_list = _cast_select(table, emit_fields, overrides)
+            con.execute(f"INSERT INTO {table} SELECT {select_list} FROM xer_raw")
+        finally:
+            con.unregister("xer_raw")
+
+        counts[table] = len(emit_rows)
+    return counts

p6flow/output.py

@@ -0,0 +1,173 @@
+"""Parquet writer and manifest builder.
+
+Writes one Parquet file per loaded DuckDB table plus a manifest.json
+that captures provenance, row counts, and schema fingerprints. The
+manifest is the cheap "did anything change" diff target across runs.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import duckdb
+
+from .tokenizer import Header, read_header
+from .validate import FkViolation
+
+
+def _sha256(path: Path, chunk_size: int = 1 << 20) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _schema_fingerprint(con: duckdb.DuckDBPyConnection, table: str) -> str:
+    """Hash of (column_name, type) tuples in declaration order. Stable
+    across runs as long as DDL doesn't drift."""
+    rows = con.execute(f'DESCRIBE "{table}"').fetchall()
+    payload = "|".join(f"{r[0]}:{r[1]}" for r in rows)
+    return hashlib.sha256(payload.encode()).hexdigest()[:16]
+
+
+def _violation_to_dict(v: FkViolation) -> dict[str, Any]:
+    return {
+        "child_table": v.child_table,
+        "fk_name": v.fk_name,
+        "target_table": v.target_table,
+        "local_fields": list(v.local_fields),
+        "target_fields": list(v.target_fields),
+        "kind": v.kind,
+        "violating_rows": v.violating_rows,
+        "sample_keys": [list(k) for k in v.sample_keys],
+    }
+
+
+def _sql_str_lit(s: str) -> str:
+    """Single-quote a value for safe inlining into a DuckDB SQL literal."""
+    return "'" + s.replace("'", "''") + "'"
+
+
+def _dw_select(
+    table: str,
+    cols_in_table: set[str],
+    src_tag: str,
+    sha: str,
+    export_date_lit: str,
+) -> str:
+    """Build the SELECT projection that adds DW columns to one table.
+
+    Source tables (parsed from %T blocks) already carry `_raw` from
+    load_xer. Derived tables (CALENDAR_WEEKLY, *_UDF, UDF_COLUMN_MAP)
+    don't, so we emit `_raw = NULL` to keep every parquet's DW column
+    shape uniform across the corpus.
+    """
+    raw_clause = (
+        '"_raw"' if "_raw" in cols_in_table else "CAST(NULL AS VARCHAR) AS _raw"
+    )
+    base = (
+        'SELECT * EXCLUDE ("_raw")'
+        if "_raw" in cols_in_table
+        else "SELECT *"
+    )
+    return (
+        f'{base}, '
+        f'{_sql_str_lit(src_tag)} AS source_xer, '
+        f'{_sql_str_lit(sha)} AS source_sha256, '
+        f'CAST((CURRENT_TIMESTAMP AT TIME ZONE \'UTC\') AS TIMESTAMP) '
+        f'AS flow_published_at, '
+        f'{export_date_lit} AS exported_at, '
+        f'{raw_clause} '
+        f'FROM "{table}"'
+    )
+
+
+def write_parquet(
+    con: duckdb.DuckDBPyConnection,
+    out_dir: Path,
+    source_xer: Path,
+    counts: dict[str, int],
+    violations: list[FkViolation] | None = None,
+    udf_unknown_types: list[dict] | None = None,
+    add_dw_columns: bool = False,
+) -> dict[str, Any]:
+    """Write each DuckDB table to a Parquet file in out_dir, plus manifest.json.
+
+    Returns the manifest dict for inspection in callers/tests. If
+    violations is None, the manifest's "fk_violations" key is omitted
+    (validation was not requested for this run). udf_unknown_types is
+    only emitted when non-empty.
+
+    When add_dw_columns is True, every emitted parquet gets data-
+    warehouse columns suitable for cross-XER joins:
+
+      source_xer         basename of the input file (no extension)
+      source_sha256      hex digest matching the manifest entry
+      flow_published_at  TIMESTAMP at COPY time (ingestion clock)
+      exported_at         TIMESTAMP from ERMHDR export_date
+      _raw               reconstructed `%R\\t...` line per row
+                         (added in load_xer; only on source tables —
+                         derived tables like CALENDAR_WEEKLY do not
+                         have an upstream %R line).
+
+    The natural cross-XER join key becomes (source_xer, <p6_id>) since
+    P6 IDs are stable inside one export but collide across exports.
+    """
+    out_dir.mkdir(parents=True, exist_ok=True)
+    header: Header = read_header(source_xer)
+    sha = _sha256(source_xer)
+    src_tag = source_xer.stem
+    export_date_lit = (
+        f"CAST({_sql_str_lit(header.export_date)} AS TIMESTAMP)"
+        if header.export_date
+        else "CAST(NULL AS TIMESTAMP)"
+    )
+
+    table_entries: list[dict[str, Any]] = []
+    for table in sorted(counts):
+        parquet_path = out_dir / f"{table}.parquet"
+        if add_dw_columns:
+            cols_in_table = {
+                r[0] for r in con.execute(f'DESCRIBE "{table}"').fetchall()
+            }
+            select_sql = _dw_select(
+                table, cols_in_table, src_tag, sha, export_date_lit
+            )
+            con.execute(
+                f"COPY ({select_sql}) TO '{parquet_path}' "
+                "(FORMAT PARQUET, COMPRESSION ZSTD)"
+            )
+        else:
+            con.execute(
+                f"COPY \"{table}\" TO '{parquet_path}' "
+                "(FORMAT PARQUET, COMPRESSION ZSTD)"
+            )
+        table_entries.append(
+            {
+                "name": table,
+                "rows": counts[table],
+                "parquet": parquet_path.name,
+                "schema_fingerprint": _schema_fingerprint(con, table),
+            }
+        )
+
+    manifest: dict[str, Any] = {
+        "source_xer": str(source_xer.resolve()),
+        "source_sha256": _sha256(source_xer),
+        "parsed_at": datetime.now(tz=UTC).isoformat(),
+        "p6_version": header.version,
+        "export_date": header.export_date,
+        "project_label": header.project_label,
+        "tables": table_entries,
+    }
+    if violations is not None:
+        manifest["fk_violations"] = [_violation_to_dict(v) for v in violations]
+    if udf_unknown_types:
+        manifest["udf_unknown_types"] = udf_unknown_types
+    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    return manifest