Merge PR #493: 减少日K同步临时读写

backend/app/services/kline_sync.py

@@ -438,8 +438,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):

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()

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")