perf(ext-data): CSV 编码转换改为分块进行, 峰值内存不再随文件线性增长

backend/app/services/ext_data.py

@@ -1,6 +1,7 @@
 """扩展数据服务 — 配置管理 + 文件解析 + Parquet 存储。"""
 from __future__ import annotations
 
+import codecs
 import copy
 import json
 import logging
@@ -453,6 +454,44 @@ def apply_config_mapping(df: pl.DataFrame, config: ExtConfig, data_dir: Path) ->
     return df
 
 
+# 编码识别与转换的分块大小，与 ext_data 上传写入用的块大小一致。
+_TRANSCODE_CHUNK_BYTES = 1024 * 1024
+
+
+def _decodes_as(file_path: Path, encoding: str) -> bool:
+    """整个文件能否按 encoding 完整解码，逐块判断，不把文件读进内存。"""
+    decoder = codecs.getincrementaldecoder(encoding)()
+    try:
+        with file_path.open("rb") as src:
+            while chunk := src.read(_TRANSCODE_CHUNK_BYTES):
+                decoder.decode(chunk)
+            decoder.decode(b"", True)  # 结尾处的半个字符也算解码失败
+    except UnicodeDecodeError:
+        return False
+    return True
+
+
+def _transcode_to_utf8(file_path: Path, out_path: Path, encoding: str) -> bool:
+    """按 encoding 逐块转成 UTF-8 写入 out_path；解码失败则删除半成品返回 False。
+
+    增量解码器负责跨块边界的多字节字符：GBK 一个汉字两字节，正好落在块边界
+    上时前半截会被留到下一块，不会被误判成解码失败。
+    """
+    decoder = codecs.getincrementaldecoder(encoding)()
+    try:
+        with (
+            file_path.open("rb") as src,
+            out_path.open("w", encoding="utf-8", newline="") as dst,
+        ):
+            while chunk := src.read(_TRANSCODE_CHUNK_BYTES):
+                dst.write(decoder.decode(chunk))
+            dst.write(decoder.decode(b"", True))
+    except UnicodeDecodeError:
+        out_path.unlink(missing_ok=True)
+        return False
+    return True
+
+
 def ensure_utf8_csv(file_path: Path) -> Path:
     """确保 CSV 文件以 UTF-8 编码可读，非 UTF-8（如 GBK/GB18030）则转换。
 
@@ -463,27 +502,14 @@ def ensure_utf8_csv(file_path: Path) -> Path:
     返回值：若已是 UTF-8 则返回原路径；否则在同目录写一个 *.utf8 文件并返回它
     （调用方用临时目录，随目录一起清理）。
     """
-    raw = file_path.read_bytes()
     # BOM 处理：UTF-8-SIG 等带 BOM 文件直接交给 Polars（它认识 BOM）
-    try:
-        raw.decode("utf-8")
+    if _decodes_as(file_path, "utf-8"):
         return file_path  # 已是合法 UTF-8
-    except UnicodeDecodeError:
-        pass
     # 依次尝试常见中文编码，第一个能完整解码的即为命中
     for enc in ("gb18030", "gbk", "gb2312", "big5"):
-        try:
-            text = raw.decode(enc)
-        except UnicodeDecodeError:
-            continue
         out_path = file_path.with_suffix(file_path.suffix + ".utf8")
-        # newline="" 关闭写入时的换行转换。默认转换在 Windows 上把文本里的 \n
-        # 写成 \r\n，源文件本来就是 CRLF 时就变成 \r\r\n，多出来的 \r 被 Polars
-        # 当作最后一列内容的一部分：列名变成 "收盘价\r"，每行的值变成 "12.34\r"，
-        # 该列于是被推断为字符串而不是数值。本函数针对的同花顺/东财/通达信和
-        # Windows Excel 导出文件用的正是 CRLF。
-        with out_path.open("w", encoding="utf-8", newline="") as f:
-            f.write(text)
+        if not _transcode_to_utf8(file_path, out_path, enc):
+            continue
         logger.info("CSV 编码转换 %s → %s (%s)", file_path.name, out_path.name, enc)
         return out_path
     # 都无法解码：返回原路径，让 Polars 抛出更精确的原始错误

backend/tests/test_ext_csv_transcode_memory.py

@@ -0,0 +1,72 @@
+"""CSV 编码转换的内存占用回归测试。
+
+上传路径专门用 `_write_upload_capped` 分块落盘（issue #204），docstring 写明
+是为了「避免 `await file.read()` 把整个文件读入内存(大文件可能触发高内存占用、
+进程 OOM 或服务不可用)」。但紧接着的 `ensure_utf8_csv` 曾用 `read_bytes()`
+把同一个文件整个读回内存再整体解码，把这层保护抵消掉：50MB 上限的文件实测
+峰值 302MB（6.05×），解码出的 str 比原字节还大。
+
+这里用 tracemalloc 量峰值，判据是「与文件大小无关」而不是某个绝对值：转换
+按块进行时峰值只跟块大小有关，文件翻倍不会让峰值翻倍。
+"""
+from __future__ import annotations
+
+import tracemalloc
+from pathlib import Path
+
+from app.services.ext_data import ensure_utf8_csv
+
+_ROW = "浦发银行,600000,12.34,上海证券交易所\r\n"
+_HEADER = "名称,代码,收盘价,交易所\r\n"
+
+
+def _gbk_csv(path: Path, size_bytes: int) -> Path:
+    body = _ROW * (size_bytes // len(_ROW.encode("gb18030")))
+    path.write_bytes((_HEADER + body).encode("gb18030"))
+    return path
+
+
+def _peak_bytes(path: Path) -> int:
+    tracemalloc.start()
+    try:
+        ensure_utf8_csv(path)
+        _, peak = tracemalloc.get_traced_memory()
+    finally:
+        tracemalloc.stop()
+    return peak
+
+
+def test_transcode_peak_memory_is_far_below_the_file(tmp_path: Path) -> None:
+    path = _gbk_csv(tmp_path / "big.csv", 8 * 1024 * 1024)
+
+    peak = _peak_bytes(path)
+
+    # 分块转换时峰值只跟块大小有关。修复前是文件的 6 倍。
+    assert peak < path.stat().st_size
+
+
+def test_transcode_peak_memory_does_not_grow_with_the_file(tmp_path: Path) -> None:
+    small = _peak_bytes(_gbk_csv(tmp_path / "small.csv", 2 * 1024 * 1024))
+    large = _peak_bytes(_gbk_csv(tmp_path / "large.csv", 8 * 1024 * 1024))
+
+    # 文件大 4 倍，峰值不应跟着涨：留一倍余量给解释器噪声。
+    assert large < small * 2
+
+
+def test_multibyte_character_on_a_chunk_boundary_survives(tmp_path: Path) -> None:
+    # GBK 一个汉字两字节，分块时可能正好被切开。增量解码器负责把半个字符
+    # 留到下一块；若改成逐块独立 decode，这里会解码失败并整份回退。
+    import polars as pl
+
+    from app.services.ext_data import _TRANSCODE_CHUNK_BYTES
+
+    filler = "浦发银行" * ((_TRANSCODE_CHUNK_BYTES // 8) + 1)
+    text = f"名称,备注\r\n浦发银行,{filler}\r\n"
+    path = tmp_path / "boundary.csv"
+    path.write_bytes(text.encode("gb18030"))
+    assert path.stat().st_size > _TRANSCODE_CHUNK_BYTES
+
+    df = pl.read_csv(ensure_utf8_csv(path), infer_schema_length=10000)
+
+    assert df.columns == ["名称", "备注"]
+    assert df["备注"][0] == filler