diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 3669b28e..3271bf4c 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -24,6 +24,8 @@ isort = "^5.12.0" pyright = "^1.1.324" numpy = "1.26.4" google-cloud-storage = "^3.1.0" +duckdb = "^1.5.4" +pyarrow = "^20.0.0" [tool.poetry.group.dev.dependencies] pyinstrument = "^4.6.1" diff --git a/backend/requirements.txt b/backend/requirements.txt index 62391faf..e697fd80 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -8,6 +8,7 @@ certifi==2025.1.31 charset-normalizer==3.4.1 click==8.1.8 colorama==0.4.6 +duckdb==1.5.4 exceptiongroup==1.2.2 fastapi==0.101.1 flake8==6.1.0 @@ -36,6 +37,7 @@ platformdirs==4.3.7 proto-plus==1.26.1 protobuf==6.30.2 psycopg2==2.9.10 +pyarrow==20.0.0 pyasn1-modules==0.4.2 pyasn1==0.6.1 pycodestyle==2.11.1 diff --git a/backend/src/api/backend.py b/backend/src/api/backend.py new file mode 100644 index 00000000..0d96f2a7 --- /dev/null +++ b/backend/src/api/backend.py @@ -0,0 +1,47 @@ +import os + +if os.environ.get("API_BACKEND", "").lower() == "duckdb": + from src.db_duckdb import ( + get_event, + get_events, + get_match, + get_matches, + get_team, + get_team_event, + get_team_events, + get_team_year, + get_team_years, + get_teams, + get_year, + get_years, + ) +else: + from src.db.read import ( + get_event, + get_events, + get_match, + get_matches, + get_team, + get_team_event, + get_team_events, + get_team_year, + get_team_years, + get_teams, + get_year, + get_years, + ) + +__all__ = [ + "get_event", + "get_events", + "get_match", + "get_matches", + "get_team_event", + "get_team_events", + "get_team_year", + "get_team_years", + "get_team", + "get_teams", + "get_year", + "get_years", +] diff --git a/backend/src/api/event.py b/backend/src/api/event.py index 5149c02c..a06bdb9f 100644 --- a/backend/src/api/event.py +++ b/backend/src/api/event.py @@ -16,7 +16,7 @@ year_query, ) from src.db.models import Event -from src.db.read import get_event, get_events +from src.api.backend import get_event, get_events from src.utils.alru_cache import alru_cache from src.utils.decorators import ( async_fail_gracefully_plural, diff --git a/backend/src/api/match.py b/backend/src/api/match.py index 3ac85d3f..f0e08bcc 100644 --- a/backend/src/api/match.py +++ b/backend/src/api/match.py @@ -15,7 +15,7 @@ year_query, ) from src.db.models import Match -from src.db.read import get_match, get_matches +from src.api.backend import get_match, get_matches from src.utils.alru_cache import alru_cache from src.utils.decorators import ( async_fail_gracefully_plural, diff --git a/backend/src/api/team.py b/backend/src/api/team.py index 32e68a86..68c483ed 100644 --- a/backend/src/api/team.py +++ b/backend/src/api/team.py @@ -14,7 +14,7 @@ state_query, ) from src.db.models import Team -from src.db.read import get_team, get_teams +from src.api.backend import get_team, get_teams from src.utils.alru_cache import alru_cache from src.utils.decorators import ( async_fail_gracefully_plural, diff --git a/backend/src/api/team_event.py b/backend/src/api/team_event.py index ad81153d..0f3e95ae 100644 --- a/backend/src/api/team_event.py +++ b/backend/src/api/team_event.py @@ -18,7 +18,7 @@ year_query, ) from src.db.models import TeamEvent -from src.db.read import get_team_event, get_team_events +from src.api.backend import get_team_event, get_team_events from src.utils.alru_cache import alru_cache from src.utils.decorators import ( async_fail_gracefully_plural, diff --git a/backend/src/api/team_year.py b/backend/src/api/team_year.py index 7db7bb21..12531d22 100644 --- a/backend/src/api/team_year.py +++ b/backend/src/api/team_year.py @@ -15,7 +15,7 @@ year_query, ) from src.db.models import TeamYear -from src.db.read import get_team_year, get_team_years +from src.api.backend import get_team_year, get_team_years from src.utils.alru_cache import alru_cache from src.utils.decorators import ( async_fail_gracefully_plural, diff --git a/backend/src/api/year.py b/backend/src/api/year.py index 0764a523..1dddd626 100644 --- a/backend/src/api/year.py +++ b/backend/src/api/year.py @@ -5,7 +5,7 @@ from src.api.query import ascending_query, limit_query, metric_query, offset_query from src.db.models import Year -from src.db.read import get_year, get_years +from src.api.backend import get_year, get_years from src.utils.alru_cache import alru_cache from src.utils.decorators import ( async_fail_gracefully_plural, diff --git a/backend/src/data/main.py b/backend/src/data/main.py index 7d959689..c0129385 100644 --- a/backend/src/data/main.py +++ b/backend/src/data/main.py @@ -39,6 +39,7 @@ update_team_years as update_team_years_db, update_teams as update_teams_db, ) +from src.google.parquet import build_parquet_uploads, write_parquet from src.google.snapshot import read_snapshot, write_snapshot from src.google.storage import write_objs as write_objs_storage @@ -83,7 +84,8 @@ def process_year( write_snapshot(year_num, objs, teams) timer.print(str(year_num) + " Write Snapshot") - write_objs_storage(objs, orig_objs if partial else None, teams) + parquet_uploads = build_parquet_uploads(year_num, objs, teams) + write_objs_storage(objs, orig_objs if partial else None, teams, parquet_uploads) timer.print(str(year_num) + " Write Storage") try: @@ -96,6 +98,10 @@ def process_year( write_objs_db(year_num, objs, orig_objs if partial else None, not partial) timer.print(str(year_num) + " Write DB") + if not DISABLE_GCS: + write_parquet(year_num, objs, teams) + timer.print(str(year_num) + " Write Parquet") + return teams diff --git a/backend/src/db_duckdb/__init__.py b/backend/src/db_duckdb/__init__.py new file mode 100644 index 00000000..e6c1f310 --- /dev/null +++ b/backend/src/db_duckdb/__init__.py @@ -0,0 +1,29 @@ +from src.db_duckdb.main import ( + get_event, + get_events, + get_match, + get_matches, + get_team, + get_team_event, + get_team_events, + get_team_year, + get_team_years, + get_teams, + get_year, + get_years, +) + +__all__ = [ + "get_event", + "get_events", + "get_match", + "get_matches", + "get_team_event", + "get_team_events", + "get_team_year", + "get_team_years", + "get_team", + "get_teams", + "get_year", + "get_years", +] diff --git a/backend/src/db_duckdb/main.py b/backend/src/db_duckdb/main.py new file mode 100644 index 00000000..1809b7a5 --- /dev/null +++ b/backend/src/db_duckdb/main.py @@ -0,0 +1,397 @@ +import os +import shutil +import tempfile +import threading +import time +from typing import Any, Dict, List, Optional, Tuple, Type + +import duckdb + +from src.db.models.main import Model, ModelORM +from src.db_duckdb.schema import PARQUET_PREFIX, SPECS, columns, from_row + +SYNC_TTL = float(os.environ.get("DUCKDB_SYNC_TTL", "30")) + +_con: Optional[duckdb.DuckDBPyConnection] = None +_cache_root: Optional[str] = None +_current_dir: Optional[str] = None +_prev_dir: Optional[str] = None +_keys: Dict[str, str] = {} +_manifest_generation: Optional[int] = None +_last_sync: float = 0.0 +_sync_lock = threading.Lock() + + +def _connection() -> duckdb.DuckDBPyConnection: + global _con + if _con is None: + _con = duckdb.connect() + return _con + + +def _sync() -> str: + global _cache_root, _current_dir, _prev_dir, _keys + global _manifest_generation, _last_sync + + now = time.monotonic() + if _current_dir is not None and _last_sync and now - _last_sync < SYNC_TTL: + return _current_dir + + with _sync_lock: + now = time.monotonic() + if _current_dir is not None and _last_sync and now - _last_sync < SYNC_TTL: + return _current_dir + + from src.google.publish import MANIFEST_OBJECT, Manifest + from src.google.storage import _bucket + + if _cache_root is None: + _cache_root = tempfile.mkdtemp(prefix="duckdb_parquet_") + + blob = _bucket().blob(MANIFEST_OBJECT) + try: + blob.reload() + except Exception: + _last_sync = now + if _current_dir is None: + _current_dir = tempfile.mkdtemp(prefix="gen_", dir=_cache_root) + return _current_dir + + if blob.generation == _manifest_generation and _current_dir is not None: + _last_sync = now + return _current_dir + + manifest = Manifest.from_json(blob.download_as_bytes()) + entries = { + logical: key + for logical, key in manifest.blobs.items() + if logical.startswith(PARQUET_PREFIX + "/") + } + + gen_dir = tempfile.mkdtemp(prefix="gen_", dir=_cache_root) + for logical, key in entries.items(): + dest = os.path.join(gen_dir, logical) + os.makedirs(os.path.dirname(dest), exist_ok=True) + prev = os.path.join(_current_dir, logical) if _current_dir else None + if _keys.get(logical) == key and prev is not None and os.path.exists(prev): + os.link(prev, dest) + else: + _bucket().blob(key).download_to_filename(dest) + + old_prev = _prev_dir + _prev_dir = _current_dir + _current_dir = gen_dir + _keys = entries + _manifest_generation = blob.generation + _last_sync = now + if old_prev is not None: + shutil.rmtree(old_prev, ignore_errors=True) + return _current_dir + + +def _years(base: str) -> List[int]: + root = os.path.join(base, PARQUET_PREFIX) + if not os.path.isdir(root): + return [] + return sorted(int(y) for y in os.listdir(root) if y.isdigit()) + + +def _source(base: str, table: str, year: Optional[int] = None) -> str: + if year is None: + return f"read_parquet('{base}/{PARQUET_PREFIX}/*/{table}.parquet')" + return f"read_parquet('{base}/{PARQUET_PREFIX}/{year}/{table}.parquet')" + + +def _teams_source(base: str) -> str: + years = _years(base) + if not years: + # No parquet yet (first db-less deploy): glob resolves to no files, which + # _query degrades to an empty result instead of crashing on max([]). + return f"read_parquet('{base}/{PARQUET_PREFIX}/*/teams.parquet')" + return f"read_parquet('{base}/{PARQUET_PREFIX}/{max(years)}/teams.parquet')" + + +def _query( + model_cls: Type[Model], + orm_type: Type[ModelORM], + source: str, + where: List[str], + params: List[Any], + metric: Optional[str] = None, + ascending: Optional[bool] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, +) -> List[Model]: + cols = columns(orm_type) + sql = f"SELECT * FROM {source}" + clauses = list(where) + if metric is not None: + if metric not in {c[0] for c in cols}: + raise ValueError(f"invalid metric: {metric}") + clauses.append(f'"{metric}" IS NOT NULL') + if clauses: + sql += " WHERE " + " AND ".join(clauses) + if metric is not None: + direction = "ASC" if ascending else "DESC" + sql += f' ORDER BY "{metric}" {direction}' + if limit is not None: + sql += f" LIMIT {int(limit)}" + if offset is not None: + sql += f" OFFSET {int(offset)}" + + cursor = _connection().cursor() + cursor.execute(sql, params) + names = [d[0] for d in cursor.description] + return [ + from_row(model_cls, cols, dict(zip(names, row))) for row in cursor.fetchall() + ] + + +def _one( + model_cls: Type[Model], + orm_type: Type[ModelORM], + source: str, + where: List[str], + params: List[Any], +) -> Optional[Model]: + rows = _query(model_cls, orm_type, source, where, params, limit=1) + return rows[0] if rows else None + + +def get_team(team: int) -> Optional[Model]: + model, orm = SPECS["teams"] + return _one(model, orm, _teams_source(_sync()), ['"team" = ?'], [team]) + + +def get_teams( + country: Optional[str] = None, + state: Optional[str] = None, + district: Optional[str] = None, + active: Optional[bool] = None, + metric: Optional[str] = None, + ascending: Optional[bool] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, +) -> List[Model]: + model, orm = SPECS["teams"] + where, params = _eq( + {"country": country, "state": state, "district": district, "active": active} + ) + return _query( + model, + orm, + _teams_source(_sync()), + where, + params, + metric, + ascending, + limit, + offset, + ) + + +def get_year(year: int) -> Optional[Model]: + model, orm = SPECS["year"] + return _one(model, orm, _source(_sync(), "year", year), ['"year" = ?'], [year]) + + +def get_years( + metric: Optional[str] = None, + ascending: Optional[bool] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, +) -> List[Model]: + model, orm = SPECS["year"] + return _query( + model, orm, _source(_sync(), "year"), [], [], metric, ascending, limit, offset + ) + + +def get_team_year(team: int, year: int) -> Optional[Model]: + model, orm = SPECS["team_years"] + return _one( + model, + orm, + _source(_sync(), "team_years", year), + ['"team" = ?', '"year" = ?'], + [team, year], + ) + + +def get_team_years( + team: Optional[int] = None, + teams: Optional[List[str]] = None, + year: Optional[int] = None, + country: Optional[str] = None, + state: Optional[str] = None, + district: Optional[str] = None, + metric: Optional[str] = None, + ascending: Optional[bool] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, +) -> List[Model]: + model, orm = SPECS["team_years"] + where, params = _eq( + { + "team": team, + "year": year, + "country": country, + "state": state, + "district": district, + } + ) + if teams is not None: + where.append(f'"team" IN ({", ".join("?" for _ in teams)})') + params.extend(teams) + return _query( + model, + orm, + _source(_sync(), "team_years", year), + where, + params, + metric, + ascending, + limit, + offset, + ) + + +def get_event(event_id: str) -> Optional[Model]: + model, orm = SPECS["events"] + return _one(model, orm, _source(_sync(), "events"), ['"key" = ?'], [event_id]) + + +def get_events( + year: Optional[int] = None, + country: Optional[str] = None, + state: Optional[str] = None, + district: Optional[str] = None, + type: Optional[str] = None, + week: Optional[int] = None, + metric: Optional[str] = None, + ascending: Optional[bool] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, +) -> List[Model]: + model, orm = SPECS["events"] + where, params = _eq( + { + "year": year, + "country": country, + "state": state, + "district": district, + "type": type, + "week": week, + } + ) + return _query( + model, + orm, + _source(_sync(), "events", year), + where, + params, + metric, + ascending, + limit, + offset, + ) + + +def get_team_event(team: int, event: str) -> Optional[Model]: + model, orm = SPECS["team_events"] + return _one( + model, + orm, + _source(_sync(), "team_events"), + ['"team" = ?', '"event" = ?'], + [team, event], + ) + + +def get_team_events( + team: Optional[int] = None, + teams: Optional[List[str]] = None, + year: Optional[int] = None, + event: Optional[str] = None, + country: Optional[str] = None, + state: Optional[str] = None, + district: Optional[str] = None, + type: Optional[str] = None, + week: Optional[int] = None, + metric: Optional[str] = None, + ascending: Optional[bool] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, +) -> List[Model]: + model, orm = SPECS["team_events"] + where, params = _eq( + { + "team": team, + "year": year, + "event": event, + "country": country, + "state": state, + "district": district, + "type": type, + "week": week, + } + ) + if teams is not None: + where.append(f'"team" IN ({", ".join("?" for _ in teams)})') + params.extend(teams) + return _query( + model, + orm, + _source(_sync(), "team_events", year), + where, + params, + metric, + ascending, + limit, + offset, + ) + + +def get_match(match: str) -> Optional[Model]: + model, orm = SPECS["matches"] + return _one(model, orm, _source(_sync(), "matches"), ['"key" = ?'], [match]) + + +def get_matches( + team: Optional[int] = None, + year: Optional[int] = None, + event: Optional[str] = None, + week: Optional[int] = None, + elim: Optional[bool] = None, + metric: Optional[str] = None, + ascending: Optional[bool] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, +) -> List[Model]: + model, orm = SPECS["matches"] + where, params = _eq({"year": year, "event": event, "week": week, "elim": elim}) + if team is not None: + cols = ["red_1", "red_2", "red_3", "blue_1", "blue_2", "blue_3"] + where.append("(" + " OR ".join(f'"{c}" = ?' for c in cols) + ")") + params.extend([team] * len(cols)) + return _query( + model, + orm, + _source(_sync(), "matches", year), + where, + params, + metric, + ascending, + limit, + offset, + ) + + +def _eq(filters: Dict[str, Any]) -> Tuple[List[str], List[Any]]: + where: List[str] = [] + params: List[Any] = [] + for name, value in filters.items(): + if value is not None: + where.append(f'"{name}" = ?') + params.append(value) + return where, params diff --git a/backend/src/db_duckdb/schema.py b/backend/src/db_duckdb/schema.py new file mode 100644 index 00000000..8d2dd811 --- /dev/null +++ b/backend/src/db_duckdb/schema.py @@ -0,0 +1,79 @@ +import json +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple, Type + +from sqlalchemy import JSON, Boolean, Float, Integer, inspect +from sqlalchemy.sql.sqltypes import Enum as SQLEnum + +from src.db.models import Event, Match, Team, TeamEvent, TeamYear, Year +from src.db.models.event import EventORM +from src.db.models.main import Model, ModelORM +from src.db.models.match import MatchORM +from src.db.models.team import TeamORM +from src.db.models.team_event import TeamEventORM +from src.db.models.team_year import TeamYearORM +from src.db.models.year import YearORM + +PARQUET_PREFIX = "parquet" + +SPECS: Dict[str, Tuple[Type[Model], Type[ModelORM]]] = { + "team_years": (TeamYear, TeamYearORM), + "events": (Event, EventORM), + "team_events": (TeamEvent, TeamEventORM), + "matches": (Match, MatchORM), + "teams": (Team, TeamORM), + "year": (Year, YearORM), +} + +Column = Tuple[str, str, Optional[Type[Enum]]] + + +def _kind(column: Any) -> Tuple[str, Optional[Type[Enum]]]: + type_ = column.type + if isinstance(type_, SQLEnum) and type_.enum_class is not None: + return "enum", type_.enum_class + if isinstance(type_, JSON): + return "json", None + if isinstance(type_, Boolean): + return "bool", None + if isinstance(type_, Integer): + return "int", None + if isinstance(type_, Float): + return "float", None + return "str", None + + +def columns(orm_type: Type[ModelORM]) -> List[Column]: + out: List[Column] = [] + for column in inspect(orm_type).columns: + kind, enum_cls = _kind(column) + out.append((column.name, kind, enum_cls)) + return out + + +def to_row(obj: Model, cols: List[Column]) -> Tuple[Any, ...]: + values: List[Any] = [] + for name, kind, _ in cols: + value = getattr(obj, name, None) + if value is None: + values.append(None) + elif kind == "enum": + values.append(value.value if isinstance(value, Enum) else value) + elif kind == "json": + values.append(json.dumps(value)) + else: + values.append(value) + return tuple(values) + + +def from_row(model_cls: Type[Model], cols: List[Column], row: Dict[str, Any]) -> Model: + data: Dict[str, Any] = {} + for name, kind, enum_cls in cols: + value = row.get(name) + if value is not None: + if kind == "json": + value = json.loads(value) + elif kind == "enum" and enum_cls is not None: + value = enum_cls(value) + data[name] = value + return model_cls.from_dict(data) diff --git a/backend/src/google/parquet.py b/backend/src/google/parquet.py new file mode 100644 index 00000000..8f111378 --- /dev/null +++ b/backend/src/google/parquet.py @@ -0,0 +1,90 @@ +import io +from typing import Any, Dict, List, Type + +import pyarrow as pa +import pyarrow.parquet as pq + +from src.data.utils import objs_type +from src.db.models import Team +from src.db.models.event import EventORM +from src.db.models.main import Model, ModelORM +from src.db.models.match import MatchORM +from src.db.models.team import TeamORM +from src.db.models.team_event import TeamEventORM +from src.db.models.team_year import TeamYearORM +from src.db.models.year import YearORM +from src.db_duckdb.schema import PARQUET_PREFIX, columns, to_row +from src.google.publish import content_hash, versioned_key +from src.google.storage import IMMUTABLE_CACHE, _bucket, read_manifest, write_manifest + +ARROW_TYPES = { + "int": pa.int64(), + "float": pa.float64(), + "bool": pa.bool_(), + "str": pa.string(), + "enum": pa.string(), + "json": pa.string(), +} + + +def parquet_logical(year: int, table: str) -> str: + return f"{PARQUET_PREFIX}/{year}/{table}.parquet" + + +def serialize_table(objs: List[Model], orm_type: Type[ModelORM]) -> bytes: + cols = columns(orm_type) + rows = [to_row(o, cols) for o in sorted(objs, key=lambda o: o.pk())] + columns_data = list(zip(*rows)) if rows else [() for _ in cols] + arrays = [ + pa.array(list(values), type=ARROW_TYPES[kind]) + for (_, kind, _), values in zip(cols, columns_data) + ] + table = pa.Table.from_arrays(arrays, names=[name for name, _, _ in cols]) + buf = io.BytesIO() + pq.write_table(table, buf) + return buf.getvalue() + + +def build_parquet_uploads( + year: int, objs: objs_type, teams: List[Team] +) -> Dict[str, bytes]: + sources: List[Any] = [ + ("team_years", list(objs[1].values()), TeamYearORM), + ("events", list(objs[2].values()), EventORM), + ("team_events", list(objs[3].values()), TeamEventORM), + ("matches", list(objs[4].values()), MatchORM), + ("teams", teams, TeamORM), + ("year", [objs[0]], YearORM), + ] + return { + parquet_logical(year, table): serialize_table(rows, orm_type) + for table, rows, orm_type in sources + } + + +def write_parquet(year: int, objs: objs_type, teams: List[Team]) -> None: + # Standalone path for historical (non-current) years, which publish parquet + # without a co-occurring site-blob publish. The current-year cycle folds parquet + # into the single site manifest write (see storage.write_objs). + manifest = read_manifest() + if manifest is None: + print("skipped parquet manifest update: manifest unavailable") + return + + bucket = _bucket() + blobs = dict(manifest.blobs) + changed = False + for logical, data in build_parquet_uploads(year, objs, teams).items(): + digest = content_hash(data) + if manifest.hash_for(logical) == digest: + continue + key = versioned_key(logical, digest) + blob = bucket.blob(key) + blob.cache_control = IMMUTABLE_CACHE + blob.upload_from_string(data, "application/octet-stream") + blobs[logical] = key + changed = True + + if changed: + manifest.blobs = blobs + write_manifest(manifest, bucket) diff --git a/backend/src/google/storage.py b/backend/src/google/storage.py index 53822478..9f75e9ad 100644 --- a/backend/src/google/storage.py +++ b/backend/src/google/storage.py @@ -20,8 +20,10 @@ VERSION_PREFIX, Manifest, UploadPlan, + content_hash, historical_key, plan_uploads, + versioned_key, ) from src.site.event import _read_all_events, _read_events, _read_event from src.site.match import _read_noteworthy_matches, _read_upcoming_matches @@ -117,6 +119,7 @@ def write_objs( objs: objs_type, orig_objs: Optional[objs_type] = None, teams: Optional[List[Team]] = None, + parquet: Optional[Dict[str, bytes]] = None, ) -> None: year = CURR_YEAR year_obj = objs[0] @@ -246,6 +249,17 @@ def add(object_name: str, data: Any) -> None: cycle = datetime.now(timezone.utc).isoformat() plan = plan_uploads(rendered, prev, cycle, hist_epoch=HIST_EPOCH) + if parquet: + # Fold parquet into the same manifest so the cycle writes it exactly once, + # last: DuckDB and the frontend never disagree by a cycle, and a partial + # read can never clobber the site-blob references. + prev_manifest = prev or Manifest() + for logical, data in parquet.items(): + digest = content_hash(data) + key = versioned_key(logical, digest) + if prev_manifest.hash_for(logical) != digest: + plan.uploads[key] = data + plan.manifest.blobs[logical] = key _publish(plan) return