380 lines
14 KiB
Python
380 lines
14 KiB
Python
"""把各种格式读成原始网格（CSV、Excel）或按列的记录（JSON、JSONL），Parquet 直接读成 Arrow 表
|
|
|
|
CSV 和 Excel 可能有标题行、多行表头、合计行，所以先读成不带表头的网格，再由 tables.py 识别表区域
|
|
文本格式边读边转成列（grid.py），不把整个文件解码进内存
|
|
"""
|
|
|
|
import codecs
|
|
import csv
|
|
import io
|
|
import json
|
|
from dataclasses import dataclass, field
|
|
from itertools import islice
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.csv as pacsv
|
|
import pyarrow.parquet as pq
|
|
from python_calamine import (
|
|
CalamineWorkbook,
|
|
PasswordError,
|
|
SheetTypeEnum,
|
|
SheetVisibleEnum,
|
|
)
|
|
|
|
from app.datafiles.base import (
|
|
MAX_TABLE_ROWS,
|
|
DataFileError,
|
|
Reporter,
|
|
brief,
|
|
no_report,
|
|
rows_exceeded,
|
|
scaled,
|
|
)
|
|
from app.datafiles.cells import as_text
|
|
from app.datafiles.grid import CHUNK_ROWS, Column, Grid, GridBuilder
|
|
from app.schemas.data_file import ParseWarning
|
|
|
|
# 原始网格比数据多出的标题、表头、空行，读取时先粗略拦住明显超限的文件，精确的行数在拆表后检查
|
|
ROW_SLACK = 1000
|
|
|
|
# JSON 嵌套对象展开的层数，更深的部分存成 JSON 字符串
|
|
JSON_FLATTEN_LEVELS = 2
|
|
|
|
# JSONL 坏行超过这个比例时整个文件判为失败
|
|
MAX_BAD_LINE_RATIO = 0.01
|
|
|
|
# 逐行读取时每隔这么多行报告一次进度
|
|
PROGRESS_ROWS = 20_000
|
|
|
|
# Arrow 读 CSV 每次处理的字节数；单行超过它时 Arrow 报错，改用标准库
|
|
ARROW_BLOCK_BYTES = 4 * 1024 * 1024
|
|
|
|
# 单元格最长的字符数；csv 模块默认 128KB，长文本单元格会报错
|
|
csv.field_size_limit(16 * 1024 * 1024)
|
|
|
|
|
|
@dataclass
|
|
class Records:
|
|
"""已有列结构的数据，按列存放"""
|
|
|
|
names: list[str]
|
|
columns: list[Column]
|
|
warnings: list[ParseWarning] = field(default_factory=list)
|
|
|
|
|
|
def read_delimited(
|
|
path: Path, codec: str, delimiter: str, report: Reporter = no_report
|
|
) -> Grid:
|
|
"""codec 是 detect 确认过的编码，整个文件都能按它解码
|
|
|
|
先用 Arrow 读（C++ 实现，快一个数量级）；它要求每行列数相同，遇到标题行、说明行这类长短不齐的行，
|
|
或引号不成对、单行过长，就报错。这时改用标准库 csv 从头读一遍。Arrow 读得出来时两者结果相同，测试里对比
|
|
"""
|
|
report = _increasing(report)
|
|
try:
|
|
return _read_delimited_arrow(path, codec, delimiter, report)
|
|
except pa.ArrowInvalid:
|
|
return _read_delimited_python(path, codec, delimiter, report)
|
|
|
|
|
|
def _read_delimited_arrow(
|
|
path: Path, codec: str, delimiter: str, report: Reporter
|
|
) -> Grid:
|
|
with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
|
|
first = next(csv.reader(text, delimiter=delimiter), [])
|
|
if not first:
|
|
raise pa.ArrowInvalid("第一行是空行，列数要从后面的行推断")
|
|
names = [f"c{i}" for i in range(len(first))]
|
|
builder = GridBuilder(strings=True)
|
|
size = max(path.stat().st_size, 1)
|
|
with path.open("rb") as raw:
|
|
if codec == "utf-8-sig" and raw.read(3) != codecs.BOM_UTF8:
|
|
raw.seek(0)
|
|
reader = pacsv.open_csv(
|
|
raw,
|
|
read_options=pacsv.ReadOptions(
|
|
column_names=names,
|
|
block_size=ARROW_BLOCK_BYTES,
|
|
encoding="utf8" if codec.startswith("utf-8") else codec,
|
|
),
|
|
parse_options=pacsv.ParseOptions(
|
|
delimiter=delimiter, newlines_in_values=True, ignore_empty_lines=False
|
|
),
|
|
convert_options=pacsv.ConvertOptions(
|
|
column_types={name: pa.string() for name in names},
|
|
strings_can_be_null=False,
|
|
quoted_strings_can_be_null=False,
|
|
),
|
|
)
|
|
batches: list[pa.RecordBatch] = []
|
|
waiting = 0
|
|
for batch in reader:
|
|
batches.append(batch)
|
|
waiting += len(batch)
|
|
# 攒到一个块再转：每段太短时，类型推断的固定开销会变大
|
|
if waiting >= CHUNK_ROWS:
|
|
_add_batches(builder, batches)
|
|
batches, waiting = [], 0
|
|
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
|
|
raise rows_exceeded(None)
|
|
report(raw.tell() / size)
|
|
if batches:
|
|
_add_batches(builder, batches)
|
|
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
|
|
raise rows_exceeded(None)
|
|
return builder.finish(None)
|
|
|
|
|
|
def _add_batches(builder: GridBuilder, batches: list[pa.RecordBatch]) -> None:
|
|
builder.add_strings(
|
|
[
|
|
pa.concat_arrays([batch.column(i) for batch in batches])
|
|
for i in range(batches[0].num_columns)
|
|
]
|
|
)
|
|
|
|
|
|
def _read_delimited_python(
|
|
path: Path, codec: str, delimiter: str, report: Reporter
|
|
) -> Grid:
|
|
builder = GridBuilder(strings=True)
|
|
size = max(path.stat().st_size, 1)
|
|
with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
|
|
reader = csv.reader(text, delimiter=delimiter)
|
|
while rows := list(islice(reader, PROGRESS_ROWS)):
|
|
builder.add_rows(rows)
|
|
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
|
|
raise rows_exceeded(None)
|
|
report(raw.tell() / size)
|
|
return builder.finish(None)
|
|
|
|
|
|
def _increasing(report: Reporter) -> Reporter:
|
|
"""Arrow 读到一半失败、换标准库从头读时，进度不往回走"""
|
|
reached = 0.0
|
|
|
|
def inner(fraction: float) -> None:
|
|
nonlocal reached
|
|
if fraction > reached:
|
|
reached = fraction
|
|
report(fraction)
|
|
|
|
return inner
|
|
|
|
|
|
def read_excel(
|
|
path: Path, report: Reporter = no_report
|
|
) -> tuple[list[Grid], list[ParseWarning]]:
|
|
try:
|
|
book = CalamineWorkbook.from_path(str(path))
|
|
except PasswordError:
|
|
raise DataFileError(422, "Excel 文件设置了密码，请去掉密码后重新上传") from None
|
|
except Exception:
|
|
raise DataFileError(422, "无法读取 Excel 文件，文件可能已损坏") from None
|
|
grids, warnings = [], []
|
|
try:
|
|
sheets = book.sheets_metadata
|
|
for number, meta in enumerate(sheets):
|
|
report(number / len(sheets))
|
|
# 图表页、宏表等没有单元格数据
|
|
if meta.typ != SheetTypeEnum.WorkSheet:
|
|
continue
|
|
if meta.visible != SheetVisibleEnum.Visible:
|
|
warnings.append(
|
|
ParseWarning(
|
|
code="hidden_sheet_skipped",
|
|
message=f"跳过了隐藏的 sheet「{meta.name}」",
|
|
)
|
|
)
|
|
continue
|
|
sheet = book.get_sheet_by_name(meta.name)
|
|
if sheet.total_height > MAX_TABLE_ROWS + ROW_SLACK:
|
|
raise rows_exceeded(meta.name)
|
|
# 解析 sheet 的 XML 占读取的大部分时间，逐行转成列占剩下的
|
|
rows = scaled(
|
|
report, (number + 0.8) / len(sheets), (number + 1) / len(sheets)
|
|
)
|
|
# iter_rows 保留顶部的空行，行号和 Excel 里看到的一致；左侧的空列不影响识别表区域；
|
|
# 没有单元格的 sheet 上 iter_rows 会崩溃，不调用
|
|
builder = GridBuilder()
|
|
cells = sheet.iter_rows() if sheet.height else iter(())
|
|
while batch := list(islice(cells, PROGRESS_ROWS)):
|
|
builder.add_rows(batch)
|
|
rows(builder.height / sheet.height)
|
|
grid = builder.finish(meta.name)
|
|
if not any(column.present().any() for column in grid.columns):
|
|
warnings.append(
|
|
ParseWarning(
|
|
code="empty_sheet_skipped",
|
|
message=f"sheet「{meta.name}」是空的，已跳过",
|
|
)
|
|
)
|
|
continue
|
|
grids.append(grid)
|
|
finally:
|
|
book.close()
|
|
return grids, warnings
|
|
|
|
|
|
def read_json(value: Any, report: Reporter = no_report) -> Grid | Records:
|
|
"""对象数组按记录读取；数组的数组当作网格，第一行通常是表头"""
|
|
warnings = []
|
|
if isinstance(value, dict):
|
|
records, path = _find_records(value)
|
|
if path:
|
|
warnings.append(
|
|
ParseWarning(
|
|
code="json_path_selected",
|
|
message=f"使用了 JSON 中「{path}」下的 {len(records)} 条记录",
|
|
)
|
|
)
|
|
else:
|
|
records = value
|
|
if not records:
|
|
raise DataFileError(422, "JSON 里没有找到数据")
|
|
if all(isinstance(item, list) for item in records):
|
|
if len(records) > MAX_TABLE_ROWS + ROW_SLACK:
|
|
raise rows_exceeded(None)
|
|
return Grid.from_rows(
|
|
None, ([_scalar(cell) for cell in row] for row in records)
|
|
)
|
|
if all(isinstance(item, dict) for item in records):
|
|
flattener = _Flattener()
|
|
for number, record in enumerate(records):
|
|
if number % PROGRESS_ROWS == 0:
|
|
report(number / len(records))
|
|
flattener.add(record)
|
|
result = flattener.finish()
|
|
result.warnings = warnings
|
|
return result
|
|
if not any(isinstance(item, (dict, list)) for item in records):
|
|
if len(records) > MAX_TABLE_ROWS:
|
|
raise rows_exceeded(None)
|
|
return Records(["值"], [Column.from_values(records)], warnings)
|
|
raise DataFileError(422, "JSON 数组里的元素类型不一致，需要都是对象或都是数组")
|
|
|
|
|
|
def read_jsonl(path: Path, codec: str, report: Reporter = no_report) -> Records:
|
|
"""逐行读取；codec 是 detect 确认过的编码"""
|
|
flattener, bad, total = _Flattener(), [], 0
|
|
size = max(path.stat().st_size, 1)
|
|
with path.open("rb") as raw, io.TextIOWrapper(raw, codec) as text:
|
|
for number, line in enumerate(text, start=1):
|
|
if number % PROGRESS_ROWS == 0:
|
|
report(raw.tell() / size)
|
|
if not line.strip():
|
|
continue
|
|
total += 1
|
|
try:
|
|
record = json.loads(line)
|
|
except ValueError:
|
|
record = None
|
|
if isinstance(record, dict):
|
|
flattener.add(record)
|
|
else:
|
|
bad.append(number)
|
|
if not flattener.count:
|
|
raise DataFileError(422, "JSONL 里没有找到数据")
|
|
if len(bad) > total * MAX_BAD_LINE_RATIO:
|
|
raise DataFileError(
|
|
422, f"JSONL 有 {len(bad)} 行无法解析（例如第 {bad[0]} 行），超过 1%"
|
|
)
|
|
result = flattener.finish()
|
|
if bad:
|
|
result.warnings.append(
|
|
ParseWarning(
|
|
code="bad_lines_skipped",
|
|
message=f"跳过了 {len(bad)} 行无法解析的内容：{brief([f'第 {n} 行' for n in bad])}",
|
|
)
|
|
)
|
|
return result
|
|
|
|
|
|
def read_parquet(path: Path) -> pa.Table:
|
|
try:
|
|
table = pq.read_table(path)
|
|
except Exception:
|
|
raise DataFileError(422, "无法读取 Parquet 文件，文件可能已损坏") from None
|
|
if table.num_rows > MAX_TABLE_ROWS:
|
|
raise rows_exceeded(None)
|
|
# struct 展开一层，列名写成 a.b；更深的嵌套由 columns.py 存成 JSON 字符串
|
|
return table.flatten()
|
|
|
|
|
|
def _find_records(value: dict) -> tuple[list, str | None]:
|
|
"""在对象的前两层里找最长的对象数组，例如 {"data": [...]}；找不到时按列式或单条记录处理"""
|
|
candidates = []
|
|
for key, item in value.items():
|
|
if _is_record_list(item):
|
|
candidates.append((len(item), str(key), item))
|
|
elif isinstance(item, dict):
|
|
for sub_key, sub in item.items():
|
|
if _is_record_list(sub):
|
|
candidates.append((len(sub), f"{key}.{sub_key}", sub))
|
|
if candidates:
|
|
_, path, records = max(candidates, key=lambda candidate: candidate[0])
|
|
return records, path
|
|
lists = list(value.values())
|
|
# 列式：{"地区": [...], "销售额": [...]}
|
|
if (
|
|
lists
|
|
and all(isinstance(item, list) for item in lists)
|
|
and len({len(item) for item in lists}) == 1
|
|
):
|
|
return [dict(zip(value, row)) for row in zip(*lists)], None
|
|
return [value], None
|
|
|
|
|
|
def _is_record_list(value: Any) -> bool:
|
|
return (
|
|
isinstance(value, list)
|
|
and bool(value)
|
|
and all(isinstance(item, dict) for item in value)
|
|
)
|
|
|
|
|
|
class _Flattener:
|
|
"""逐条展开记录，按列存放：列按首次出现的顺序排列，记录里没有的字段为空"""
|
|
|
|
def __init__(self) -> None:
|
|
self._index: dict[str, int] = {}
|
|
self._builder = GridBuilder()
|
|
|
|
@property
|
|
def count(self) -> int:
|
|
return self._builder.height
|
|
|
|
def add(self, record: dict) -> None:
|
|
if self._builder.height >= MAX_TABLE_ROWS:
|
|
raise rows_exceeded(None)
|
|
flat: dict[str, Any] = {}
|
|
_flatten(record, "", 0, flat)
|
|
row: list[Any] = [None] * len(self._index)
|
|
for name, value in flat.items():
|
|
position = self._index.setdefault(name, len(self._index))
|
|
if position < len(row):
|
|
row[position] = value
|
|
else:
|
|
row.append(value)
|
|
self._builder.add(row)
|
|
|
|
def finish(self) -> Records:
|
|
grid = self._builder.finish(None)
|
|
return Records(list(self._index), grid.columns)
|
|
|
|
|
|
def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None:
|
|
for key, value in record.items():
|
|
name = f"{prefix}.{key}" if prefix else str(key)
|
|
if isinstance(value, dict) and value and level < JSON_FLATTEN_LEVELS:
|
|
_flatten(value, name, level + 1, out)
|
|
else:
|
|
out[name] = _scalar(value)
|
|
|
|
|
|
def _scalar(value: Any) -> Any:
|
|
# 嵌套的对象、数组存成 JSON 字符串，其余原样交给类型推断
|
|
return as_text(value) if isinstance(value, (dict, list)) else value
|