Spaces:
Sleeping
Sleeping
File size: 2,251 Bytes
e85a672 | 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 | """Build the Space's read-only DuckDB from the published daily Parquet files."""
from __future__ import annotations
import argparse
import os
from pathlib import Path
import duckdb
from huggingface_hub import hf_hub_download
DEFAULT_REPO = "egoh33/ibkr-daily-stock-data"
FILES = ("polygon_bars.parquet", "silver_stock_features.parquet")
def materialize_database(files: dict[str, Path], db_path: Path) -> None:
"""Atomically create a DuckDB containing the two public daily-stock tables."""
db_path.parent.mkdir(parents=True, exist_ok=True)
temp_path = db_path.with_suffix(db_path.suffix + ".tmp")
temp_path.unlink(missing_ok=True)
with duckdb.connect(str(temp_path)) as conn:
for filename, source in files.items():
table = Path(filename).stem
safe_source = source.as_posix().replace("'", "''")
conn.execute(f"CREATE TABLE {table} AS SELECT * FROM read_parquet('{safe_source}')")
conn.execute("CHECKPOINT")
db_path.unlink(missing_ok=True)
temp_path.replace(db_path)
def download_files(repo_id: str, token: str | None = None) -> dict[str, Path]:
return {
filename: Path(
hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename, token=token)
)
for filename in FILES
}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--repo", default=os.getenv("HF_DATASET_REPO", DEFAULT_REPO))
parser.add_argument("--db-path", default=os.getenv("DB_PATH", "./data/equity.duckdb"))
parser.add_argument("--source-dir", help="Use local Parquet files instead of downloading")
args = parser.parse_args()
if args.source_dir:
source_dir = Path(args.source_dir)
files = {filename: source_dir / filename for filename in FILES}
else:
files = download_files(args.repo, token=os.getenv("HF_TOKEN"))
missing = [str(path) for path in files.values() if not path.is_file()]
if missing:
raise FileNotFoundError(f"Missing required daily-stock files: {', '.join(missing)}")
materialize_database(files, Path(args.db_path))
print(f"Materialized {args.db_path} from {args.repo}")
if __name__ == "__main__":
main()
|