File size: 6,767 Bytes
111ef6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """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")
#: A column is listed only when the catalog says it has at most this many values.
MAX_DISTINCT = 25
#: Tables queried per column name (largest first): `section` sits on 15 tables with
#: 4-6 values each; three cover the union without 15 queries.
TABLES_PER_COLUMN = 3
MAX_QUERIES = 12
#: Measured on the production source (2026-10-29): 9.7 s for the first query on a
#: cold connection, ~2 s each after. The load runs alongside the planner call
#: (5-20 s), so this cap rarely adds wait.
TIMEOUT_S = 20.0
#: Per-process cache keyed by the catalog's sources and `updated_at`; the TTL bounds
#: staleness of rows Go does not re-catalog. Only complete loads are cached.
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: # a partial or failed load is never cached
_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
|