186 lines
7.7 KiB
Python
186 lines
7.7 KiB
Python
"""
|
|
SCHEMA.md: what each table looks like, without its rows.
|
|
|
|
Written for reading, by a person or a web LLM that has to understand the data
|
|
before anything else: columns, an inferred type, how many are empty, one
|
|
example value, and how big the table really is. Types are inferred from what
|
|
pandas read, so they are a starting point, not a DDL.
|
|
|
|
Text columns, which is every column with all_text, are described by what their
|
|
values look like: integer-like, a code with leading zeros, a date in a given
|
|
format. Nothing is converted; the point is to have the evidence in one place
|
|
when the real schema is decided, in the loader.
|
|
"""
|
|
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
|
|
from sqlgen import human_bytes, sanitize_identifier
|
|
|
|
EXAMPLE_MAX = 40
|
|
# Text shapes are judged on at most this many values, spread through the column.
|
|
SHAPE_SAMPLE = 200_000
|
|
# A shape that fits this share of the values is reported as "mostly".
|
|
MOSTLY = 0.95
|
|
|
|
SHAPES = (
|
|
# (name, pattern), most specific first.
|
|
("integer-like", r"[+-]?\d+"),
|
|
("decimal-like", r"[+-]?(?:\d+\.\d*|\.\d+|\d+)(?:[eE][+-]?\d+)?"),
|
|
("timestamp-like (YYYY-MM-DD hh:mm)", r"\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:?\d{2})?"),
|
|
("date-like (YYYY-MM-DD)", r"\d{4}-\d{2}-\d{2}"),
|
|
("date-like (DDMONYYYY)", r"\d{1,2}[A-Za-z]{3}\d{4}"),
|
|
("timestamp-like (DDMONYYYY:hh:mm)", r"\d{1,2}[A-Za-z]{3}\d{4}[: ]\d{2}:\d{2}(?::\d{2})?"),
|
|
("date-like (D/M/Y or M/D/Y)", r"\d{1,2}/\d{1,2}/\d{2,4}"),
|
|
("boolean-like", r"(?i:y|n|yes|no|true|false)"),
|
|
)
|
|
|
|
|
|
def text_shape(values: pd.Series) -> str:
|
|
"""What a column of strings looks like, for choosing its real type later."""
|
|
text = values.astype(str)
|
|
longest = int(text.str.len().max())
|
|
if len(text) > SHAPE_SAMPLE:
|
|
text = text.iloc[:: len(text) // SHAPE_SAMPLE]
|
|
stripped = text.str.strip()
|
|
stripped = stripped[stripped != ""]
|
|
if stripped.empty:
|
|
return f"text, blank (max {longest})"
|
|
best = None
|
|
for name, pattern in SHAPES:
|
|
fits = stripped.str.fullmatch(pattern)
|
|
share = fits.mean()
|
|
if name == "integer-like" and share and stripped[fits].str.fullmatch(r"[+-]?0\d+").any():
|
|
name = "code with leading zeros"
|
|
if share == 1:
|
|
return f"text, {name} (max {longest})"
|
|
if share >= MOSTLY:
|
|
others = int((~fits).sum())
|
|
return f"text, mostly {name}: {others} other{'s' if others != 1 else ''} (max {longest})"
|
|
if best is None or share > best[1]:
|
|
best = (name, share)
|
|
# Half or more is still worth knowing when the type is being decided: a
|
|
# visit number that is "UNS" a third of the time is a decision, not text.
|
|
if best is not None and best[1] >= 0.5:
|
|
return f"text, {best[1]:.0%} {best[0]} (max {longest})"
|
|
return f"text (max {longest})"
|
|
|
|
|
|
def infer_type(series: pd.Series) -> str:
|
|
values = series.dropna()
|
|
if values.empty:
|
|
return "unknown (all null)"
|
|
if pd.api.types.is_bool_dtype(values):
|
|
return "boolean"
|
|
if pd.api.types.is_integer_dtype(values):
|
|
return "integer"
|
|
if pd.api.types.is_float_dtype(values):
|
|
return "integer" if (values % 1 == 0).all() else "numeric"
|
|
if pd.api.types.is_datetime64_any_dtype(values):
|
|
return "timestamp"
|
|
# Object columns from spreadsheets mix types; name the one that is there.
|
|
kinds = {type(v).__name__ for v in values}
|
|
if kinds <= {"int", "bool"}:
|
|
return "integer"
|
|
if kinds <= {"int", "float"}:
|
|
return "numeric"
|
|
if kinds <= {"datetime", "Timestamp"}:
|
|
return "timestamp"
|
|
if kinds == {"str"}:
|
|
return text_shape(values)
|
|
return f"mixed ({', '.join(sorted(kinds))})"
|
|
|
|
|
|
def example(series: pd.Series) -> str:
|
|
values = series.dropna()
|
|
if values.empty:
|
|
return ""
|
|
text = str(values.iloc[0]).replace("\n", " ").replace("|", "\\|")
|
|
return text if len(text) <= EXAMPLE_MAX else text[: EXAMPLE_MAX - 1] + "…"
|
|
|
|
|
|
class SchemaReport:
|
|
"""Collects one entry per written table, then writes them as one document."""
|
|
|
|
def __init__(self):
|
|
self.entries = []
|
|
|
|
def add(self, source, table, filename, df, total_rows, full_bytes, exact, shown_rows, parts=1):
|
|
self.entries.append(dict(
|
|
source=source, table=table, filename=filename, df=df,
|
|
total_rows=total_rows, full_bytes=full_bytes, exact=exact, shown_rows=shown_rows,
|
|
parts=parts,
|
|
))
|
|
|
|
def tables(self):
|
|
"""
|
|
One entry per output file. Several sources can feed the same table, and
|
|
the file holds all of them, so the report does too: rows and sizes are
|
|
summed, and the columns come from the first source.
|
|
"""
|
|
merged = {}
|
|
for e in self.entries:
|
|
m = merged.get(e["filename"])
|
|
if m is None:
|
|
merged[e["filename"]] = dict(e, sources=[e["source"]])
|
|
continue
|
|
m["sources"].append(e["source"])
|
|
m["total_rows"] += e["total_rows"]
|
|
m["full_bytes"] += e["full_bytes"]
|
|
m["shown_rows"] += e["shown_rows"]
|
|
m["exact"] = m["exact"] and e["exact"]
|
|
# The part count is read off the disk after each write, so the
|
|
# latest one already includes everything written before it.
|
|
m["parts"] = max(m["parts"], e["parts"])
|
|
return list(merged.values())
|
|
|
|
def write(self, out_dir: Path, max_rows):
|
|
tables = self.tables()
|
|
lines = ["# Data schema", ""]
|
|
total_rows = sum(t["total_rows"] for t in tables)
|
|
total_bytes = sum(t["full_bytes"] for t in tables)
|
|
lines.append(
|
|
f"{len(tables)} tables · {total_rows} rows · full seed files "
|
|
f"{'~' if any(not t['exact'] for t in tables) else ''}{human_bytes(total_bytes)}"
|
|
)
|
|
lines.append("")
|
|
if max_rows is not None:
|
|
lines += [
|
|
f"The .sql files beside this one hold at most {max_rows} rows per table from each",
|
|
"source: they are samples for understanding the data, not seeds to load. Sizes",
|
|
"marked ~ are estimated from the rows that were written.",
|
|
"",
|
|
]
|
|
lines += ["| table | file | rows | full size | in the file |", "|---|---|---:|---:|---|"]
|
|
for t in tables:
|
|
size = ("" if t["exact"] else "~") + human_bytes(t["full_bytes"])
|
|
kept = "all" if t["exact"] else f"{t['shown_rows']} rows"
|
|
file = f"`{t['filename']}`"
|
|
if t["parts"] > 1:
|
|
stem = t["filename"][:-len(".sql")]
|
|
file = f"`{stem}.001.sql` … `{stem}.{t['parts']:03d}.sql` ({t['parts']} parts)"
|
|
lines.append(f"| `{t['table']}` | {file} | {t['total_rows']} | {size} | {kept} |")
|
|
lines.append("")
|
|
|
|
for t in tables:
|
|
df = t["df"]
|
|
sources = ", ".join(f"`{src}`" for src in t["sources"])
|
|
folder = t["filename"].rsplit("/", 1)[0] if "/" in t["filename"] else ""
|
|
heading = f"{t['table']} · {folder}" if folder else t["table"]
|
|
lines += [f"## {heading}", "", f"From {sources} · {t['total_rows']} rows · {len(df.columns)} columns", ""]
|
|
lines += ["| column | source header | type | nulls | example |", "|---|---|---|---:|---|"]
|
|
for col in df.columns:
|
|
series = df[col]
|
|
header = str(col).replace("|", "\\|")
|
|
lines.append(
|
|
f"| `{sanitize_identifier(col)}` | {header} | {infer_type(series)} | "
|
|
f"{int(series.isna().sum())} | {example(series)} |"
|
|
)
|
|
lines.append("")
|
|
|
|
path = out_dir / "SCHEMA.md"
|
|
path.write_text("\n".join(lines), encoding="utf-8")
|
|
return path
|