From 4d53f66cd11c5ea56dce7dcaa40b07db33519d36 Mon Sep 17 00:00:00 2001 From: Zhangergou Date: Mon, 19 Oct 2026 15:40:56 +0800 Subject: [PATCH] =?UTF-8?q?=E5=87=8F=E5=B0=91=E6=97=A5K=E5=90=8C=E6=AD=A5?= =?UTF-8?q?=E4=B8=B4=E6=97=B6=E8=AF=BB=E5=86=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/services/kline_sync.py | 19 +++++- .../scripts/benchmark_daily_sync_staging.py | 64 ++++++++++++++++++ .../tests/test_fuyao_daily_streaming_sync.py | 66 ++++++++++++++++++- 3 files changed, 144 insertions(+), 5 deletions(-) create mode 100644 backend/scripts/benchmark_daily_sync_staging.py diff --git a/backend/app/services/kline_sync.py b/backend/app/services/kline_sync.py index aa78ca72..354c83c5 100644 --- a/backend/app/services/kline_sync.py +++ b/backend/app/services/kline_sync.py @@ -434,8 +434,23 @@ def _persist_daily_chunks(chunks, repo: KlineRepository) -> int: for date_dir in sorted(root.glob("date=*")): files = sorted(date_dir.glob("*.parquet")) - if files: - repo.append_daily(pl.scan_parquet(files).collect(engine="streaming")) + if not files: + continue + target_dir = repo.store.data_dir / "kline_daily" / date_dir.name + target_dir.mkdir(parents=True, exist_ok=True) + existing = sorted(target_dir.glob("*.parquet")) + if not existing: + for index, staged in enumerate(files): + replace_with_retry(staged, target_dir / f"part-stream-{index}.parquet") + continue + base = pl.scan_parquet(existing).collect(engine="streaming") + incoming = pl.scan_parquet(files).collect(engine="streaming") + merged = pl.concat([base, incoming], how="diagonal_relaxed").unique( + subset=["symbol", "date"], keep="last" + ) + repo.append_daily(merged) + for part in target_dir.glob("part-stream-*.parquet"): + part.unlink(missing_ok=True) finally: shutil.rmtree(root, ignore_errors=True) with contextlib.suppress(OSError): diff --git a/backend/scripts/benchmark_daily_sync_staging.py b/backend/scripts/benchmark_daily_sync_staging.py new file mode 100644 index 00000000..20c80a6d --- /dev/null +++ b/backend/scripts/benchmark_daily_sync_staging.py @@ -0,0 +1,64 @@ +"""Compare daily sync staging read-back with direct promotion. + +Run: uv run --frozen python -m scripts.benchmark_daily_sync_staging +Uses synthetic data only; no external data source is contacted. +""" +from __future__ import annotations + +import shutil +import tempfile +import time +from pathlib import Path + +import polars as pl + +ROWS = 100_000 +PARTS = 4 + + +def make_parts(root: Path) -> list[Path]: + symbols = [f"{i:06d}.SZ" for i in range(ROWS)] + frame = pl.DataFrame({ + "symbol": symbols, + "date": ["2026-09-18"] * ROWS, + "open": [10.0] * ROWS, + "high": [11.0] * ROWS, + "low": [9.0] * ROWS, + "close": [10.5] * ROWS, + "volume": [100.0] * ROWS, + "amount": [1050.0] * ROWS, + }).with_columns(pl.Series("bucket", [i % PARTS for i in range(ROWS)])) + paths = [] + for index, part in enumerate(frame.partition_by("bucket")): + part = part.drop("bucket") + path = root / f"part-{index}.parquet" + part.write_parquet(path) + paths.append(path) + return paths + + +def main() -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + parts = make_parts(root) + parts = parts[:PARTS] + target = root / "target.parquet" + + started = time.perf_counter() + pl.scan_parquet(parts).collect(engine="streaming").write_parquet(target) + read_back = time.perf_counter() - started + + target.unlink() + started = time.perf_counter() + for index, part in enumerate(parts): + shutil.copyfile(part, root / f"promoted-{index}.parquet") + promoted = time.perf_counter() - started + + print(f"parts={len(parts)}, rows={ROWS}") + print(f"read-back-and-rewrite: {read_back:.4f}s") + print(f"direct-promotion-copy: {promoted:.4f}s") + print(f"speedup: {read_back / promoted:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/backend/tests/test_fuyao_daily_streaming_sync.py b/backend/tests/test_fuyao_daily_streaming_sync.py index 4e67d277..28575ea5 100644 --- a/backend/tests/test_fuyao_daily_streaming_sync.py +++ b/backend/tests/test_fuyao_daily_streaming_sync.py @@ -50,6 +50,30 @@ def _route_fuyao(monkeypatch, provider): monkeypatch.setattr(custom, "get_provider", lambda name: provider) +def test_sync_promotes_staged_files_without_reading_them_back(monkeypatch, repo): + provider = _StreamingProvider([ + _daily("000001.SZ", date(2020, 1, 2)), + _daily("000002.SZ", date(2020, 1, 2)), + ]) + _route_fuyao(monkeypatch, provider) + monkeypatch.setattr( + kline_sync.pl, + "scan_parquet", + lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("不应读回临时文件")), + ) + + written = kline_sync.sync_and_persist_daily_batch( + ["000001.SZ", "000002.SZ"], repo, object(), + start_date=datetime(2020, 1, 1), end_date=datetime(2020, 1, 2), + ) + + assert written == 2 + parts = sorted((repo.store.data_dir / "kline_daily" / "date=2020-01-02").glob("*.parquet")) + assert len(parts) == 2 + stored = pl.concat([pl.read_parquet(part) for part in parts]) + assert set(stored["symbol"].to_list()) == {"000001.SZ", "000002.SZ"} + + def test_sync_persists_streamed_fuyao_chunks_only_after_fetch(monkeypatch, repo): provider = _StreamingProvider([ _daily("000001.SZ", date(2020, 1, 2)), @@ -64,13 +88,49 @@ def test_sync_persists_streamed_fuyao_chunks_only_after_fetch(monkeypatch, repo) assert written == 2 assert provider.get_daily_called is False - stored = pl.read_parquet( - repo.store.data_dir / "kline_daily" / "date=2020-01-02" / "part.parquet" - ) + parts = sorted((repo.store.data_dir / "kline_daily" / "date=2020-01-02").glob("*.parquet")) + stored = pl.concat([pl.read_parquet(part) for part in parts]) assert set(stored["symbol"].to_list()) == {"000001.SZ", "000002.SZ"} assert not (repo.store.data_dir / ".daily_sync_staging").exists() +def test_sync_merges_after_direct_promotion_without_duplicate_rows(monkeypatch, repo): + provider = _StreamingProvider([_daily("000001.SZ", date(2020, 1, 2))]) + _route_fuyao(monkeypatch, provider) + kline_sync.sync_and_persist_daily_batch( + ["000001.SZ"], repo, object(), + start_date=datetime(2020, 1, 1), end_date=datetime(2020, 1, 2), + ) + kline_sync.sync_and_persist_daily_batch( + ["000001.SZ"], repo, object(), + start_date=datetime(2020, 1, 1), end_date=datetime(2020, 1, 2), + ) + + files = sorted((repo.store.data_dir / "kline_daily" / "date=2020-01-02").glob("*.parquet")) + stored = pl.concat([pl.read_parquet(file) for file in files]) + assert stored.unique(subset=["symbol", "date"]).height == 1 + assert stored.height == 1 + + +def test_sync_merges_into_existing_daily_partition(monkeypatch, repo): + target = repo.store.data_dir / "kline_daily" / "date=2020-01-02" + target.mkdir(parents=True) + pl.DataFrame({ + "symbol": ["000001.SZ"], "date": [date(2020, 1, 2)], "close": [9.0], + }).write_parquet(target / "part.parquet") + provider = _StreamingProvider([_daily("000002.SZ", date(2020, 1, 2))]) + _route_fuyao(monkeypatch, provider) + + written = kline_sync.sync_and_persist_daily_batch( + ["000002.SZ"], repo, object(), + start_date=datetime(2020, 1, 1), end_date=datetime(2020, 1, 2), + ) + + assert written == 1 + stored = pl.read_parquet(target / "part.parquet") + assert set(stored["symbol"].to_list()) == {"000001.SZ", "000002.SZ"} + + def test_sync_discards_staging_when_streaming_fails(monkeypatch, repo): provider = _StreamingProvider( [_daily("000001.SZ", date(2020, 1, 2))], error=RuntimeError("network lost")