137 lines
4.3 KiB
Python
137 lines
4.3 KiB
Python
"""解析入口：识别格式 → 读取 → 拆表 → 推断类型 → 每张表写一个 Parquet → 列概况
|
||
|
||
只读写本地文件，不碰数据库；上传流程在子进程里调用（worker.py），超时直接终止子进程
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import pyarrow.parquet as pq
|
||
from pydantic import BaseModel
|
||
|
||
from app.datafiles.base import DataFileError
|
||
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
|
||
|
||
|
||
# 子进程把解析结果转成 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) -> ParsedFile:
|
||
"""解析 path，每张表在 out_dir 下写一个 Parquet：1.parquet、2.parquet……"""
|
||
detected = detect(path)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
warnings: list[ParseWarning] = []
|
||
tables: list[ParsedTable] = []
|
||
|
||
def add_table(
|
||
sheet: str | None,
|
||
part: int | None,
|
||
typed: TypedColumns,
|
||
table_warnings: list[ParseWarning],
|
||
) -> None:
|
||
target = out_dir / f"{len(tables) + 1}.parquet"
|
||
pq.write_table(typed.table, target)
|
||
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,
|
||
)
|
||
)
|
||
|
||
def add_grid(grid: Grid) -> None:
|
||
found, dropped = split_tables(grid)
|
||
warnings.extend(dropped)
|
||
for table in found:
|
||
names, renamed = clean_names(table.header, table.sheet)
|
||
typed, kept = build_columns(names, list(zip(*table.rows)), table.sheet)
|
||
add_table(
|
||
table.sheet,
|
||
table.part,
|
||
typed,
|
||
[*table.warnings, *_present(renamed, kept)],
|
||
)
|
||
|
||
def add_records(records: Records) -> None:
|
||
warnings.extend(records.warnings)
|
||
names, renamed = clean_names(records.names, None)
|
||
typed, kept = build_columns(names, records.columns, None)
|
||
add_table(None, None, typed, _present(renamed, kept))
|
||
|
||
match detected.format:
|
||
case "csv" | "tsv":
|
||
add_grid(read_delimited(detected.text, detected.delimiter))
|
||
case "xlsx" | "xls" | "xlsb":
|
||
grids, sheet_warnings = read_excel(path)
|
||
warnings.extend(sheet_warnings)
|
||
for grid in grids:
|
||
add_grid(grid)
|
||
case "json":
|
||
result = read_json(detected.value)
|
||
if isinstance(result, Grid):
|
||
add_grid(result)
|
||
else:
|
||
add_records(result)
|
||
case "jsonl":
|
||
add_records(read_jsonl(detected.text))
|
||
case "parquet":
|
||
typed, renamed = arrow_columns(read_parquet(path))
|
||
if typed.table.num_rows:
|
||
add_table(None, None, typed, _present(renamed))
|
||
|
||
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]
|