# ============================================================ # SPID — DuckDB store # # Source corpus (filings / sections / companies) is exposed as # read-only views over the existing parquet. SPID writes its # own tables: mentions, projects, timeline, graph, embeddings. # ============================================================ from __future__ import annotations from pathlib import Path import duckdb from spid.config import Settings, get_settings from spid.extract.base import ProjectMention SCHEMA = """ CREATE TABLE IF NOT EXISTS project_mentions ( mention_id VARCHAR PRIMARY KEY, accession_number VARCHAR, cik VARCHAR, section_id VARCHAR, section_name VARCHAR, project_name VARCHAR, project_type VARCHAR, description VARCHAR, objective VARCHAR, amount_usd DOUBLE, amount_raw VARCHAR, locations VARCHAR[], technologies VARCHAR[], partners VARCHAR[], suppliers VARCHAR[], benefits VARCHAR[], risks VARCHAR[], years INTEGER[], status VARCHAR, confidence DOUBLE, backend VARCHAR, -- denormalised filing context (filled on load) ticker VARCHAR, company_name VARCHAR, form_type VARCHAR, filing_date DATE, fiscal_year INTEGER, sector VARCHAR ); CREATE TABLE IF NOT EXISTS projects ( project_id VARCHAR PRIMARY KEY, cik VARCHAR, ticker VARCHAR, company_name VARCHAR, sector VARCHAR, project_type VARCHAR, project_name VARCHAR, canonical_location VARCHAR, technologies VARCHAR[], status VARCHAR, total_amount_usd DOUBLE, n_mentions INTEGER, n_filings INTEGER, first_seen DATE, last_seen DATE, first_year INTEGER, last_year INTEGER, avg_confidence DOUBLE, description VARCHAR ); CREATE TABLE IF NOT EXISTS project_timeline ( project_id VARCHAR, accession_number VARCHAR, filing_date DATE, form_type VARCHAR, status VARCHAR, amount_usd DOUBLE, confidence DOUBLE, snippet VARCHAR ); CREATE TABLE IF NOT EXISTS kg_nodes ( node_id VARCHAR PRIMARY KEY, node_type VARCHAR, -- Company | Project | Technology | Location | Partner label VARCHAR, props JSON ); CREATE TABLE IF NOT EXISTS kg_edges ( src VARCHAR, rel VARCHAR, -- owns | uses | located_in | partners_with | requires_capital dst VARCHAR, props JSON ); CREATE TABLE IF NOT EXISTS project_embeddings ( project_id VARCHAR PRIMARY KEY, model VARCHAR, dim INTEGER, vector FLOAT[] ); """ class Store: def __init__(self, settings: Settings | None = None, read_only: bool = False): self.s = settings or get_settings() self.s.db_path.parent.mkdir(parents=True, exist_ok=True) self.con = duckdb.connect(str(self.s.db_path), read_only=read_only) if not read_only: self._init() def _init(self) -> None: self.con.execute( "CREATE TABLE IF NOT EXISTS spid_sections (" "section_id VARCHAR PRIMARY KEY, accession_number VARCHAR, section_name VARCHAR," "item_number VARCHAR, text VARCHAR, word_count INTEGER)") self._register_sources() self.con.execute(SCHEMA) def _register_sources(self) -> None: c = self.con c.execute( "CREATE TABLE IF NOT EXISTS extracted_sections (" "section_id VARCHAR PRIMARY KEY, accession_number VARCHAR, backend VARCHAR," "n_mentions INTEGER, extracted_at TIMESTAMP DEFAULT current_timestamp)") c.execute(f"CREATE OR REPLACE VIEW src_filings AS SELECT * FROM read_parquet('{self.s.filings_parquet}')") c.execute(f"CREATE OR REPLACE VIEW src_sections AS SELECT * FROM read_parquet('{self.s.sections_parquet}')") c.execute(f"CREATE OR REPLACE VIEW src_companies AS SELECT * FROM read_parquet('{self.s.companies_parquet}')") # unified, enriched section feed: pre-parsed sections (all forms since 2026-10-27) # UNION ingested 10-K items (legacy S&P 500 run); a 10-K already ingested in # spid_sections is not read a second time from the parsed corpus. scope = self.s.section_scope scope_sql = "" if scope: parts = [] for form, items in scope.items(): q = ",".join("'" + i + "'" for i in items) parts.append(f"(form_type = '{form}' AND section_name IN ({q}))") others = ",".join("'" + f + "'" for f in scope) scope_sql = " WHERE " + " OR ".join(parts) + f" OR form_type NOT IN ({others})" c.execute(f""" CREATE OR REPLACE VIEW v_sections AS WITH parsed AS ( SELECT s.section_id, s.accession_number, s.section_name, s.item_number, s.text, s.word_count, f.form_type FROM src_sections s JOIN src_filings f USING (accession_number) WHERE NOT (f.form_type = '10-K' AND s.accession_number IN (SELECT DISTINCT accession_number FROM spid_sections)) ), all_sec AS ( SELECT section_id, accession_number, section_name, item_number, text, word_count FROM parsed{scope_sql} UNION ALL SELECT section_id, accession_number, section_name, item_number, text, word_count FROM spid_sections ) SELECT s.section_id, s.accession_number, s.section_name, s.item_number, s.text, s.word_count, f.cik, f.ticker, f.company_name AS filing_company, f.form_type, CAST(f.filing_date AS DATE) AS filing_date, f.fiscal_year, co.company_name AS company_name, co.sector, co.sub_industry, co.sic_code, co.state FROM all_sec s JOIN src_filings f USING (accession_number) LEFT JOIN src_companies co ON f.cik = co.cik """) # ---- read side ---------------------------------------------------- def companies(self) -> list[tuple]: return self.con.execute( "SELECT DISTINCT cik, ticker FROM src_filings ORDER BY cik" ).fetchall() def sections_for(self, *, ciks: list[str] | None = None, forms: list[str] | None = None, limit: int | None = None, min_words: int = 30, skip_done: bool = True) -> list[dict]: where = [f"word_count >= {int(min_words)}"] if skip_done: where.append("section_id NOT IN (SELECT section_id FROM extracted_sections)") if ciks: quoted = ",".join("'" + c.replace("'", "") + "'" for c in ciks) where.append(f"cik IN ({quoted})") if forms: quoted = ",".join("'" + f.replace("'", "") + "'" for f in forms) where.append(f"form_type IN ({quoted})") sql = "SELECT * FROM v_sections WHERE " + " AND ".join(where) sql += " ORDER BY filing_date" if limit: sql += f" LIMIT {int(limit)}" res = self.con.execute(sql) table = res.to_arrow_table() if hasattr(res, "to_arrow_table") else res.fetch_arrow_table() # DuckDB >= 2.0 / < 2.0 return table.to_pylist() # ---- write side --------------------------------------------------- def upsert_mentions(self, mentions: list[ProjectMention], ctx: dict) -> int: """ctx maps accession_number -> filing context dict.""" if not mentions: return 0 rows = [] for m in mentions: cx = ctx.get(m.accession_number, {}) rows.append(( m.mention_id, m.accession_number, m.cik, m.section_id, m.section_name, m.project_name, m.project_type, m.description, m.objective, m.amount_usd, m.amount_raw, m.locations, m.technologies, m.partners, m.suppliers, m.benefits, m.risks, m.years, m.status, m.confidence, m.backend, cx.get("ticker"), cx.get("company_name") or cx.get("filing_company"), cx.get("form_type"), cx.get("filing_date"), cx.get("fiscal_year"), cx.get("sector"), )) self.con.executemany( "INSERT OR REPLACE INTO project_mentions VALUES " "(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", rows) return len(rows) def mark_extracted(self, rows: list[tuple]) -> None: """rows: (section_id, accession_number, backend, n_mentions).""" if rows: self.con.executemany( "INSERT OR REPLACE INTO extracted_sections (section_id, accession_number, backend, n_mentions) " "VALUES (?,?,?,?)", rows) def backfill_extracted(self, backend: str = "llm") -> int: """Record every section currently visible in v_sections as already extracted. Used once, before the Russell 3000 corpus is merged, so the legacy S&P 500 run (whose sections were processed without a ledger) is not re-billed. """ self.con.execute(f""" INSERT OR REPLACE INTO extracted_sections (section_id, accession_number, backend, n_mentions) SELECT v.section_id, v.accession_number, '{backend}', (SELECT COUNT(*) FROM project_mentions m WHERE m.section_id = v.section_id) FROM v_sections v """) return self.count("extracted_sections") def count(self, table: str) -> int: return self.con.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] def close(self) -> None: self.con.close()