|
| 1 | +"""Project a LOKF bundle to a set of linked tables — the tabular counterpart to |
| 2 | +its RDF graph. |
| 3 | +
|
| 4 | +A bundle becomes **one table per concept type** (the nodes) plus a single |
| 5 | +``relations`` table (the typed edges: ``source``, ``predicate``, ``target``). |
| 6 | +From there the same well-modeled data is analysable as DataFrames, landable as |
| 7 | +CSV/Parquet, persistable as SQL, or registerable as a lakehouse via |
| 8 | +``CREATE EXTERNAL TABLE`` DDL for BigQuery or Athena. |
| 9 | +
|
| 10 | +pandas (and, optionally, polars / pyarrow) are only needed here, so they live in |
| 11 | +the ``tables`` extra:: |
| 12 | +
|
| 13 | + pip install "lokf[tables]" |
| 14 | +""" |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import json |
| 18 | +import pathlib |
| 19 | +from typing import Any |
| 20 | + |
| 21 | +from lokf.model import Bundle, load_bundle |
| 22 | +from lokf.schema import vocabulary |
| 23 | + |
| 24 | +#: Name of the edge table that holds every typed relation in the bundle. |
| 25 | +RELATIONS_TABLE = "relations" |
| 26 | + |
| 27 | + |
| 28 | +def _make_frame(rows: list[dict], engine: str): |
| 29 | + """Build a DataFrame of *rows* with the chosen engine (pandas | polars).""" |
| 30 | + if engine == "polars": |
| 31 | + try: |
| 32 | + import polars as pl |
| 33 | + except ModuleNotFoundError as exc: # pragma: no cover |
| 34 | + raise ModuleNotFoundError( |
| 35 | + "engine='polars' needs polars: pip install polars." |
| 36 | + ) from exc |
| 37 | + return pl.DataFrame(rows) |
| 38 | + try: |
| 39 | + import pandas as pd |
| 40 | + except ModuleNotFoundError as exc: # pragma: no cover |
| 41 | + raise ModuleNotFoundError( |
| 42 | + "lokf.tables needs pandas (and pyarrow for Parquet). " |
| 43 | + "Install the extra: pip install 'lokf[tables]'." |
| 44 | + ) from exc |
| 45 | + return pd.DataFrame(rows) |
| 46 | + |
| 47 | + |
| 48 | +def _scalarize(value: Any) -> Any: |
| 49 | + """Flatten a non-relation frontmatter value into a single table cell.""" |
| 50 | + if isinstance(value, list): |
| 51 | + return "; ".join(str(v) for v in value) |
| 52 | + if isinstance(value, dict): |
| 53 | + return json.dumps(value, ensure_ascii=False) |
| 54 | + return value |
| 55 | + |
| 56 | + |
| 57 | +def _split(doc: dict, rel_slots: set[str]) -> tuple[dict, list[dict]]: |
| 58 | + """Split one concept into a (node row, edge rows) pair. |
| 59 | +
|
| 60 | + Concept-ranged relation slots (and reified ``relations``) become edges; |
| 61 | + everything else is a scalar column on the node's type table. |
| 62 | + """ |
| 63 | + source = doc.get("id") |
| 64 | + node: dict[str, Any] = {} |
| 65 | + edges: list[dict] = [] |
| 66 | + for key, value in doc.items(): |
| 67 | + if key in rel_slots: |
| 68 | + for target in value if isinstance(value, list) else [value]: |
| 69 | + edges.append({"source": source, "predicate": key, "target": target}) |
| 70 | + elif key == "relations": # reified Relation objects |
| 71 | + for rel in value or []: |
| 72 | + edges.append({ |
| 73 | + "source": source, |
| 74 | + "predicate": rel.get("predicate"), |
| 75 | + "target": rel.get("target"), |
| 76 | + }) |
| 77 | + else: |
| 78 | + node[key] = _scalarize(value) |
| 79 | + return node, edges |
| 80 | + |
| 81 | + |
| 82 | +def to_frames(bundle: Bundle | str | pathlib.Path, engine: str = "pandas") -> dict: |
| 83 | + """Project *bundle* to ``{type_name: nodes_frame, "relations": edges_frame}``. |
| 84 | +
|
| 85 | + *bundle* may be a loaded :class:`~lokf.model.Bundle` or a path to a bundle |
| 86 | + directory. One frame per concept ``type`` holds that type's scalar fields; |
| 87 | + the ``relations`` frame holds every typed edge as ``(source, predicate, |
| 88 | + target)``. |
| 89 | + """ |
| 90 | + if not isinstance(bundle, Bundle): |
| 91 | + bundle = load_bundle(bundle) |
| 92 | + rel_slots = set(vocabulary().relation_slots) |
| 93 | + by_type: dict[str, list[dict]] = {} |
| 94 | + edges: list[dict] = [] |
| 95 | + for doc in bundle.docs(): |
| 96 | + node, doc_edges = _split(doc, rel_slots) |
| 97 | + by_type.setdefault(str(doc.get("type", "Concept")), []).append(node) |
| 98 | + edges.extend(doc_edges) |
| 99 | + frames = {name: _make_frame(rows, engine) for name, rows in by_type.items()} |
| 100 | + frames[RELATIONS_TABLE] = _make_frame(edges, engine) |
| 101 | + return frames |
| 102 | + |
| 103 | + |
| 104 | +# --------------------------------------------------------------------------- |
| 105 | +# Writers |
| 106 | +# --------------------------------------------------------------------------- |
| 107 | +def _to_pandas(df): |
| 108 | + return df.to_pandas() if hasattr(df, "to_pandas") else df |
| 109 | + |
| 110 | + |
| 111 | +def write_csv(frames: dict, outdir) -> pathlib.Path: |
| 112 | + """Write each frame to ``outdir/<name>.csv``.""" |
| 113 | + out = pathlib.Path(outdir) |
| 114 | + out.mkdir(parents=True, exist_ok=True) |
| 115 | + for name, df in frames.items(): |
| 116 | + if hasattr(df, "write_csv"): # polars |
| 117 | + df.write_csv(str(out / f"{name}.csv")) |
| 118 | + else: |
| 119 | + df.to_csv(out / f"{name}.csv", index=False) |
| 120 | + return out |
| 121 | + |
| 122 | + |
| 123 | +def write_parquet(frames: dict, outdir) -> pathlib.Path: |
| 124 | + """Write each frame to ``outdir/<name>.parquet`` (needs pyarrow).""" |
| 125 | + out = pathlib.Path(outdir) |
| 126 | + out.mkdir(parents=True, exist_ok=True) |
| 127 | + for name, df in frames.items(): |
| 128 | + if hasattr(df, "write_parquet"): # polars |
| 129 | + df.write_parquet(str(out / f"{name}.parquet")) |
| 130 | + else: |
| 131 | + df.to_parquet(out / f"{name}.parquet") |
| 132 | + return out |
| 133 | + |
| 134 | + |
| 135 | +def to_sqlite(frames: dict, path) -> pathlib.Path: |
| 136 | + """Write every frame as a table in a SQLite database at *path*.""" |
| 137 | + import sqlite3 |
| 138 | + |
| 139 | + path = pathlib.Path(path) |
| 140 | + con = sqlite3.connect(str(path)) |
| 141 | + try: |
| 142 | + for name, df in frames.items(): |
| 143 | + _to_pandas(df).to_sql(name, con, if_exists="replace", index=False) |
| 144 | + finally: |
| 145 | + con.close() |
| 146 | + return path |
| 147 | + |
| 148 | + |
| 149 | +# --------------------------------------------------------------------------- |
| 150 | +# External-table DDL (the lakehouse story) |
| 151 | +# --------------------------------------------------------------------------- |
| 152 | +_BQ_TYPE = { |
| 153 | + "object": "STRING", "string": "STRING", "str": "STRING", |
| 154 | + "int64": "INT64", "Int64": "INT64", |
| 155 | + "float64": "FLOAT64", "float": "FLOAT64", |
| 156 | + "bool": "BOOL", "boolean": "BOOL", |
| 157 | + "datetime64[ns]": "TIMESTAMP", |
| 158 | +} |
| 159 | +_ATHENA_TYPE = { |
| 160 | + "STRING": "string", "INT64": "bigint", "FLOAT64": "double", |
| 161 | + "BOOL": "boolean", "TIMESTAMP": "timestamp", |
| 162 | +} |
| 163 | + |
| 164 | + |
| 165 | +def _columns(df) -> list[tuple[str, str]]: |
| 166 | + """(column, BigQuery type) pairs inferred from a frame's dtypes.""" |
| 167 | + if hasattr(df, "dtypes") and hasattr(df.dtypes, "items"): # pandas |
| 168 | + items = [(c, str(t)) for c, t in df.dtypes.items()] |
| 169 | + else: # polars |
| 170 | + items = [(c, str(t).lower()) for c, t in df.schema.items()] |
| 171 | + return [(c, _BQ_TYPE.get(t, "STRING")) for c, t in items] |
| 172 | + |
| 173 | + |
| 174 | +def external_table_ddl(frames: dict, dialect: str = "bigquery", |
| 175 | + location: str = "gs://your-bucket/lokf", |
| 176 | + dataset: str = "lokf") -> str: |
| 177 | + """``CREATE EXTERNAL TABLE`` DDL registering the bundle's Parquet as a lakehouse. |
| 178 | +
|
| 179 | + Pair with :func:`write_parquet` (land the files at *location*), then run this |
| 180 | + DDL to expose one external table per concept type + the relations table. |
| 181 | + """ |
| 182 | + base = location.rstrip("/") |
| 183 | + stmts = [] |
| 184 | + for name, df in frames.items(): |
| 185 | + cols = _columns(df) |
| 186 | + if dialect == "bigquery": |
| 187 | + defs = ",\n ".join(f"`{c}` {t}" for c, t in cols) |
| 188 | + stmts.append( |
| 189 | + f"CREATE OR REPLACE EXTERNAL TABLE `{dataset}.{name}` (\n {defs}\n)\n" |
| 190 | + f"OPTIONS (format = 'PARQUET', uris = ['{base}/{name}.parquet']);" |
| 191 | + ) |
| 192 | + elif dialect == "athena": |
| 193 | + defs = ",\n ".join(f"`{c}` {_ATHENA_TYPE.get(t, 'string')}" for c, t in cols) |
| 194 | + stmts.append( |
| 195 | + f"CREATE EXTERNAL TABLE IF NOT EXISTS {name} (\n {defs}\n)\n" |
| 196 | + f"STORED AS PARQUET\nLOCATION '{base}/{name}/';" |
| 197 | + ) |
| 198 | + else: |
| 199 | + raise ValueError(f"unknown dialect: {dialect!r} (use 'bigquery' or 'athena')") |
| 200 | + return "\n\n".join(stmts) |
0 commit comments