203 lines
6.8 KiB
Python
203 lines
6.8 KiB
Python
"""解析入口：识别格式 → 读取 → 拆表 → 推断类型 → 每张表写一个 Parquet → 列概况
|
||
|
||
只读写本地文件，不碰数据库；上传流程在子进程里调用（worker.py），超时直接终止子进程
|
||
"""
|
||
|
||
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.profile import profile_table
|
||
from app.datafiles.readers import (
|
||
Grid,
|
||
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 的实测：各环节真实耗时和进度大致成正比，预计时间才准
|
||
STAGE_ENDS: dict[str, tuple[float, float]] = {
|
||
"csv": (0.01, 0.08),
|
||
"tsv": (0.01, 0.08),
|
||
"xlsx": (0.01, 0.30),
|
||
"xls": (0.01, 0.30),
|
||
"xlsb": (0.01, 0.30),
|
||
"json": (0.10, 0.35),
|
||
"jsonl": (0.03, 0.47),
|
||
"parquet": (0.02, 0.25),
|
||
}
|
||
|
||
# 建表各步骤完成时占其所在范围的比例：拆表（只有网格要）占网格的这么多，
|
||
# 然后每张表推断完列类型、写完 Parquet，剩下的是算列概况
|
||
SPLIT_AT = 0.35
|
||
TYPED_AT = 0.8
|
||
WRITTEN_AT = 0.85
|
||
|
||
|
||
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 在识别、读取、建表过程中被调用，用来向调用方报告进度
|
||
"""
|
||
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)
|
||
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))
|
||
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,
|
||
list(zip(*table.rows)),
|
||
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":
|
||
grid = read_delimited(detected.text, detected.delimiter, reading)
|
||
add_grid(grid, building)
|
||
case "xlsx" | "xls" | "xlsb":
|
||
grids, sheet_warnings = read_excel(path, reading)
|
||
warnings.extend(sheet_warnings)
|
||
# 各 sheet 按行数分摊建表的进度
|
||
total_rows = sum(len(grid.rows) for grid in grids)
|
||
done_rows = 0
|
||
for grid in grids:
|
||
start = done_rows / total_rows
|
||
done_rows += len(grid.rows)
|
||
add_grid(grid, 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(detected.text, 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]
|