| """Value lists of the catalog's small categorical columns, for resolving scope. |
| |
| Go's catalog keeps 3 sample values per column, so "COAL LOADER" and site codes like |
| "SDJ" are often invisible to the planner, and the planner resolved an invisible focus |
| value on 0 of 3 gate runs (2026-10-29) while visible ones passed. Code therefore reads |
| the full value list of each SMALL string column — once per process per catalog — |
| and matches the objective and the user's words against it. |
| |
| Read-only and through the normal query pipeline: one `retrieve_data` GROUP BY per |
| (table, column), validated and executed like any planner query (all five defense |
| layers apply), never SQL. Capped, time-boxed and never-throw. The result says whether |
| the load was COMPLETE: a focus resolved from a partial index is not pinned, and a |
| partial index is not cached (adversarial review 2026-10-29: a timed-out site list |
| otherwise pinned a focus without the site for the whole analysis). When Go emits full |
| value lists for low-cardinality columns (handoff), this becomes a catalog read. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import asyncio |
| import time |
| from collections.abc import Awaitable, Callable |
| from typing import Any |
|
|
| from src.middlewares.logging import get_logger |
|
|
| from ...catalog.models import Catalog, Column, Source, Table |
|
|
| logger = get_logger("scope_values") |
|
|
| |
| MAX_DISTINCT = 25 |
| |
| |
| TABLES_PER_COLUMN = 3 |
| MAX_QUERIES = 12 |
| |
| |
| |
| TIMEOUT_S = 20.0 |
| |
| |
| CACHE_TTL_S = 900.0 |
|
|
| Invoke = Callable[[str, dict[str, Any]], Awaitable[Any]] |
|
|
|
|
| class ValueIndex(dict): |
| """`{column name: [values]}` plus how it was loaded. |
| |
| `complete`: every queried column answered in time. `sources`: lower-cased column |
| name → the source ids whose lists were read, so a same-named column in another |
| source (a file's own `section`) is not scoped by values it never held. |
| """ |
|
|
| def __init__(self, *args: Any, complete: bool = True, **kwargs: Any) -> None: |
| super().__init__(*args, **kwargs) |
| self.complete = complete |
| self.sources: dict[str, set[str]] = {} |
|
|
|
|
| _CACHE: dict[tuple, tuple[float, ValueIndex]] = {} |
|
|
|
|
| def _small(col: Column) -> bool: |
| dc = col.stats.distinct_count if col.stats else None |
| return ( |
| col.data_type == "string" |
| and not col.pii_flag |
| and dc is not None |
| and 0 < dc <= MAX_DISTINCT |
| ) |
|
|
|
|
| def _targets(catalog: Catalog) -> list[tuple[Source, Table, Column]]: |
| """Every column NAME's best table first, then second-best tables, … — so the cap |
| can never starve a column of its only query (review: 4 columns on 3 tables each |
| used all 12 queries and `site_code` was never read).""" |
| by_name: dict[str, list[tuple[Source, Table, Column]]] = {} |
| for s in catalog.sources: |
| if s.source_type not in ("schema", "tabular"): |
| continue |
| for t in s.tables: |
| for c in t.columns: |
| if _small(c): |
| by_name.setdefault(c.name.lower(), []).append((s, t, c)) |
| for items in by_name.values(): |
| items.sort( |
| key=lambda x: (x[2].stats.distinct_count or 0, x[1].row_count or 0), |
| reverse=True, |
| ) |
| out: list[tuple[Source, Table, Column]] = [] |
| for rank in range(TABLES_PER_COLUMN): |
| out.extend(items[rank] for items in by_name.values() if len(items) > rank) |
| return out[:MAX_QUERIES] |
|
|
|
|
| async def _one(invoke: Invoke, s: Source, t: Table, c: Column) -> list[str] | None: |
| """The column's values, or None when the query failed.""" |
| ir = { |
| "source_id": s.source_id, |
| "table_id": t.table_id, |
| "select": [{"kind": "column", "column_id": c.column_id, "alias": c.name}], |
| "group_by": [c.column_id], |
| "limit": MAX_DISTINCT * 2, |
| } |
| out = await invoke("retrieve_data", {"ir": ir}) |
| if getattr(out, "kind", None) != "table": |
| return None |
| rows = getattr(out, "rows", None) or [] |
| return [str(r[0]) for r in rows if r and isinstance(r[0], str) and r[0].strip()] |
|
|
|
|
| def _cache_key(catalog: Catalog) -> tuple: |
| return ( |
| catalog.user_id, |
| tuple(sorted((s.source_id, str(s.updated_at)) for s in catalog.sources)), |
| ) |
|
|
|
|
| async def load_value_index(catalog: Catalog, invoke: Invoke) -> ValueIndex: |
| """The value lists of the small string columns. Never raises.""" |
| try: |
| key = _cache_key(catalog) |
| hit = _CACHE.get(key) |
| if hit is not None and time.monotonic() - hit[0] < CACHE_TTL_S: |
| return hit[1] |
| index = await _load(catalog, invoke) |
| if index and index.complete: |
| _CACHE[key] = (time.monotonic(), index) |
| return index |
| except Exception as e: |
| logger.warning( |
| "scope value index failed", degraded_seam="scope_values", error=repr(e) |
| ) |
| return ValueIndex(complete=False) |
|
|
|
|
| async def _load(catalog: Catalog, invoke: Invoke) -> ValueIndex: |
| targets = _targets(catalog) |
| if not targets: |
| return ValueIndex() |
| tasks = [asyncio.create_task(_one(invoke, s, t, c)) for s, t, c in targets] |
| done, pending = await asyncio.wait(tasks, timeout=TIMEOUT_S) |
| for p in pending: |
| p.cancel() |
| index = ValueIndex(complete=not pending) |
| for task, (s, _t, c) in zip(tasks, targets, strict=True): |
| if task not in done: |
| continue |
| if task.exception() is not None: |
| logger.info("scope value list failed", error=repr(task.exception())) |
| index.complete = False |
| continue |
| values = task.result() |
| if values is None: |
| index.complete = False |
| continue |
| bucket = index.setdefault(c.name, []) |
| for v in values: |
| if v not in bucket: |
| bucket.append(v) |
| index.sources.setdefault(c.name.lower(), set()).add(s.source_id) |
| for name in [k for k, v in index.items() if not v]: |
| del index[name] |
| logger.info( |
| "scope value index loaded", |
| columns=len(index), |
| queries=len(targets), |
| pending=len(pending), |
| complete=index.complete, |
| values=sum(len(v) for v in index.values()), |
| ) |
| return index |
|
|
