diff --git a/backend/app/services/ext_data.py b/backend/app/services/ext_data.py index 2cf4c9c..2f081b6 100644 --- a/backend/app/services/ext_data.py +++ b/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 抛出更精确的原始错误 diff --git a/backend/tests/test_ext_csv_transcode_memory.py b/backend/tests/test_ext_csv_transcode_memory.py new file mode 100644 index 0000000..d9308c9 --- /dev/null +++ b/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