| 8 | """ |
8 | """ |
| 9 | from __future__ import annotations |
9 | from __future__ import annotations |
| 10 | |
10 | |
| |
11 | +import re |
| 11 | from datetime import date |
12 | from datetime import date |
| 12 | from typing import Any |
13 | from typing import Any |
| 13 | |
14 | |
| 126 | raise ValueError("statement must start with SELECT / WITH") |
127 | raise ValueError("statement must start with SELECT / WITH") |
| 127 | is_explain = bool(re.match(r"^\s*explain\b", sql, re.I)) |
128 | is_explain = bool(re.match(r"^\s*explain\b", sql, re.I)) |
| 128 | con = _con() |
129 | con = _con() |
| |
130 | + scoped_used: list[str] = [] |
| 129 | dirs = [] |
131 | dirs = [] |
| 130 | for r in Lake().roots(): |
132 | for r in Lake().roots(): |
| 131 | rr = str(r.resolve()).replace("'", "''") |
133 | rr = str(r.resolve()).replace("'", "''") |
| 137 | if name not in _referenced_views(sql, {name: spec}): |
139 | if name not in _referenced_views(sql, {name: spec}): |
| 138 | continue # only materialize views the statement uses (directory globs are not free) |
140 | continue # only materialize views the statement uses (directory globs are not free) |
| 139 | if isinstance(spec, dict) and spec.get("sql"): |
141 | if isinstance(spec, dict) and spec.get("sql"): |
| 140 | − con.execute(f"CREATE VIEW {name} AS {spec['sql']}") |
142 | + fast = None |
| |
143 | + if spec.get("dataset_id"): |
| |
144 | + try: |
| |
145 | + fast = scoped_view_sql(spec["dataset_id"], spec.get("kind", ""), extract_scope(sql), Lake()) |
| |
146 | + except Exception: |
| |
147 | + fast = None |
| |
148 | + con.execute(f"CREATE VIEW {name} AS {fast or spec['sql']}") |
| |
149 | + scoped_used.append(name if fast else f"{name} (full scan)") |
| 141 | elif isinstance(spec, list) and spec: # legacy explicit file list |
150 | elif isinstance(spec, list) and spec: # legacy explicit file list |
| 142 | lst = ", ".join("'" + g.replace("'", "''") + "'" for g in spec) |
151 | lst = ", ".join("'" + g.replace("'", "''") + "'" for g in spec) |
| 143 | con.execute(f"CREATE VIEW {name} AS SELECT * FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false)") |
152 | con.execute(f"CREATE VIEW {name} AS SELECT * FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false)") |
| 154 | finally: |
163 | finally: |
| 155 | timer.cancel() |
164 | timer.cancel() |
| 156 | truncated = tbl.num_rows > max_rows |
165 | truncated = tbl.num_rows > max_rows |
| 157 | − return tbl.slice(0, max_rows), {"truncated": truncated} |
166 | + return tbl.slice(0, max_rows), {"truncated": truncated, "views": scoped_used} |
| 158 | |
167 | |
| 159 | |
168 | |
| 160 | _views_cache: dict[str, Any] = {"at": 0.0, "value": None} |
169 | _views_cache: dict[str, Any] = {"at": 0.0, "value": None} |
| 191 | if "delta" in layouts: # session files: derive the same columns so UNION ALL BY NAME lines up |
200 | if "delta" in layouts: # session files: derive the same columns so UNION ALL BY NAME lines up |
| 192 | parts.append(f"SELECT *, symbol_as_reported AS symbol, year(session_date)::BIGINT AS year " |
201 | parts.append(f"SELECT *, symbol_as_reported AS symbol, year(session_date)::BIGINT AS year " |
| 193 | f"FROM read_parquet('{layouts['delta']}', hive_partitioning=false)") |
202 | f"FROM read_parquet('{layouts['delta']}', hive_partitioning=false)") |
| 194 | − out[d["id"].replace(".", "_")] = {"sql": " UNION ALL BY NAME ".join(parts), "layouts": sorted(layouts), "root": str(root), "kind": d["kind"]} |
203 | + out[d["id"].replace(".", "_")] = {"sql": " UNION ALL BY NAME ".join(parts), "layouts": sorted(layouts), "root": str(root), |
| |
204 | + "kind": d["kind"], "dataset_id": d["id"]} |
| 195 | _views_cache.update(at=time.monotonic(), value=out) |
205 | _views_cache.update(at=time.monotonic(), value=out) |
| 196 | return out |
206 | return out |
| 197 | |
207 | |
| 198 | |
208 | |
| |
209 | +_SYM_RE = re.compile(r"\b(?:symbol|symbol_as_reported|underlying)\s*=\s*'([A-Za-z0-9.\-_]+)'", re.I) |
| |
210 | +_SYM_IN_RE = re.compile(r"\b(?:symbol|symbol_as_reported|underlying)\s+IN\s*\(([^)]+)\)", re.I) |
| |
211 | +_YEAR_RE = re.compile(r"\byear\s*=\s*(\d{4})\b", re.I) |
| |
212 | +_YEAR_IN_RE = re.compile(r"\byear\s+IN\s*\(([^)]+)\)", re.I) |
| |
213 | +_DATE_RE = re.compile(r"\b(?:session_date|trade_date)\s*(>=|<=|=|>|<)\s*(?:DATE\s*)?'(\d{4}-\d{2}-\d{2})'", re.I) |
| |
214 | +_BETWEEN_RE = re.compile(r"\b(?:session_date|trade_date)\s+BETWEEN\s+(?:DATE\s*)?'(\d{4}-\d{2}-\d{2})'\s+AND\s+(?:DATE\s*)?'(\d{4}-\d{2}-\d{2})'", re.I) |
| |
215 | + |
| |
216 | + |
| |
217 | +def extract_scope(sql: str) -> dict[str, Any]: |
| |
218 | + """Pull simple equality/IN/range predicates on partition columns out of a statement so that the file list can |
| |
219 | + be resolved from the catalog (indexed) instead of globbing hundreds of thousands of files.""" |
| |
220 | + from datetime import date as _date |
| |
221 | + syms: set[str] = {m.group(1).upper() for m in _SYM_RE.finditer(sql)} |
| |
222 | + for m in _SYM_IN_RE.finditer(sql): |
| |
223 | + syms |= {x.strip().strip("'\"").upper() for x in m.group(1).split(",") if x.strip()} |
| |
224 | + years: set[int] = {int(m.group(1)) for m in _YEAR_RE.finditer(sql)} |
| |
225 | + for m in _YEAR_IN_RE.finditer(sql): |
| |
226 | + years |= {int(x) for x in re.findall(r"\d{4}", m.group(1))} |
| |
227 | + start = end = None |
| |
228 | + for op, d in _DATE_RE.findall(sql): |
| |
229 | + dd = _date.fromisoformat(d) |
| |
230 | + if op in (">=", ">", "="): |
| |
231 | + start = dd if start is None else max(start, dd) |
| |
232 | + if op in ("<=", "<", "="): |
| |
233 | + end = dd if end is None else min(end, dd) |
| |
234 | + for a, b in _BETWEEN_RE.findall(sql): |
| |
235 | + start = max(start, _date.fromisoformat(a)) if start else _date.fromisoformat(a) |
| |
236 | + end = min(end, _date.fromisoformat(b)) if end else _date.fromisoformat(b) |
| |
237 | + if years: |
| |
238 | + ys, ye = _date(min(years), 1, 1), _date(max(years), 12, 31) |
| |
239 | + start = ys if start is None else max(start, ys) |
| |
240 | + end = ye if end is None else min(end, ye) |
| |
241 | + return {"symbols": sorted(syms), "start": start, "end": end} |
| |
242 | + |
| |
243 | + |
| |
244 | +def scoped_view_sql(dataset_id: str, kind: str, scope: dict[str, Any], lake: Lake) -> str | None: |
| |
245 | + """Explicit file list from the catalog for symbol-scoped statements (fast path). None → use the glob view.""" |
| |
246 | + if not scope["symbols"]: |
| |
247 | + return None |
| |
248 | + with connection() as conn: |
| |
249 | + files = live_files(conn, dataset_id, scope["symbols"], scope["start"], scope["end"]) |
| |
250 | + if len(files) > 20000: |
| |
251 | + return None |
| |
252 | + if not files: |
| |
253 | + return "SELECT * FROM (VALUES (NULL)) t(empty) WHERE false" |
| |
254 | + lst = ", ".join("'" + str(lake.resolve(f["path"])).replace("'", "''") + "'" for f in files) |
| |
255 | + if kind == "options": |
| |
256 | + return (f"SELECT * EXCLUDE (filename), regexp_extract(filename, 'period=([^/]+)', 1) AS period " |
| |
257 | + f"FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false, filename=true)") |
| |
258 | + return (f"SELECT *, symbol_as_reported AS symbol, year(session_date)::BIGINT AS year " |
| |
259 | + f"FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false)") |
| |
260 | + |
| |
261 | + |
| 199 | def _referenced_views(sql: str, views: dict[str, Any]) -> list[str]: |
262 | def _referenced_views(sql: str, views: dict[str, Any]) -> list[str]: |
| 200 | import re as _re |
263 | import re as _re |
| 201 | low = sql.lower() |
264 | low = sql.lower() |