"""解析入口:识别格式 → 读取 → 拆表 → 推断类型 → 每张表写一个 Parquet → 列概况 只读写本地文件,不碰数据库;上传流程在子进程里调用(worker.py),超时直接终止子进程 """ import gc from collections.abc import Callable from pathlib import Path from typing import Any, Literal import pyarrow.parquet as pq from pydantic import BaseModel from app.datafiles.base import DataFileError, Reporter, scaled from app.datafiles.columns import ( TypedColumns, arrow_columns, build_columns, clean_names, ) from app.datafiles.detect import detect from app.datafiles.grid import Grid from app.datafiles.profile import profile_table from app.datafiles.readers import ( Records, read_delimited, read_excel, read_json, read_jsonl, read_parquet, ) from app.datafiles.tables import split_tables from app.schemas.data_file import DataColumn, DataFormat, ParseWarning, TableProfile ParseStage = Literal["detecting", "reading", "building"] # 解析进度的回调:当前环节和整个解析已完成的比例(0 到 1,只增不减) Progress = Callable[[ParseStage, float], None] # 各格式识别、读取两个环节结束时的整体进度,剩下的是建表(拆表、推断类型、写 Parquet、算列概况) # 数字来自 scripts/measure_parse_progress.py 在 100 万行上的实测:各环节真实耗时和进度大致成正比,预计时间才准 STAGE_ENDS: dict[str, tuple[float, float]] = { "csv": (0.06, 0.28), "tsv": (0.06, 0.28), "xlsx": (0.01, 0.84), "xls": (0.01, 0.84), "xlsb": (0.01, 0.84), "json": (0.27, 0.84), "jsonl": (0.03, 0.90), "parquet": (0.01, 0.75), } # 建表各步骤完成时占其所在范围的比例:拆表(只有网格要)占网格的这么多, # 然后每张表推断完列类型、写完 Parquet,剩下的是算列概况 SPLIT_AT = 0.2 TYPED_AT = 0.4 WRITTEN_AT = 0.8 # Parquet 每个行组的行数。DuckDB 按行组并行扫描,pyarrow 默认一个行组放 100 多万行,整张表只能单线程读; # 行组越小并行度越高,但字典编码按行组重复,文件变大、写得也慢,这个大小两头折中 ROW_GROUP_ROWS = 128 * 1024 def no_progress(stage: ParseStage, value: float) -> None: pass # 子进程把解析结果转成 JSON 交回上传流程,所以用 pydantic 模型 class ParsedTable(BaseModel): sheet: str | None # 同一 sheet 切出多张表时的序号,从 1 开始 part: int | None rows: int columns: list[DataColumn] # out_dir 下的 Parquet 文件名 parquet: str profile: TableProfile preview: list[list[Any]] warnings: list[ParseWarning] class ParsedFile(BaseModel): format: DataFormat encoding: str | None tables: list[ParsedTable] # 不属于某张表的警告,例如跳过的 sheet warnings: list[ParseWarning] def parse_file( path: Path, out_dir: Path, progress: Progress = no_progress ) -> ParsedFile: """解析 path,每张表在 out_dir 下写一个 Parquet:1.parquet、2.parquet…… progress 在识别、读取、建表过程中被调用,用来向调用方报告进度 """ # 读取时成批创建短命的行列表,分代垃圾回收反复扫描它们却回收不到什么,关掉后读取快约三成; # 解析不产生循环引用,结束后照常回收 enabled = gc.isenabled() gc.disable() try: return _parse_file(path, out_dir, progress) finally: if enabled: gc.enable() def _parse_file(path: Path, out_dir: Path, progress: Progress) -> ParsedFile: progress("detecting", 0.0) detected = detect(path) detect_end, read_end = STAGE_ENDS[detected.format] out_dir.mkdir(parents=True, exist_ok=True) warnings: list[ParseWarning] = [] tables: list[ParsedTable] = [] def stage(name: ParseStage, start: float, end: float) -> Reporter: return scaled(lambda value: progress(name, value), start, end) reading = stage("reading", detect_end, read_end) building = stage("building", read_end, 1.0) progress("reading", detect_end) def add_table( sheet: str | None, part: int | None, typed: TypedColumns, table_warnings: list[ParseWarning], report: Reporter, ) -> None: target = out_dir / f"{len(tables) + 1}.parquet" pq.write_table(typed.table, target, row_group_size=ROW_GROUP_ROWS) report(WRITTEN_AT) profile, preview = profile_table(target, typed.columns, typed.notes) tables.append( ParsedTable( sheet=sheet, part=part, rows=typed.table.num_rows, columns=typed.columns, parquet=target.name, profile=profile, preview=preview, warnings=table_warnings, ) ) report(1.0) def add_grid(grid: Grid, report: Reporter) -> None: found, dropped = split_tables(grid, scaled(report, 0.0, SPLIT_AT)) # 拆出的表只引用用到的行;build_columns 推断完一列就释放一列,网格不能再留着 del grid warnings.extend(dropped) tables = scaled(report, SPLIT_AT, 1.0) for number, table in enumerate(found): # 这个 sheet 里每张表各占一份 share = scaled(tables, number / len(found), (number + 1) / len(found)) names, renamed = clean_names(table.header, table.sheet) typed, kept = build_columns( names, table.columns, table.sheet, scaled(share, 0.0, TYPED_AT), ) add_table( table.sheet, table.part, typed, [*table.warnings, *_present(renamed, kept)], scaled(share, TYPED_AT, 1.0), ) report(1.0) def add_records(records: Records, report: Reporter) -> None: warnings.extend(records.warnings) names, renamed = clean_names(records.names, None) typed, kept = build_columns( names, records.columns, None, scaled(report, 0.0, TYPED_AT) ) add_table( None, None, typed, _present(renamed, kept), scaled(report, TYPED_AT, 1.0) ) match detected.format: case "csv" | "tsv": # 不留局部变量,网格才能在建表时用完即释放 add_grid( read_delimited(path, detected.codec, detected.delimiter, reading), building, ) case "xlsx" | "xls" | "xlsb": grids, sheet_warnings = read_excel(path, reading) warnings.extend(sheet_warnings) # 各 sheet 按行数分摊建表的进度 total_rows = sum(grid.height for grid in grids) done_rows = 0 while grids: start = done_rows / total_rows done_rows += grids[0].height # 取出后不再引用,这个 sheet 的网格才能在建表时用完即释放 add_grid(grids.pop(0), scaled(building, start, done_rows / total_rows)) case "json": result = read_json(detected.value, reading) if isinstance(result, Grid): add_grid(result, building) else: add_records(result, building) case "jsonl": add_records(read_jsonl(path, detected.codec, reading), building) case "parquet": typed, renamed = arrow_columns(read_parquet(path)) if typed.table.num_rows: add_table(None, None, typed, _present(renamed), building) if not tables: raise DataFileError(422, "文件里没有找到数据") return ParsedFile( format=detected.format, encoding=detected.encoding, tables=tables, warnings=warnings, ) def _present(*warnings: ParseWarning | None) -> list[ParseWarning]: return [warning for warning in warnings if warning is not None]