| 8 |
8 |
""" |
| 9 |
9 |
from __future__ import annotations |
| 10 |
10 |
|
|
11 |
+import re |
| 11 |
12 |
from datetime import date |
| 12 |
13 |
from typing import Any |
| 13 |
14 |
|
| 126 |
127 |
raise ValueError("statement must start with SELECT / WITH") |
| 127 |
128 |
is_explain = bool(re.match(r"^\s*explain\b", sql, re.I)) |
| 128 |
129 |
con = _con() |
|
130 |
+ scoped_used: list[str] = [] |
| 129 |
131 |
dirs = [] |
| 130 |
132 |
for r in Lake().roots(): |
| 131 |
133 |
rr = str(r.resolve()).replace("'", "''") |
| 137 |
139 |
if name not in _referenced_views(sql, {name: spec}): |
| 138 |
140 |
continue # only materialize views the statement uses (directory globs are not free) |
| 139 |
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 |
150 |
elif isinstance(spec, list) and spec: # legacy explicit file list |
| 142 |
151 |
lst = ", ".join("'" + g.replace("'", "''") + "'" for g in spec) |
| 143 |
152 |
con.execute(f"CREATE VIEW {name} AS SELECT * FROM read_parquet([{lst}], union_by_name=true, hive_partitioning=false)") |
| 154 |
163 |
finally: |
| 155 |
164 |
timer.cancel() |
| 156 |
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 |
169 |
_views_cache: dict[str, Any] = {"at": 0.0, "value": None} |
| 191 |
200 |
if "delta" in layouts: # session files: derive the same columns so UNION ALL BY NAME lines up |
| 192 |
201 |
parts.append(f"SELECT *, symbol_as_reported AS symbol, year(session_date)::BIGINT AS year " |
| 193 |
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 |
205 |
_views_cache.update(at=time.monotonic(), value=out) |
| 196 |
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 |
262 |
def _referenced_views(sql: str, views: dict[str, Any]) -> list[str]: |
| 200 |
263 |
import re as _re |
| 201 |
264 |
low = sql.lower() |