feat：优化增量执行、工具调用容错与评测体系
This commit is contained in:
+182
-63
@@ -1,16 +1,20 @@
|
||||
"""把各种格式读成原始网格（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,
|
||||
@@ -28,7 +32,8 @@ from app.datafiles.base import (
|
||||
rows_exceeded,
|
||||
scaled,
|
||||
)
|
||||
from app.datafiles.columns import as_text
|
||||
from app.datafiles.cells import as_text
|
||||
from app.datafiles.grid import CHUNK_ROWS, Column, Grid, GridBuilder
|
||||
from app.schemas.data_file import ParseWarning
|
||||
|
||||
# 原始网格比数据多出的标题、表头、空行，读取时先粗略拦住明显超限的文件，精确的行数在拆表后检查
|
||||
@@ -43,37 +48,120 @@ 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 Grid:
|
||||
"""不带表头的原始网格；rows[i] 是文件或 sheet 中的第 i + 1 行"""
|
||||
|
||||
sheet: str | None
|
||||
rows: list[list[Any]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Records:
|
||||
"""已有列结构的数据，按列存放"""
|
||||
|
||||
names: list[str]
|
||||
columns: list[list[Any]]
|
||||
columns: list[Column]
|
||||
warnings: list[ParseWarning] = field(default_factory=list)
|
||||
|
||||
|
||||
def read_delimited(text: str, delimiter: str, report: Reporter = no_report) -> Grid:
|
||||
rows = []
|
||||
stream = io.StringIO(text, newline="")
|
||||
for row in csv.reader(stream, delimiter=delimiter):
|
||||
rows.append(row)
|
||||
if len(rows) > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(None)
|
||||
if len(rows) % PROGRESS_ROWS == 0:
|
||||
report(stream.tell() / len(text))
|
||||
return Grid(None, rows)
|
||||
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(
|
||||
@@ -104,9 +192,19 @@ def read_excel(
|
||||
sheet = book.get_sheet_by_name(meta.name)
|
||||
if sheet.total_height > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(meta.name)
|
||||
# 保留左上角的空白区域，行号和 Excel 里看到的一致
|
||||
rows = sheet.to_python(skip_empty_area=False)
|
||||
if not any(cell != "" for row in rows for cell in row):
|
||||
# 解析 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",
|
||||
@@ -114,7 +212,7 @@ def read_excel(
|
||||
)
|
||||
)
|
||||
continue
|
||||
grids.append(Grid(meta.name, rows))
|
||||
grids.append(grid)
|
||||
finally:
|
||||
book.close()
|
||||
return grids, warnings
|
||||
@@ -137,42 +235,53 @@ def read_json(value: Any, report: Reporter = no_report) -> Grid | Records:
|
||||
if not records:
|
||||
raise DataFileError(422, "JSON 里没有找到数据")
|
||||
if all(isinstance(item, list) for item in records):
|
||||
return Grid(None, [[_scalar(cell) for cell in row] for row 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):
|
||||
result = _flatten_records(records, report)
|
||||
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):
|
||||
return Records(["值"], [list(records)], warnings)
|
||||
if len(records) > MAX_TABLE_ROWS:
|
||||
raise rows_exceeded(None)
|
||||
return Records(["值"], [Column.from_values(records)], warnings)
|
||||
raise DataFileError(422, "JSON 数组里的元素类型不一致，需要都是对象或都是数组")
|
||||
|
||||
|
||||
def read_jsonl(text: str, report: Reporter = no_report) -> Records:
|
||||
records, bad, total = [], [], 0
|
||||
lines = text.splitlines()
|
||||
# 解析每行 JSON 和展开成列各占一半
|
||||
parsing = scaled(report, 0.0, 0.5)
|
||||
for number, line in enumerate(lines, start=1):
|
||||
if number % PROGRESS_ROWS == 0:
|
||||
parsing(number / len(lines))
|
||||
if not line.strip():
|
||||
continue
|
||||
total += 1
|
||||
try:
|
||||
record = json.loads(line)
|
||||
except ValueError:
|
||||
record = None
|
||||
if isinstance(record, dict):
|
||||
records.append(record)
|
||||
else:
|
||||
bad.append(number)
|
||||
if not records:
|
||||
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 = _flatten_records(records, scaled(report, 0.5, 1.0))
|
||||
result = flattener.finish()
|
||||
if bad:
|
||||
result.warnings.append(
|
||||
ParseWarning(
|
||||
@@ -226,24 +335,34 @@ def _is_record_list(value: Any) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def _flatten_records(records: list[dict], report: Reporter = no_report) -> Records:
|
||||
# 列按首次出现的顺序排列，记录里没有的字段为空
|
||||
index: dict[str, int] = {}
|
||||
columns: list[list[Any]] = []
|
||||
for row_number, record in enumerate(records):
|
||||
if row_number % PROGRESS_ROWS == 0:
|
||||
report(row_number / len(records))
|
||||
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():
|
||||
if name not in index:
|
||||
index[name] = len(columns)
|
||||
columns.append([None] * row_number)
|
||||
columns[index[name]].append(value)
|
||||
for column in columns:
|
||||
if len(column) == row_number:
|
||||
column.append(None)
|
||||
return Records(list(index), columns)
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user