spb/hfmarketdata

1"""DuckDB query plans: top-N-per-group queries must scan the Parquet files ONCE, with the datetime filter pushed2into the scan (no late-materialization re-scan), and the process must hold a single DuckDB database."""3from __future__ import annotations45import threading67import pytest8910@pytest.fixture11def paths(app):12    import main13    idx = main.ticker_index("stock", "1min", "UNADJUSTED")14    return [idx[t] for t in ("AAPL", "MSFT", "SMCP")]151617def _plan(sql, params):18    from core.duck import con, explain19    return explain(con(), sql, params)202122def test_topn_plan_single_filtered_scan(paths):23    from core.duck import scan_count, topn_per_group24    lo, hi = "2025-06-30 09:30:00", "2025-06-30 09:34:59"25    sql = topn_per_group(None, paths, "datetime >= ? AND datetime <= ?", [lo, hi], "ticker", "datetime DESC", 100,26                         order_by="ticker, datetime DESC", limit=300)27    plan = _plan(sql, [paths, lo, hi, 100, 300])28    assert scan_count(plan) == 1, plan29    assert "Filters" in plan and "datetime>=" in plan.replace(" ", "") , plan30    assert "HASH_JOIN" not in plan, plan313233def test_naive_qualify_would_double_scan(paths):34    """Documents WHY the helper exists: the bare form is rewritten into two scans + a join."""35    from core.duck import scan_count36    lo, hi = "2025-06-30 09:30:00", "2025-06-30 09:34:59"37    sql = ("SELECT * FROM read_parquet(?) WHERE datetime >= ? AND datetime <= ? "38           "QUALIFY row_number() OVER (PARTITION BY ticker ORDER BY datetime DESC) <= ? ORDER BY ticker, datetime DESC LIMIT ?")39    plan = _plan(sql, [paths, lo, hi, 100, 300])40    assert scan_count(plan) >= 2 or "HASH_JOIN" in plan, plan414243def test_snapshot_plan_single_scan(paths):44    from core.duck import scan_count, topn_per_group45    sql = topn_per_group(None, paths, "datetime <= ?", ["2025-06-30 10:00:00"], "ticker", "datetime DESC", 1,46                         order_by="ticker", limit=3)47    plan = _plan(sql, [paths, "2025-06-30 10:00:00", 1, 3])48    assert scan_count(plan) == 1 and "HASH_JOIN" not in plan, plan495051def test_futures_merged_plan_single_scan_per_file(app):52    from core.duck import scan_count53    from futures.lake import files_for, merged_sql54    files = files_for("ESH25", "1day")55    assert len(files) == 2, files          # update + archive overlap for 2025 contracts (daily: both cover 2024-12)56    sql = f"SELECT * EXCLUDE (ticker) FROM ({merged_sql(files, where='datetime >= ? AND datetime < ?')}) ORDER BY datetime LIMIT ?"57    plan = _plan(sql, [*files, "2024-12-01", "2024-12-31", 5001])58    assert scan_count(plan) == 2, plan     # one READ_PARQUET per file, none duplicated59    assert "HASH_JOIN" not in plan, plan60    from core.duck import con61    df = con().execute(sql, [*files, "2024-12-01", "2024-12-31", 5001]).df()62    assert len(df) == 21 and df["datetime"].is_unique       # dedup on datetime across the two buckets636465def test_run_topn_results(paths):66    from core.duck import con, run_topn67    df = run_topn(con(), paths, "datetime >= ? AND datetime <= ?", ["2025-06-30 09:30:00", "2025-06-30 09:34:59"],68                  "ticker", "datetime DESC", 2, order_by="ticker, datetime DESC", limit=100)69    assert len(df) == 6 and list(df["ticker"].unique()) == ["AAPL", "MSFT", "SMCP"]70    assert str(df["datetime"].iloc[0]).startswith("2025-06-30 09:34")717273def test_single_database_shared_across_threads(app):74    from core import duck75    seen = {}7677    def worker(i):78        c = duck.con()79        assert duck.con() is c                                   # stable within the thread80        seen[i] = (c, c.execute("SELECT current_setting('memory_limit')").fetchone()[0])8182    ts = [threading.Thread(target=worker, args=(i,)) for i in range(4)]83    for t in ts:84        t.start()85    for t in ts:86        t.join()87    assert len({id(v[0]) for v in seen.values()}) == 4      # one cursor per thread…88    assert len({v[1] for v in seen.values()}) == 1          # …one configured database89    limit = seen[0][1]90    assert limit.upper().startswith(("3.", "4", "3.7"))  # 4GB → DuckDB reports ~3.7 GiB91    threads = duck.con().execute("SELECT current_setting('threads')").fetchone()[0]92    assert int(threads) == duck.settings.duck_threads939495def test_cache_has_ttl_and_maxsize():96    from core.cache import TTLCache97    c = TTLCache(maxsize=3, ttl=1000)98    for k in "abcd":99        c.set(k, k)100    assert len(c) == 3 and c.get("a") is None and c.get("d") == "d"101    calls = {"n": 0}102103    def build():104        calls["n"] += 1105        return "v"106    assert c.get_or_build("x", build) == "v" and c.get_or_build("x", build) == "v" and calls["n"] == 1107    assert c.get_or_build("x", build, ttl=-1) == "v" and calls["n"] == 2      # caller-specific horizon expired108    assert c.invalidate("x") == 1 and c.get("x") is None109