feat:优化增量执行、工具调用容错与评测体系

This commit is contained in:
2026-10-11 04:40:25 -07:00
parent cd5a2b4d8f
commit dadd797391
22 changed files with 2198 additions and 578 deletions
+516
View File
@@ -0,0 +1,516 @@
"""单元格的类型规则:单个值用 classify 判断,一段 Arrow 数组用 classify_piece 批量判断
classify 是规则的唯一定义:一个值是什么类型、转换成什么值。大表不能逐个调用它,classify_piece 对常见写法
(空值记号、布尔、ASCII 数字和百分数、编号、ISO 日期和常见的日期时间)用正则和 Arrow 转换批量得出结论,
其余的值去重后逐个交给 classify。批量规则只认它有把握的写法,认不出的一律交给 classify,所以结论和逐个
调用 classify 相同;测试里逐值对比两者
"""
import json
import math
import re
from dataclasses import dataclass
from datetime import UTC, date, datetime, time
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
NULL_TOKENS = {"", "-", "--", "—", "——", "N/A", "n/a", "NULL", "null", "#N/A"}
EXCEL_ERROR = re.compile(r"^#(?:NULL!|DIV/0!|VALUE!|REF!|NAME\?|NUM!|N/A)$")
CURRENCY_SYMBOLS = "¥¥$€"
NUMBER = re.compile(r"^[+-]?(?:\d+|\d{1,3}(?:,\d{3})+)(?:\.\d+)?(?:[eE][+-]?\d+)?$")
INTEGER = re.compile(r"^[+-]?\d+$")
# 16 位及以上的纯数字是证件号、卡号一类的编号,转成数值会丢失精度
MIN_ID_DIGITS = 16
# 按年份后的第一个字符选格式;%m、%d 也接受一位数。没有对应格式的直接判为不是日期:
# strptime 每次失败都要抛异常,一个值试完所有格式约 16 μs,还要抢一把模块级的锁,多线程时更慢
DATE_FORMATS = {
"-": ("%Y-%m-%d", "%Y-%m"),
"/": ("%Y/%m/%d", "%Y/%m"),
".": ("%Y.%m.%d",),
"年": ("%Y年%m月%d日", "%Y年%m月"),
}
DATETIME_FORMATS = {
"-": ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M"),
"/": ("%Y/%m/%d %H:%M:%S", "%Y/%m/%d %H:%M"),
}
INT64_RANGE = range(-(2**63), 2**63)
# 超过这个范围的整数存成 float 会丢精度
EXACT_FLOAT_INT = 2**53
# str.isspace() 为真的全部字符,也就是 str.strip() 去掉的字符;Arrow 自带的去空白函数按另一套定义,
# 读取时显式传入这些字符,结果才和 Python 一致
WHITESPACE = "\t\n\x0b\x0c\r\x1c\x1d\x1e\x1f \x85\xa0            

   "
# 正则里匹配上面这些字符的写法(RE2 的 \s 只认 ASCII 空白)
WHITESPACE_CLASS = "".join(f"\\x{{{ord(char):x}}}" for char in WHITESPACE)
# 类型代码,classify 返回的类型名按这个顺序编号
KINDS = ("null", "bool", "int", "float", "percent", "date", "datetime", "text")
NULL, BOOL, INT, FLOAT, PERCENT, DATE, DATETIME, TEXT = range(len(KINDS))
KIND_CODES = {name: code for code, name in enumerate(KINDS)}
MICROS_PER_DAY = 86_400_000_000
EPOCH = datetime(1970, 1, 1)
def classify(value: Any) -> tuple[str, Any]:
"""单个值的类型和转换后的值;类型是 KINDS 之一"""
if value is None:
return "null", None
if isinstance(value, bool):
return "bool", value
if isinstance(value, int):
return ("int", value) if value in INT64_RANGE else ("text", str(value))
if isinstance(value, float):
return (
("null", None)
if math.isnan(value) or math.isinf(value)
else ("float", value)
)
if isinstance(value, datetime):
# Excel 只设了日期格式的单元格读出来是零点的 datetime
if value.time() == time() and value.tzinfo is None:
return "date", value.date()
return "datetime", _naive(value)
if isinstance(value, date):
return "date", value
if not isinstance(value, str):
return "text", as_text(value)
text = value.strip()
if text in NULL_TOKENS or EXCEL_ERROR.match(text):
return "null", None
if text.lower() in ("true", "false"):
return "bool", text.lower() == "true"
number = parse_number(text)
if number is not None:
if isinstance(number, int):
# 带千分位的超长整数和超出范围的原生整数一样保留文本
return ("int", number) if number in INT64_RANGE else ("text", text)
return "float", number
if text.endswith("%"):
number = parse_number(text[:-1].rstrip())
if number is not None:
return "percent", number / 100
moment = parse_datetime(text)
if moment is not None:
return ("datetime" if isinstance(moment, datetime) else "date"), moment
return "text", text
def parse_number(text: str) -> int | float | None:
"""千分位和货币符号可以有;有前导零或 16 位以上的数字串是编号,不算数值
判断数字一律用 isdecimal:isdigit 对 ²、① 也返回 True,但 int() 解析不了它们
"""
if text and text[0] in CURRENCY_SYMBOLS:
text = text[1:].lstrip()
# 最常见的纯数字不用正则
plain = text.isdecimal()
if not plain and not NUMBER.match(text):
return None
digits = text.lstrip("+-")
if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal():
return None
if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS:
return None
if plain:
return int(text)
text = text.replace(",", "")
return int(text) if INTEGER.match(text) else float(text)
def parse_datetime(text: str) -> date | datetime | None:
if len(text) < 6 or not text[:4].isdecimal():
return None
# 标准 ISO 写法先用 fromisoformat,比 strptime 快一个数量级
if ":" in text:
try:
return _naive(datetime.fromisoformat(text))
except ValueError:
pass
for pattern in DATETIME_FORMATS.get(text[4], ()):
try:
return datetime.strptime(text, pattern)
except ValueError:
continue
return None
if len(text) == 10 and text[4] == "-":
try:
return date.fromisoformat(text)
except ValueError:
pass
for pattern in DATE_FORMATS.get(text[4], ()):
try:
return datetime.strptime(text, pattern).date()
except ValueError:
continue
return None
def as_text(value: Any) -> str:
"""单元格转成文本:列名、保留为文本的值、警告里引用的内容都用它;Excel 里年份等数字读出来是 2025.0"""
if value is None:
return ""
if isinstance(value, str):
return value.strip()
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, float) and value.is_integer():
return str(int(value))
if isinstance(value, datetime):
return value.isoformat(sep=" ")
if isinstance(value, (date, time)):
return value.isoformat()
if isinstance(value, (dict, list)):
# Parquet 的嵌套值里可能有日期、Decimal、bytes,JSON 表示不了的写成字符串
return json.dumps(value, ensure_ascii=False, default=str)
return str(value)
def _naive(value: datetime) -> datetime:
# 带时区的时间统一换成 UTC 再去掉时区,同一列里的值才能比较
if value.tzinfo is None:
return value
return value.astimezone(UTC).replace(tzinfo=None)
# ---- 批量判断 ----
# classify 只会把这些开头的字符串判成文本以外的类型(空值记号和布尔另外判断):
# 货币符号,或可选的正负号加数字(数值、百分数、日期都以数字开头)
CANDIDATE = r"^(?:[¥¥$€]|[+-]?\p{Nd})"
# 最多 15 位整数部分的 ASCII 数值,保证整数在 int64 内、百分数除以 100 与 Python 结果相同;
# 有前导零的(编号)、货币符号后带空格的等写法不在这里,交给 classify
_DIGITS = r"[¥¥$€]?[+-]?(?:0|[1-9][0-9]{0,14}|[1-9][0-9]{0,2}(?:,[0-9]{3}){1,4})(?:\.[0-9]+)?(?:[eE][+-]?[0-9]{1,3})?"
FAST_NUMBER = f"^{_DIGITS}$"
FAST_PERCENT = f"^{_DIGITS}%$"
# 有前导零或 16 位以上的纯数字是编号:classify 判成文本
FAST_ID = r"^0[0-9]+$|^[0-9]{16,}$"
FAST_DATE = r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}$"
# fromisoformat 能解析的写法,以及 DATETIME_FORMATS 里分隔符为 - 或 /、各部分一两位数字的写法
FAST_DATETIMES = (
r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}[T ][0-9]{2}:[0-9]{2}(?::[0-9]{2}(?:\.[0-9]{3}(?:[0-9]{3})?)?)?$",
r"^[1-9][0-9]{3}-[0-9]{1,2}-[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
r"^[1-9][0-9]{3}/[0-9]{1,2}/[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
)
NULL_SET = pa.array(sorted(NULL_TOKENS))
BOOL_SET = pa.array(["true", "false"])
NULL_TEXT = pa.scalar(None, pa.string())
@dataclass
class Classified:
"""一段单元格逐个分类的结果,和对每个值调用 classify 的结论相同
kinds 是每个值的类型代码;ints、reals、moments、flags 和这段等长,只在对应类型的位置上有意义,
这段里没有这种类型时为空
"""
piece: pa.Array
kinds: np.ndarray
# int
ints: np.ndarray | None = None
# float、percent(已除以 100)
reals: np.ndarray | None = None
# date、datetime:自 1970 年起的微秒数,日期按零点
moments: np.ndarray | None = None
# bool
flags: np.ndarray | None = None
def values(self, name: str) -> np.ndarray:
"""取某类转换后的值,没有时补零"""
array = getattr(self, name)
if array is None:
array = np.zeros(len(self.kinds), _DTYPES[name])
setattr(self, name, array)
return array
def count(self) -> np.ndarray:
return np.bincount(self.kinds, minlength=len(KINDS))
_DTYPES = {"ints": np.int64, "reals": np.float64, "moments": np.int64, "flags": bool}
def union_layout(piece: pa.UnionArray) -> tuple[np.ndarray, np.ndarray]:
"""dense union 每个位置属于第几个子数组、是子数组里的第几个
切片后的 union,pyarrow 的 type_codes、offsets 不计切片的偏移,所以直接读缓冲区
"""
start, stop = piece.offset, piece.offset + len(piece)
_, codes, offsets = piece.buffers()[:3]
return (
np.frombuffer(codes, np.int8)[start:stop],
np.frombuffer(offsets, np.int32)[start:stop],
)
def mask(values: pa.Array | pa.ChunkedArray) -> np.ndarray:
"""Arrow 布尔数组转成 numpy,空值当作 False"""
return pc.fill_null(values, False).to_numpy(zero_copy_only=False)
def classify_piece(piece: pa.Array, order: list | None = None) -> Classified:
"""批量判断一段单元格;段的类型是读取时产生的几种之一,见 grid.py
order 是字符串段依次尝试的批量规则,同一列的各段传同一个列表:这些规则认的写法互不重叠,先试哪个
结论都一样,所以每段之后把认出最多的规则挪到前面,后面的段大多一次就认完
"""
kind = piece.type
if pa.types.is_union(kind):
return _classify_union(piece, list(FAST_RULES) if order is None else order)
valid = mask(piece.is_valid())
if pa.types.is_string(kind):
return _classify_strings(
piece, valid, list(FAST_RULES) if order is None else order
)
result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
if pa.types.is_null(kind):
return result
if pa.types.is_boolean(kind):
result.kinds[valid] = BOOL
result.flags = mask(piece)
elif pa.types.is_integer(kind):
result.kinds[valid] = INT
result.ints = pc.fill_null(piece, 0).to_numpy()
elif pa.types.is_floating(kind):
reals = pc.fill_null(piece, 0.0).to_numpy()
result.kinds[valid] = np.where(np.isfinite(reals[valid]), FLOAT, NULL)
result.reals = reals
elif pa.types.is_date32(kind):
result.kinds[valid] = DATE
result.moments = (
pc.fill_null(piece, 0).cast(pa.int32()).to_numpy().astype(np.int64)
* MICROS_PER_DAY
)
elif pa.types.is_timestamp(kind):
moments = pc.fill_null(piece.cast(pa.int64()), 0).to_numpy()
midnight = moments % MICROS_PER_DAY == 0
result.kinds[valid] = np.where(midnight[valid], DATE, DATETIME)
result.moments = moments
else:
raise TypeError(f"不支持的单元格类型:{kind}")
return result
def _classify_union(piece: pa.UnionArray, order: list) -> Classified:
"""类型混杂的段(见 grid.py):每个子数组各判断一次,再按位置拼回"""
result = Classified(piece, np.full(len(piece), NULL, np.int8))
codes, offsets = union_layout(piece)
for code in range(piece.type.num_fields):
at = np.flatnonzero(codes == code)
if not len(at):
continue
child = classify_piece(piece.field(code), order)
rows = offsets[at]
result.kinds[at] = child.kinds[rows]
for name in _DTYPES:
values = getattr(child, name)
if values is not None:
result.values(name)[at] = values[rows]
return result
def _classify_strings(piece: pa.Array, valid: np.ndarray, order: list) -> Classified:
"""字符串段:按 classify 的顺序(空值记号 → 布尔 → 数值 → 百分数 → 日期时间 → 文本)批量判断"""
result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
# 空值记号、Excel 错误值、布尔都不以货币符号或数字开头,只在其余的值里找;
# 这些值剩下的就是文本,类型已经记好。数值、日期列整段都是候选,省掉这几步
pending = valid & mask(pc.match_substring_regex(piece, CANDIDATE))
others = valid & ~pending
if others.any():
hit = others & mask(
pc.or_(
pc.is_in(piece, value_set=NULL_SET),
pc.match_substring_regex(piece, EXCEL_ERROR.pattern),
)
)
result.kinds[hit] = NULL
others &= ~hit
if others.any():
lower = pc.utf8_lower(piece)
hit = others & mask(pc.is_in(lower, value_set=BOOL_SET))
if hit.any():
result.kinds[hit] = BOOL
result.values("flags")[hit] = mask(pc.equal(lower, "true"))[hit]
claimed = {}
for rule in order:
left = int(pending.sum())
if not left:
break
rule(piece, result, pending)
claimed[rule] = left - int(pending.sum())
order.sort(key=lambda rule: -claimed.get(rule, 0))
if pending.any():
_delegate(piece, result, pending)
return result
def _matched(piece: pa.Array, pattern: str, pending: np.ndarray) -> np.ndarray:
"""还没有结论、并且符合 pattern 的值的位置"""
return np.flatnonzero(pending & mask(pc.match_substring_regex(piece, pattern)))
def _plain_number(strings: pa.Array) -> pa.Array:
"""去掉货币符号、千分位和正号,剩下 Arrow 能直接转换的数字
只用于符合 FAST_NUMBER 的值:货币符号只在开头,正号在开头或指数里(1e+5 去掉后是 1e5,值不变)
"""
return pc.replace_substring_regex(strings, r"[¥¥$€,+]", "")
def _fast_numbers(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
rows = _matched(piece, FAST_NUMBER, pending)
if not len(rows):
return
plain = _plain_number(piece.take(rows))
integer = mask(pc.match_substring_regex(plain, r"^-?[0-9]+$"))
if integer.any():
at = rows[integer]
result.kinds[at] = INT
result.values("ints")[at] = pc.cast(
plain.filter(integer), pa.int64()
).to_numpy()
if not integer.all():
at = rows[~integer]
result.kinds[at] = FLOAT
result.values("reals")[at] = pc.cast(
plain.filter(~integer), pa.float64()
).to_numpy()
pending[rows] = False
def _fast_percents(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
rows = _matched(piece, FAST_PERCENT, pending)
if not len(rows):
return
plain = _plain_number(pc.utf8_slice_codeunits(piece.take(rows), 0, -1))
result.kinds[rows] = PERCENT
result.values("reals")[rows] = pc.cast(plain, pa.float64()).to_numpy() / 100
pending[rows] = False
def _fast_ids(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
# 类型已经记为文本,只需标记为已判断
pending[_matched(piece, FAST_ID, pending)] = False
def _fast_dates(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
rows = _matched(piece, FAST_DATE, pending)
if len(rows):
_put_moments(result, pending, rows, piece.take(rows), pa.date32(), DATE)
def _fast_datetimes(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
for pattern in FAST_DATETIMES:
rows = _matched(piece, pattern, pending)
if not len(rows):
continue
# 统一成 Arrow 认识的 ISO 写法:/ 换成 -,一位数补零(相邻的两位数各占一次,所以替换两遍)
strings = pc.replace_substring(piece.take(rows), "/", "-")
for _ in range(2):
strings = pc.replace_substring_regex(
strings, r"([-: ])([0-9])([-: ]|$)", r"\10\2\3"
)
_put_moments(result, pending, rows, strings, pa.timestamp("us"), DATETIME)
def _put_moments(
result: Classified,
pending: np.ndarray,
rows: np.ndarray,
strings: pa.Array,
arrow_type: pa.DataType,
kind: int,
) -> None:
"""转换成日期或时间并填回;转换不了的(2 月 30 日、25 点)留给 classify"""
ok = np.zeros(len(rows), bool)
moments = np.zeros(len(rows), np.int64)
_cast_parts(strings, arrow_type, 0, ok, moments)
at = rows[ok]
result.kinds[at] = kind
result.values("moments")[at] = moments[ok]
pending[at] = False
# 二分到这么短还转换失败,就整段交给 classify
MIN_CAST_PART = 64
def _cast_parts(
strings: pa.Array,
arrow_type: pa.DataType,
offset: int,
ok: np.ndarray,
moments: np.ndarray,
) -> None:
"""Arrow 的转换遇到一个坏值就整体失败;失败时对半分开重试,只把含坏值的小段留下"""
try:
values = pc.cast(strings, arrow_type)
except pa.ArrowInvalid:
if len(strings) <= MIN_CAST_PART:
return
half = len(strings) // 2
_cast_parts(strings.slice(0, half), arrow_type, offset, ok, moments)
_cast_parts(strings.slice(half), arrow_type, offset + half, ok, moments)
return
end = offset + len(strings)
ok[offset:end] = True
if pa.types.is_date32(arrow_type):
moments[offset:end] = (
values.cast(pa.int32()).to_numpy().astype(np.int64) * MICROS_PER_DAY
)
else:
moments[offset:end] = values.cast(pa.int64()).to_numpy()
# 批量规则;认的写法互不重叠,顺序不影响结论
FAST_RULES = (_fast_numbers, _fast_percents, _fast_ids, _fast_dates, _fast_datetimes)
def _delegate(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
"""批量规则认不出的值去重后逐个交给 classify,再按位置填回"""
rows = np.flatnonzero(pending)
values = piece.take(rows)
unique = pc.unique(values)
codes = pc.index_in(values, value_set=unique).to_numpy()
outcomes = [classify(value) for value in unique.to_pylist()]
kinds = np.array([KIND_CODES[kind] for kind, _ in outcomes], np.int8)
result.kinds[rows] = kinds[codes]
for name, wanted, convert in _FIELDS:
chosen = np.isin(kinds, wanted)
if not chosen.any():
continue
converted = np.array(
[
convert(value) if pick else 0
for (_, value), pick in zip(outcomes, chosen)
],
_DTYPES[name],
)
at = chosen[codes]
result.values(name)[rows[at]] = converted[codes[at]]
pending[rows] = False
def _micros(value: date) -> int:
if isinstance(value, datetime):
return (value - EPOCH) // _MICROSECOND
return (value - EPOCH.date()).days * MICROS_PER_DAY
_MICROSECOND = datetime(1970, 1, 1, microsecond=1) - EPOCH
# classify 的结果填进 Classified 的哪个字段:字段名、对应的类型代码、怎样转换
_FIELDS = (
("ints", (INT,), int),
("reals", (FLOAT, PERCENT), float),
("moments", (DATE, DATETIME), _micros),
("flags", (BOOL,), bool),
)
+125 -211
View File
@@ -1,52 +1,46 @@
"""列名清理和类型推断,输出 Arrow 列
CSV 的单元格全是字符串;Excel 的数值、日期单元格直接采用,文本单元格再按规则推断
一列的非空值必须全部符合同一类型才转换,否则保留文本:宁可让模型在 SQL 里 TRY_CAST,也不悄悄丢掉值
单元格的类型规则在 cells.py;这里按列汇总:一列的非空值必须全部符合同一类型才转换,否则保留文本,
宁可让模型在 SQL 里 TRY_CAST,也不悄悄丢掉值。各列互不相关,用线程并行推断(Arrow 的计算不占 GIL)
"""
import json
import math
import os
import re
from collections import Counter
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import UTC, date, datetime, time
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
from app.datafiles.base import Reporter, brief, no_report, where
from app.datafiles.cells import (
BOOL,
DATE,
DATETIME,
EXACT_FLOAT_INT,
FAST_RULES,
FLOAT,
INT,
KINDS,
MICROS_PER_DAY,
NULL,
NULL_TEXT,
PERCENT,
TEXT,
Classified,
as_text,
classify_piece,
)
from app.datafiles.grid import Column
from app.schemas.data_file import DataColumn, ParseWarning
NULL_TOKENS = {"", "-", "--", "—", "——", "N/A", "n/a", "NULL", "null", "#N/A"}
EXCEL_ERROR = re.compile(r"^#(?:NULL!|DIV/0!|VALUE!|REF!|NAME\?|NUM!|N/A)$")
CURRENCY_SYMBOLS = "¥¥$€"
NUMBER = re.compile(r"^[+-]?(?:\d+|\d{1,3}(?:,\d{3})+)(?:\.\d+)?(?:[eE][+-]?\d+)?$")
INTEGER = re.compile(r"^[+-]?\d+$")
# 16 位及以上的纯数字是证件号、卡号一类的编号,转成数值会丢失精度
MIN_ID_DIGITS = 16
# 按年份后的第一个字符选格式;%m、%d 也接受一位数
DATE_FORMATS = {
"-": ("%Y-%m-%d", "%Y-%m"),
"/": ("%Y/%m/%d", "%Y/%m"),
".": ("%Y.%m.%d",),
"年": ("%Y年%m月%d日", "%Y年%m月"),
}
DATETIME_FORMATS = (
"%Y-%m-%d %H:%M:%S",
"%Y-%m-%d %H:%M",
"%Y/%m/%d %H:%M:%S",
"%Y/%m/%d %H:%M",
)
INT64_RANGE = range(-(2**63), 2**63)
# 超过这个范围的整数存成 float 会丢精度
EXACT_FLOAT_INT = 2**53
NUMERIC_KINDS = ("int", "float", "percent")
DATE_KINDS = ("date", "datetime")
NUMERIC_KINDS = (INT, FLOAT, PERCENT)
DATE_KINDS = (DATE, DATETIME)
# 同时推断的列数
INFER_THREADS = min(8, os.cpu_count() or 1)
@dataclass
@@ -84,19 +78,28 @@ def clean_names(
def build_columns(
names: list[str],
values: list[Sequence[Any]],
columns: list[Column],
sheet: str | None,
report: Reporter = no_report,
) -> tuple[TypedColumns, ParseWarning | None]:
"""按列推断类型;有列因少数无法解析的值保留为文本时,返回一条 kept_as_text 警告"""
columns, arrays, notes = [], [], {}
for number, (name, column_values) in enumerate(zip(names, values), start=1):
column, array, note = infer_column(name, column_values)
report(number / len(names))
columns.append(column)
arrays.append(array)
if note:
notes[name] = note
"""按列推断类型;有列因少数无法解析的值保留为文本时,返回一条 kept_as_text 警告
columns 会被清空:每列推断完就释放它的原始单元格,大表不必同时留着原始和转换后的两份
"""
described, arrays, notes = [], [], {}
with ThreadPoolExecutor(INFER_THREADS) as pool:
futures = [
pool.submit(infer_column, name, column)
for name, column in zip(names, columns)
]
columns.clear()
for number, (name, future) in enumerate(zip(names, futures), start=1):
column, array, note = future.result()
report(number / len(names))
described.append(column)
arrays.append(array)
if note:
notes[name] = note
warning = None
if notes:
listed = "、".join(f"「{name}」" for name in notes)
@@ -105,7 +108,7 @@ def build_columns(
message=f"{where(sheet)}{listed}有无法解析的值,整列按文本处理",
)
return TypedColumns(
columns, pa.Table.from_arrays(arrays, names=names), notes
described, pa.Table.from_arrays(arrays, names=names), notes
), warning
@@ -154,194 +157,105 @@ def _arrow_kind(arrow_type: pa.DataType) -> str | None:
def infer_column(
name: str, values: Sequence[Any]
) -> tuple[DataColumn, pa.Array, str | None]:
value_kinds, parsed = [], []
for value in values:
kind, result = classify(value)
value_kinds.append(kind)
parsed.append(result)
kinds = Counter(kind for kind in value_kinds if kind != "null")
present = set(kinds)
name: str, column: Column
) -> tuple[DataColumn, pa.ChunkedArray, str | None]:
order = list(FAST_RULES)
parts = [classify_piece(piece, order) for piece in column.pieces]
counts = sum((part.count() for part in parts), np.zeros(len(KINDS), np.int64))
present = {kind for kind in range(len(KINDS)) if counts[kind] and kind != NULL}
if present and present <= {"bool"}:
return DataColumn(name=name, type="boolean"), pa.array(parsed, pa.bool_()), None
if present and present <= {"int", "float"}:
if all(
value is None
or isinstance(value, int)
or (value.is_integer() and abs(value) < EXACT_FLOAT_INT)
for value in parsed
):
ints = [None if value is None else int(value) for value in parsed]
return (
DataColumn(name=name, type="number"),
pa.array(ints, pa.int64()),
None,
)
return (
DataColumn(name=name, type="number"),
pa.array(parsed, pa.float64()),
None,
def build(make, arrow_type: pa.DataType) -> pa.ChunkedArray:
return pa.chunked_array([make(part) for part in parts], arrow_type)
if present and present <= {BOOL}:
array = build(
lambda part: _pick(part, "flags", (BOOL,), pa.bool_()), pa.bool_()
)
if present == {"percent"}:
column = DataColumn(name=name, type="number", format="percent")
return column, pa.array(parsed, pa.float64()), None
if present and present <= {"date"}:
return DataColumn(name=name, type="date"), pa.array(parsed, pa.date32()), None
return DataColumn(name=name, type="boolean"), array, None
if present and present <= {INT, FLOAT}:
if all(_integral(part) for part in parts):
array = build(lambda part: _numbers(part, np.int64), pa.int64())
else:
array = build(lambda part: _numbers(part, np.float64), pa.float64())
return DataColumn(name=name, type="number"), array, None
if present == {PERCENT}:
array = build(
lambda part: _pick(part, "reals", (PERCENT,), pa.float64()), pa.float64()
)
return DataColumn(name=name, type="number", format="percent"), array, None
if present and present <= {DATE}:
array = build(_dates, pa.date32())
return DataColumn(name=name, type="date"), array, None
if present and present <= set(DATE_KINDS):
stamps = [
datetime.combine(value, time()) if type(value) is date else value
for value in parsed
]
return (
DataColumn(name=name, type="datetime"),
pa.array(stamps, pa.timestamp("us")),
None,
array = build(
lambda part: _pick(part, "moments", DATE_KINDS, pa.timestamp("us")),
pa.timestamp("us"),
)
return DataColumn(name=name, type="datetime"), array, None
texts = [
None if result is None else as_text(value)
for value, result in zip(values, parsed)
]
note = _text_note(texts, value_kinds, kinds)
return DataColumn(name=name, type="text"), pa.array(texts, pa.string()), note
array = build(_texts, pa.string())
return DataColumn(name=name, type="text"), array, _text_note(counts, parts)
def classify(value: Any) -> tuple[str, Any]:
"""单个值的类型和转换后的值;类型是 null、bool、int、float、percent、date、datetime、text 之一"""
if value is None:
return "null", None
if isinstance(value, bool):
return "bool", value
if isinstance(value, int):
return ("int", value) if value in INT64_RANGE else ("text", str(value))
if isinstance(value, float):
return (
("null", None)
if math.isnan(value) or math.isinf(value)
else ("float", value)
)
if isinstance(value, datetime):
# Excel 只设了日期格式的单元格读出来是零点的 datetime
if value.time() == time() and value.tzinfo is None:
return "date", value.date()
return "datetime", _naive(value)
if isinstance(value, date):
return "date", value
if not isinstance(value, str):
return "text", as_text(value)
text = value.strip()
if text in NULL_TOKENS or EXCEL_ERROR.match(text):
return "null", None
if text.lower() in ("true", "false"):
return "bool", text.lower() == "true"
number = parse_number(text)
if number is not None:
return ("int" if isinstance(number, int) else "float"), number
if text.endswith("%"):
number = parse_number(text[:-1].rstrip())
if number is not None:
return "percent", number / 100
moment = parse_datetime(text)
if moment is not None:
return ("datetime" if isinstance(moment, datetime) else "date"), moment
return "text", text
def _pick(
part: Classified, field: str, kinds: tuple[int, ...], arrow_type: pa.DataType
) -> pa.Array:
"""这段里属于 kinds 的值,其余位置为空"""
keep = np.isin(part.kinds, kinds)
return pa.array(part.values(field), arrow_type, mask=~keep)
def parse_number(text: str) -> int | float | None:
"""千分位和货币符号可以有;有前导零或 16 位以上的数字串是编号,不算数值
判断数字一律用 isdecimal:isdigit 对 ²、① 也返回 True,但 int() 解析不了它们
"""
if text and text[0] in CURRENCY_SYMBOLS:
text = text[1:].lstrip()
# 最常见的纯数字不用正则
plain = text.isdecimal()
if not plain and not NUMBER.match(text):
return None
digits = text.lstrip("+-")
if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal():
return None
if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS:
return None
if plain:
return int(text)
text = text.replace(",", "")
return int(text) if INTEGER.match(text) else float(text)
def _integral(part: Classified) -> bool:
"""小数都是整数值并且在 float 能精确表示的范围内(Excel 的整数读出来是 100.0)"""
reals = part.values("reals")[part.kinds == FLOAT]
return bool(np.all((np.floor(reals) == reals) & (np.abs(reals) < EXACT_FLOAT_INT)))
def parse_datetime(text: str) -> date | datetime | None:
if len(text) < 6 or not text[:4].isdecimal():
return None
# 标准 ISO 写法先用 fromisoformat,比 strptime 快一个数量级
if ":" in text:
try:
return _naive(datetime.fromisoformat(text))
except ValueError:
pass
for pattern in DATETIME_FORMATS:
try:
return datetime.strptime(text, pattern)
except ValueError:
continue
return None
if len(text) == 10 and text[4] == "-":
try:
return date.fromisoformat(text)
except ValueError:
pass
for pattern in DATE_FORMATS.get(text[4], ()):
try:
return datetime.strptime(text, pattern).date()
except ValueError:
continue
return None
def _numbers(part: Classified, dtype) -> pa.Array:
is_int, is_float = part.kinds == INT, part.kinds == FLOAT
# 其他位置可能是 NaN(float 段里的空值),先清零再转换
reals = np.where(is_float, part.values("reals"), 0.0)
values = np.where(is_int, part.values("ints").astype(dtype), reals.astype(dtype))
keep = is_int | is_float
return pa.array(values, pa.from_numpy_dtype(dtype), mask=~keep)
def as_text(value: Any) -> str:
"""单元格转成文本:列名、保留为文本的值、警告里引用的内容都用它;Excel 里年份等数字读出来是 2025.0"""
if value is None:
return ""
if isinstance(value, str):
return value.strip()
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, float) and value.is_integer():
return str(int(value))
if isinstance(value, datetime):
return value.isoformat(sep=" ")
if isinstance(value, (date, time)):
return value.isoformat()
if isinstance(value, (dict, list)):
# Parquet 的嵌套值里可能有日期、Decimal、bytes,JSON 表示不了的写成字符串
return json.dumps(value, ensure_ascii=False, default=str)
return str(value)
def _dates(part: Classified) -> pa.Array:
days = (part.values("moments") // MICROS_PER_DAY).astype(np.int32)
return pa.array(days, pa.date32(), mask=part.kinds != DATE)
def _naive(value: datetime) -> datetime:
# 带时区的时间统一换成 UTC 再去掉时区,同一列里的值才能比较
if value.tzinfo is None:
return value
return value.astimezone(UTC).replace(tzinfo=None)
def _texts(part: Classified) -> pa.Array:
"""保留为文本:字符串是去掉空白后的原文,其他类型按 as_text 转换;空值记号仍为空"""
keep = part.kinds != NULL
if pa.types.is_string(part.piece.type):
return pc.if_else(pa.array(keep), part.piece, NULL_TEXT)
return pa.array(
[
as_text(value) if kept else None
for value, kept in zip(part.piece.to_pylist(), keep)
],
pa.string(),
)
def _text_note(
texts: list[str | None], value_kinds: list[str], kinds: Counter[str]
) -> str | None:
def _text_note(counts: np.ndarray, parts: list[Classified]) -> str | None:
"""大部分值是数值或日期、少数无法解析时,说明原因,模型可以在 SQL 里用 TRY_CAST"""
failed = kinds["text"]
numeric = sum(kinds[kind] for kind in NUMERIC_KINDS)
dates = sum(kinds[kind] for kind in DATE_KINDS)
failed = int(counts[TEXT])
numeric = int(sum(counts[kind] for kind in NUMERIC_KINDS))
dates = int(sum(counts[kind] for kind in DATE_KINDS))
if not failed or max(numeric, dates) < failed:
return None
examples: list[str] = []
for text, kind in zip(texts, value_kinds):
if kind == "text" and text not in examples:
examples.append(text)
for part in parts:
failures = _texts(part).filter(pa.array(part.kinds == TEXT))
for text in pc.unique(failures).to_pylist():
if text not in examples:
examples.append(text)
if len(examples) == 2:
break
if len(examples) == 2:
break
shown = "、".join(f"「{example}」" for example in examples)
label = "数值" if numeric >= dates else "日期"
return f"看起来是{label},有 {failed} 个值无法解析,例如{shown}"
+122 -34
View File
@@ -1,4 +1,7 @@
"""按文件内容识别格式和文本编码,不看扩展名"""
"""按文件内容识别格式和文本编码,不看扩展名
文本文件整个解码一遍确认编码,但只留开头一段用来判断格式和分隔符;读取时再按确认的编码边读边解码
"""
import codecs
import csv
@@ -15,11 +18,16 @@ from app.schemas.data_file import DataFormat
# xlsx/xlsb 解压后的总大小上限,防止解压炸弹
MAX_UNZIPPED_BYTES = 1024 * 1024 * 1024
# 带 BOM 的文本直接按 BOM 解码;UTF-32 的 BOM 以 UTF-16 的开头,解出来含 NUL,按二进制拒绝
# 整体是一个 JSON 值的文件只能整个读进内存解析,峰值内存约为文件大小的 6~7 倍;
# 这个上限让它和最大的 CSV 一样峰值约 1 GB。JSONL 边读边解析,不受这个限制
MAX_JSON_BYTES = 150 * 1024 * 1024
# 带 BOM 的文本直接按 BOM 解码:编码名、跳过 BOM 后的解码方式、读取时的编码(会自己去掉 BOM);
# UTF-32 的 BOM 以 UTF-16 的开头,解出来含 NUL,按二进制拒绝
BOMS = (
(codecs.BOM_UTF8, "utf-8", "utf-8"),
(codecs.BOM_UTF16_LE, "utf-16-le", "utf-16"),
(codecs.BOM_UTF16_BE, "utf-16-be", "utf-16"),
(codecs.BOM_UTF8, "utf-8", "utf-8", "utf-8-sig"),
(codecs.BOM_UTF16_LE, "utf-16", "utf-16-le", "utf-16"),
(codecs.BOM_UTF16_BE, "utf-16", "utf-16-be", "utf-16"),
)
# 没有 BOM 时依次尝试;GB18030 兼容 GBK、GB2312,中文 Windows 上 Excel 另存的 CSV 多是这种
@@ -30,23 +38,42 @@ DELIMITERS = (",", "\t", ";", "|")
# 判断分隔符只看开头这么多行
SNIFF_LINES = 50
# 判断格式和分隔符用的开头文本长度;每次解码读入的字节数
HEAD_CHARS = 1024 * 1024
READ_BYTES = 1024 * 1024
# 开头像 JSON 才按 JSON 处理:对象的 { 后面只能是 " 或 },数组的 [ 后面是一个值的开头;
# CSV 的列名也可能以 [ 开头,例如 [日期],[销售额]
JSON_START = re.compile(r'\{\s*["}]|\[\s*(?:[\[{"\]\d-]|(?:true|false|null)\s*[,\]])')
UNKNOWN_ENCODING = "无法识别文件编码,请另存为 UTF-8 后重新上传"
@dataclass(frozen=True)
class Detected:
format: DataFormat
# 文本格式的编码名(返回给调用方)和读取时传给 open 的编码
encoding: str | None = None
# 文本格式解码后的内容
text: str | None = None
codec: str | None = None
# 文本格式开头的一段
head: str | None = None
# CSV/TSV 的分隔符
delimiter: str | None = None
# JSON 已解析的内容,避免读取时再解析一遍
value: Any = None
@dataclass(frozen=True)
class _Text:
encoding: str
codec: str
head: str
# 开头一段之后还有内容
truncated: bool
# 整个文件只有空白
blank: bool
def detect(path: Path) -> Detected:
with path.open("rb") as file:
head = file.read(8)
@@ -57,14 +84,19 @@ def detect(path: Path) -> Detected:
if head.startswith(b"PAR1"):
return Detected("parquet")
text, encoding = decode_text(path.read_bytes())
stripped = text.lstrip()
if not stripped:
text = scan_text(path)
if text.blank:
raise DataFileError(422, "文件是空的")
if JSON_START.match(stripped):
return _json_format(text, encoding)
delimiter = sniff_delimiter(text)
return Detected("tsv" if delimiter == "\t" else "csv", encoding, text, delimiter)
if JSON_START.match(text.head.lstrip()):
return _json_format(path, text)
delimiter = sniff_delimiter(text.head)
return Detected(
"tsv" if delimiter == "\t" else "csv",
text.encoding,
text.codec,
text.head,
delimiter,
)
def _zip_format(path: Path) -> DataFormat:
@@ -83,41 +115,76 @@ def _zip_format(path: Path) -> DataFormat:
raise DataFileError(415, UNSUPPORTED_FORMAT)
def decode_text(data: bytes) -> tuple[str, str]:
"""返回解码后的文本和编码名"""
for bom, codec, name in BOMS:
if data.startswith(bom):
def scan_text(path: Path) -> _Text:
"""确认文本编码:BOM 优先,否则依次严格尝试 FALLBACK_ENCODINGS"""
with path.open("rb") as file:
start = file.read(4)
for bom, name, decoding, codec in BOMS:
if start.startswith(bom):
try:
text = data[len(bom) :].decode(codec)
return _decode(path, name, decoding, codec, len(bom))
except UnicodeDecodeError:
raise DataFileError(
422, "无法识别文件编码,请另存为 UTF-8 后重新上传"
) from None
return _reject_binary(text), name
raise DataFileError(422, UNKNOWN_ENCODING) from None
for encoding in FALLBACK_ENCODINGS:
try:
return _reject_binary(data.decode(encoding)), encoding
return _decode(path, encoding, encoding, encoding, 0)
except UnicodeDecodeError:
continue
raise DataFileError(422, "无法识别文件编码,请另存为 UTF-8 后重新上传")
raise DataFileError(422, UNKNOWN_ENCODING)
def _reject_binary(text: str) -> str:
# 文本文件不会有 NUL;图片、PDF 等二进制文件偶尔能按 GB18030 解码成功,在这里拦下
if "\x00" in text:
raise DataFileError(415, UNSUPPORTED_FORMAT)
return text
def _decode(path: Path, name: str, decoding: str, codec: str, skip: int) -> _Text:
"""严格解码整个文件,留下开头一段;含 NUL 时按二进制文件拒绝"""
decoder = codecs.getincrementaldecoder(decoding)("strict")
parts: list[str] = []
kept, blank, truncated = 0, True, False
with path.open("rb") as file:
file.seek(skip)
while True:
chunk = file.read(READ_BYTES)
text = decoder.decode(chunk, final=not chunk)
# 文本文件不会有 NUL;图片、PDF 等二进制文件偶尔能按 GB18030 解码成功,在这里拦下
if "\x00" in text:
raise DataFileError(415, UNSUPPORTED_FORMAT)
blank = blank and not text.strip()
if kept < HEAD_CHARS:
parts.append(text)
kept += len(text)
elif text:
truncated = True
if not chunk:
break
head = "".join(parts)
if len(head) > HEAD_CHARS:
head, truncated = head[:HEAD_CHARS], True
cut = head.rfind("\n")
if truncated and cut >= 0:
# 截断处可能在一行中间,丢掉最后半行
head = head[: cut + 1]
return _Text(name, codec, head, truncated, blank)
def _json_format(text: str, encoding: str) -> Detected:
def _json_format(path: Path, text: _Text) -> Detected:
if _first_line_is_record(text):
return Detected("jsonl", text.encoding, text.codec, text.head)
# 整体是一个 JSON 值,只能整个读进来解析
if path.stat().st_size > MAX_JSON_BYTES:
raise DataFileError(
413,
f"JSON 文件不能超过 {MAX_JSON_BYTES // 1024 // 1024} MB,"
"请转成 JSONL(每行一个 JSON 对象)后上传",
)
content = path.read_bytes().decode(text.codec)
try:
return Detected("json", encoding, text, value=json.loads(text))
return Detected(
"json", text.encoding, text.codec, text.head, value=json.loads(content)
)
except ValueError as error:
json_error = error
first = next(line for line in text.splitlines() if line.strip())
first = next(line for line in content.splitlines() if line.strip())
try:
if isinstance(json.loads(first), dict):
return Detected("jsonl", encoding, text)
return Detected("jsonl", text.encoding, text.codec, text.head)
except ValueError:
pass
# 开头像 JSON 却不是合法 JSON,按 JSON 报错比当成 CSV 解析出一堆乱码更有用
@@ -128,6 +195,27 @@ def _json_format(text: str, encoding: str) -> Detected:
)
def _first_line_is_record(text: _Text) -> bool:
"""第一行是完整的 JSON 对象、后面还有内容:整体不可能是一个 JSON 值,是 JSONL
不用整个解析就能判断,大 JSONL 文件不必整个读进内存
"""
lines = text.head.splitlines(keepends=True)
index = next((i for i, line in enumerate(lines) if line.strip()), None)
if index is None:
return False
first, rest = lines[index], lines[index + 1 :]
if first == first.rstrip("\r\n") and (rest or text.truncated):
# 第一行在开头一段里没有结束
return False
if not (text.truncated or any(line.strip() for line in rest)):
return False
try:
return isinstance(json.loads(first), dict)
except ValueError:
return False
def sniff_delimiter(text: str) -> str:
"""选出让开头各行列数最一致、列数最多的分隔符;都只有一列时按逗号处理"""
lines = []
+306
View File
@@ -0,0 +1,306 @@
"""按列存放的单元格,读取、拆表和类型推断共用
每列是若干段 Arrow 数组,每段一种类型:字符串(已去掉首尾空白,空串记为空值)、布尔、int64、float64、
date32、timestamp[us],或全是空值。读取时每 CHUNK_ROWS 行转一次,同一时间只有一个块是 Python 对象;
块里同一列的类型混杂时(表头和数值在同一列、数值列里夹着"待定"),这一段存成 Arrow 的 dense union:
每种类型一个子数组,原值和原类型都保留。Arrow 的 is_valid 等函数不认 union 的子数组里的空值,
所以按段处理的地方都经过 each,对子数组分别计算再按位置拼回
"""
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from datetime import date, datetime, time
from itertools import zip_longest
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
from app.datafiles.cells import (
EXACT_FLOAT_INT,
INT64_RANGE,
NULL_TEXT,
WHITESPACE,
as_text,
mask,
union_layout,
)
CHUNK_ROWS = 65_536
NoneType = type(None)
class Column:
def __init__(self, pieces: Sequence[pa.Array]):
self.pieces = [piece for piece in pieces if len(piece)]
self.offsets = np.cumsum([0, *map(len, self.pieces)])
@classmethod
def from_values(cls, values: Sequence[Any]) -> "Column":
return cls([_convert(values)])
def __len__(self) -> int:
return int(self.offsets[-1])
def present(self) -> np.ndarray:
"""每个单元格是否有值"""
return _join(
[
each(piece, lambda part: mask(part.is_valid()), bool)
for piece in self.pieces
],
bool,
)
def matches(self, test) -> np.ndarray:
"""字符串单元格中 test(Arrow 字符串数组 → 布尔数组)为真的位置;其他类型的单元格都不算"""
def check(part: pa.Array) -> np.ndarray:
if pa.types.is_string(part.type):
return mask(test(part))
return np.zeros(len(part), bool)
return _join([each(piece, check, bool) for piece in self.pieces], bool)
def slice(self, start: int, stop: int) -> "Column":
pieces = []
for piece, offset in zip(self.pieces, self.offsets):
begin, end = max(start - offset, 0), min(stop - offset, len(piece))
if begin < end:
pieces.append(piece.slice(begin, end - begin))
return Column(pieces)
def take(self, rows: np.ndarray) -> "Column":
"""按行号(升序)取出若干行;行号连续时不复制数据"""
if len(rows) and rows[-1] - rows[0] + 1 == len(rows):
return self.slice(int(rows[0]), int(rows[-1]) + 1)
bounds = np.searchsorted(rows, self.offsets)
return Column(
[
piece.take(pa.array(rows[bounds[i] : bounds[i + 1]] - self.offsets[i]))
for i, piece in enumerate(self.pieces)
]
)
def to_pylist(self) -> list[Any]:
return [value for piece in self.pieces for value in piece.to_pylist()]
@dataclass
class Grid:
"""不带表头的原始网格;第 i 行是文件或 sheet 中的第 i + 1 行,各列一样长"""
sheet: str | None
columns: list[Column]
height: int
@classmethod
def from_rows(cls, sheet: str | None, rows: Iterable[Sequence[Any]]) -> "Grid":
builder = GridBuilder()
for row in rows:
builder.add(row)
return builder.finish(sheet)
def rows(
self, start: int, stop: int, columns: Sequence[int] | None = None
) -> list[list[Any]]:
"""[start, stop) 行的 Python 值,只取 columns 这些列"""
return self._rows(
[column.slice(start, stop) for column in self._pick(columns)], stop - start
)
def rows_at(
self, rows: np.ndarray, columns: Sequence[int] | None = None
) -> list[list[Any]]:
return self._rows(
[column.take(rows) for column in self._pick(columns)], len(rows)
)
def _pick(self, columns: Sequence[int] | None) -> list[Column]:
return self.columns if columns is None else [self.columns[c] for c in columns]
@staticmethod
def _rows(columns: list[Column], count: int) -> list[list[Any]]:
if not columns:
return [[] for _ in range(count)]
return [list(row) for row in zip(*(column.to_pylist() for column in columns))]
class GridBuilder:
"""逐行接收单元格,每 CHUNK_ROWS 行转成列;行长短不一时补空值
strings 为真表示单元格全是字符串(CSV),跳过类型检查;已经按列读好的字符串(Arrow 读 CSV)用 add_strings
"""
def __init__(self, strings: bool = False):
self._strings = strings
self._columns: list[list[pa.Array]] = []
self._pending: list[Sequence[Any]] = []
self._done = 0
@property
def height(self) -> int:
return self._done + len(self._pending)
def add(self, row: Sequence[Any]) -> None:
self._pending.append(row)
if len(self._pending) >= CHUNK_ROWS:
self._flush()
def add_rows(self, rows: Sequence[Sequence[Any]]) -> None:
self._pending.extend(rows)
if len(self._pending) >= CHUNK_ROWS:
self._flush()
def add_strings(self, columns: Sequence[pa.Array]) -> None:
"""一批按列存放的字符串单元格,各列等长"""
self._flush()
self._append([[_normalize(column)] for column in columns], len(columns[0]))
def finish(self, sheet: str | None) -> Grid:
self._flush()
return Grid(sheet, [Column(pieces) for pieces in self._columns], self._done)
def _flush(self) -> None:
rows, count = self._pending, len(self._pending)
if not count:
return
width = max(map(len, rows))
self._append(
[
[
_normalize(pa.array(values, pa.string()))
if self._strings
else _convert(values)
]
for values in (zip_longest(*rows) if width else ())
],
count,
)
self._pending = []
def _append(self, columns: list[list[pa.Array]], count: int) -> None:
for position, pieces in enumerate(columns):
if position == len(self._columns):
# 这一列在前面的块里没有出现过
self._columns.append([pa.nulls(self._done)])
self._columns[position].extend(pieces)
for pieces in self._columns[len(columns) :]:
pieces.append(pa.nulls(count))
self._done += count
def each(piece: pa.Array, compute, dtype) -> np.ndarray:
"""对一段计算 numpy 结果;union 段对每个子数组分别计算,再按位置拼回"""
if not pa.types.is_union(piece.type):
return compute(piece)
codes, offsets = union_layout(piece)
result = np.zeros(len(piece), dtype)
for code in range(piece.type.num_fields):
at = np.flatnonzero(codes == code)
if len(at):
result[at] = compute(piece.field(code))[offsets[at]]
return result
def _join(parts: list[np.ndarray], dtype) -> np.ndarray:
return np.concatenate(parts) if parts else np.zeros(0, dtype)
def _normalize(piece: pa.Array) -> pa.Array:
"""字符串去掉首尾空白(和 str.strip 相同),空串记为空值"""
if not pa.types.is_string(piece.type):
return piece
piece = pc.utf8_trim(piece, WHITESPACE)
return pc.if_else(pc.equal(piece, ""), NULL_TEXT, piece)
def _convert(values: Sequence[Any]) -> pa.Array:
"""一段单元格转成 Arrow:类型一致时转成一个数组,混杂时转成 dense union"""
types = set(map(type, values))
types.discard(NoneType)
if str in types and len(types) > 1:
# 空白单元格(Excel 读出来是 "")不算一种类型
values = [
None if type(v) is str and (not v or v.isspace()) else v for v in values
]
types = set(map(type, values))
types.discard(NoneType)
piece = _uniform(values, types)
if piece is not None:
return _normalize(piece)
families: dict[type, int] = {}
children: list[list[Any]] = []
codes = np.empty(len(values), np.int8)
offsets = np.empty(len(values), np.int32)
for index, value in enumerate(values):
family = _family(value)
code = families.get(family)
if code is None:
code = families[family] = len(children)
children.append([])
codes[index] = code
offsets[index] = len(children[code])
children[code].append(value)
return pa.UnionArray.from_dense(
pa.array(codes),
pa.array(offsets),
[_normalize(_child(child)) for child in children],
)
def _uniform(values: Sequence[Any], types: set[type]) -> pa.Array | None:
"""类型一致(或能无损地合成一种)时转成一个 Arrow 数组,否则返回空"""
if not types:
return pa.nulls(len(values))
if types == {str}:
return pa.array(values, pa.string())
if types == {bool}:
return pa.array(values, pa.bool_())
if types <= {int, float}:
if float not in types:
try:
return pa.array(values, pa.int64())
except (OverflowError, pa.ArrowException):
return None
# 整数和小数混在一起时存成 float64;整数超出 float 能精确表示的范围就不合并
if all(
-EXACT_FLOAT_INT < v < EXACT_FLOAT_INT for v in values if type(v) is int
):
return pa.array(values, pa.float64())
return None
if types <= {date, datetime}:
if any(type(v) is datetime and v.tzinfo is not None for v in values):
return None
if types == {date}:
return pa.array(values, pa.date32())
# 日期按零点存,类型推断时零点的时间仍算日期,和 classify 一致
return pa.array(
[datetime.combine(v, time()) if type(v) is date else v for v in values],
pa.timestamp("us"),
)
return None
def _family(value: Any) -> type:
"""union 里归为同一个子数组的类型:日期和不带时区的时间合在一起,超出 int64 的整数等存成文本,其余按 Python 类型分开"""
kind = type(value)
if kind is date or (kind is datetime and value.tzinfo is None):
return date
if kind is int and value not in INT64_RANGE:
return object
if kind in (NoneType, str, bool, int, float):
return kind
return object
def _child(values: list[Any]) -> pa.Array:
types = set(map(type, values))
types.discard(NoneType)
piece = _uniform(values, types)
if piece is not None:
return piece
# 超出 int64 的整数、带时区的时间、时长等:存成文本,classify 对它们的结论与文本相同
return pa.array([None if v is None else as_text(v) for v in values], pa.string())
+46 -22
View File
@@ -3,6 +3,7 @@
只读写本地文件,不碰数据库;上传流程在子进程里调用(worker.py),超时直接终止子进程
"""
import gc
from collections.abc import Callable
from pathlib import Path
from typing import Any, Literal
@@ -18,9 +19,9 @@ from app.datafiles.columns import (
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 (
Grid,
Records,
read_delimited,
read_excel,
@@ -37,23 +38,28 @@ ParseStage = Literal["detecting", "reading", "building"]
Progress = Callable[[ParseStage, float], None]
# 各格式识别、读取两个环节结束时的整体进度,剩下的是建表(拆表、推断类型、写 Parquet、算列概况)
# 数字来自 scripts/measure_parse_progress.py 的实测:各环节真实耗时和进度大致成正比,预计时间才准
# 数字来自 scripts/measure_parse_progress.py 在 100 万行上的实测:各环节真实耗时和进度大致成正比,预计时间才准
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),
"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.35
TYPED_AT = 0.8
WRITTEN_AT = 0.85
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:
@@ -89,6 +95,18 @@ def parse_file(
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]
@@ -111,7 +129,7 @@ def parse_file(
report: Reporter,
) -> None:
target = out_dir / f"{len(tables) + 1}.parquet"
pq.write_table(typed.table, target)
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(
@@ -130,6 +148,8 @@ def parse_file(
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):
@@ -138,7 +158,7 @@ def parse_file(
names, renamed = clean_names(table.header, table.sheet)
typed, kept = build_columns(
names,
list(zip(*table.rows)),
table.columns,
table.sheet,
scaled(share, 0.0, TYPED_AT),
)
@@ -163,18 +183,22 @@ def parse_file(
match detected.format:
case "csv" | "tsv":
grid = read_delimited(detected.text, detected.delimiter, reading)
add_grid(grid, building)
# 不留局部变量,网格才能在建表时用完即释放
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(len(grid.rows) for grid in grids)
total_rows = sum(grid.height for grid in grids)
done_rows = 0
for grid in grids:
while grids:
start = done_rows / total_rows
done_rows += len(grid.rows)
add_grid(grid, scaled(building, 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):
@@ -182,7 +206,7 @@ def parse_file(
else:
add_records(result, building)
case "jsonl":
add_records(read_jsonl(detected.text, reading), building)
add_records(read_jsonl(path, detected.codec, reading), building)
case "parquet":
typed, renamed = arrow_columns(read_parquet(path))
if typed.table.num_rows:
+55 -20
View File
@@ -1,12 +1,13 @@
"""列概况和预览:上传时在 Parquet 上用 DuckDB 计算,存进文件记录,供 describe_data、上传后的分析和接口返回"""
import math
import os
from pathlib import Path
from typing import Any
import duckdb
from app.datafiles.columns import EXACT_FLOAT_INT
from app.datafiles.cells import EXACT_FLOAT_INT
from app.datafiles.query import connect, json_value, quote_identifier
from app.schemas.data_file import ColumnProfile, DataColumn, TableProfile, TopValue
@@ -20,28 +21,30 @@ TOP_VALUES_MAX_DISTINCT = 50
EXAMPLE_SCAN_ROWS = 1000
PROFILE_CELL_CHARS = 50
PREVIEW_CELL_CHARS = 200
# 算列概况的线程数:在解析子进程里、只读自己刚写的文件,不受模型查询的线程限制;
# Parquet 按行组并行扫描,见 parse.py 的 ROW_GROUP_ROWS
PROFILE_THREADS = min(8, os.cpu_count() or 1)
def profile_table(
path: Path, columns: list[DataColumn], notes: dict[str, str]
) -> tuple[TableProfile, list[list[Any]]]:
"""返回列概况和接口用的预览(前 PREVIEW_ROWS 行)"""
connection = connect({"t": path})
connection = connect({"t": path}, threads=PROFILE_THREADS)
try:
connection.execute("SUMMARIZE t")
fields = [column[0] for column in connection.description]
summary = {row[0]: dict(zip(fields, row)) for row in connection.fetchall()}
total, statistics = _statistics(connection, columns)
head = connection.execute(f"FROM t LIMIT {EXAMPLE_SCAN_ROWS}").fetchall()
profiled = [
_profile_column(
connection,
column,
summary[column.name],
stats,
total,
head,
position,
notes.get(column.name),
)
for position, column in enumerate(columns)
for position, (column, stats) in enumerate(zip(columns, statistics))
]
sample = connection.execute(
f"FROM t USING SAMPLE reservoir({SAMPLE_ROWS} ROWS) REPEATABLE (42)"
@@ -84,24 +87,57 @@ def read_rows(
return [[json_value(value, PREVIEW_CELL_CHARS) for value in row] for row in rows]
def _statistics(
connection: duckdb.DuckDBPyConnection, columns: list[DataColumn]
) -> tuple[int, list[dict[str, Any]]]:
"""总行数和每列的统计值,一次扫描算完:非空数、去重数(估算),数值列的最小、最大、平均和四分位数,
日期列的最小、最大
四分位数用 quantile_disc 精确计算:近似算法在多线程下每次结果略有不同,同一个文件重新上传概况会变
"""
selects, layout = ["count(*)"], []
for column in columns:
name = quote_identifier(column.name)
fields = {
"count": f"count({name})",
"approx_unique": f"approx_count_distinct({name})",
}
if column.type == "number":
fields |= {
"min": f"min({name})",
"max": f"max({name})",
"avg": f"avg({name})",
"quartiles": f"quantile_disc({name}, [0.25, 0.5, 0.75])",
}
elif column.type in ("date", "datetime"):
fields |= {"min": f"min({name})", "max": f"max({name})"}
layout.append(list(fields))
selects.extend(fields.values())
row = connection.execute(f"SELECT {', '.join(selects)} FROM t").fetchone()
values = iter(row[1:])
return row[0], [dict(zip(fields, values)) for fields in layout]
def _profile_column(
connection: duckdb.DuckDBPyConnection,
column: DataColumn,
stats: dict[str, Any],
total: int,
head: list[tuple],
position: int,
note: str | None,
) -> ColumnProfile:
item: dict[str, Any] = {}
if column.type == "number":
item = {
key: _number(stats[key])
for key in ("min", "max", "avg", "q25", "q50", "q75")
quartiles = stats["quartiles"] or [None] * 3
item = {key: _number(stats[key]) for key in ("min", "max", "avg")}
item |= {
key: _number(value) for key, value in zip(("q25", "q50", "q75"), quartiles)
}
elif column.type in ("date", "datetime"):
# SUMMARIZE 的时间用空格分隔,和预览里的 ISO 写法统一
# 和预览里的时间写法一致
item = {
key: None if stats[key] is None else stats[key].replace(" ", "T", 1)
key: None if stats[key] is None else stats[key].isoformat()
for key in ("min", "max")
}
if (
@@ -124,23 +160,22 @@ def _profile_column(
examples.append(value)
if len(examples) == EXAMPLES:
break
null_ratio = float(stats["null_percentage"]) / 100
return ColumnProfile(
**column.model_dump(),
null_ratio=round(null_ratio, 4),
# approx_unique 是估算值,小表上可能比行数还大
distinct=min(stats["approx_unique"], round(stats["count"] * (1 - null_ratio))),
null_ratio=round(1 - stats["count"] / total, 4),
# approx_unique 是估算值,小表上可能比非空值的个数还大
distinct=min(stats["approx_unique"], stats["count"]),
examples=examples,
note=note,
**item,
)
def _number(text: str | None) -> int | float | None:
"""SUMMARIZE 的统计值是字符串;整数不带小数点,其余保留 6 位小数;Parquet 里的 inf 返回空"""
if text is None:
def _number(raw: Any) -> int | float | None:
"""整数不带小数点,其余保留 6 位小数;Parquet 里的 inf、NaN 返回空"""
if raw is None:
return None
value = float(text)
value = float(raw)
if not math.isfinite(value):
return None
if value.is_integer() and abs(value) < EXACT_FLOAT_INT:
+8 -3
View File
@@ -87,8 +87,13 @@ def run_query(
)
def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection:
"""建好视图并锁定的连接;列概况也用它,和模型查询看到的数据一致"""
def connect(
tables: Mapping[str, Path], *, threads: int = THREADS
) -> duckdb.DuckDBPyConnection:
"""建好视图并锁定的连接;列概况也用它,和模型查询看到的数据一致
threads 只给解析子进程算列概况时放宽,模型写的查询一律用 THREADS
"""
connection = duckdb.connect(
":memory:",
config={
@@ -96,7 +101,7 @@ def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection:
"autoload_known_extensions": False,
"python_enable_replacements": False,
"memory_limit": MEMORY_LIMIT,
"threads": THREADS,
"threads": threads,
},
)
try:
+182 -63
View File
@@ -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:
+109 -82
View File
@@ -1,13 +1,17 @@
"""在原始网格里识别表:按空行切块,找表头,合并两级表头,去掉标题、说明和合计行
只用规则;判断不了的地方照常产出表,同时记一条警告,调用方据此提示用户
要看每一行的规则(空行切块、合计行)对整列批量计算;表头、标题、说明只看每块开头或末尾的几行,取成 Python 值判断
"""
import math
import re
from dataclasses import dataclass, field
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
from app.datafiles.base import (
MAX_TABLE_ROWS,
Reporter,
@@ -16,8 +20,8 @@ from app.datafiles.base import (
rows_exceeded,
where,
)
from app.datafiles.columns import as_text, classify
from app.datafiles.readers import Grid
from app.datafiles.cells import WHITESPACE_CLASS, as_text, classify, mask
from app.datafiles.grid import Column, Grid
from app.schemas.data_file import ParseWarning
# 只在每块的前几行里找表头
@@ -40,6 +44,11 @@ TOTAL_LABELS = {
}
# 例如"华东小计""一季度合计"
TOTAL_SUFFIXES = ("小计", "合计", "总计")
_TOTAL_SET = pa.array(sorted(TOTAL_LABELS))
_TOTAL_STRIP = f"[{WHITESPACE_CLASS}::]"
_TOTAL_SUFFIX = f"(?:{'|'.join(TOTAL_SUFFIXES)})$"
# 每个合计字样(含后缀)至少有其中一个字符;英文按转小写之前的大小写都列出
_TOTAL_HINT = "[计和均tTmMeE]"
# 警告里引用的单元格文本最长显示这么多字
QUOTE_CHARS = 30
@@ -51,8 +60,8 @@ class RawTable:
# 同一 sheet 切出多张表时的序号,从 1 开始;只有一张表时为空
part: int | None
header: list[Any]
# 数据行,宽度与表头相同
rows: list[list[Any]]
# 数据列,和表头一一对应
columns: list[Column]
warnings: list[ParseWarning] = field(default_factory=list)
@@ -62,8 +71,8 @@ class _Piece:
# 表头所在行在网格中的下标
header_row: int
header: list[Any]
# 数据行在网格中的下标
data: list[int]
# 数据行在网格中的范围 [start, stop)
data: list[tuple[int, int]]
warnings: list[ParseWarning]
@@ -74,40 +83,41 @@ def split_tables(
几步都要扫一遍所有行,每步结束报告一次进度,比例来自大表上的实测耗时
"""
rows = _normalize(grid.rows)
report(0.25)
sheet = grid.sheet
present = [column.present() for column in grid.columns]
filled = np.logical_or.reduce(present) if present else np.zeros(grid.height, bool)
blocks = _blocks(filled)
used = _used_columns(present, blocks)
firsts = grid.rows_at(np.array([start for start, _ in blocks], np.int64))
report(0.3)
pieces: list[_Piece] = []
dropped: list[ParseWarning] = []
blocks = _blocks(rows)
report(0.3)
for number, block in enumerate(blocks, start=1):
columns = [
c for c in range(len(rows[0])) if any(rows[r][c] is not None for r in block)
]
for number, ((start, stop), columns, first) in enumerate(
zip(blocks, used, firsts), start=1
):
last = pieces[-1] if pieces else None
# 表中间夹着空行时,后面的块以数据行开头、列也落在上一张表内,接在上一张表后面
if (
last
and not _is_label_row([rows[block[0]][c] for c in columns])
and not _is_label_row([first[c] for c in columns])
and set(columns) <= set(last.columns)
):
last.data.extend(block)
last.data.append((start, stop))
continue
if len(block) < 2:
if stop - start < 2:
dropped.append(
ParseWarning(
code="block_dropped",
message=f"{where(sheet)}第 {block[0] + 1} 行的{_quote(rows[block[0]])}不是表格,已跳过",
message=f"{where(sheet)}第 {start + 1} 行的{_quote(first)}不是表格,已跳过",
)
)
continue
pieces.append(_read_block(rows, block, columns, sheet))
pieces.append(_read_block(grid, start, stop, columns, sheet))
report(0.3 + 0.25 * number / len(blocks))
tables = []
for number, piece in enumerate(pieces, start=1):
table = _finish(rows, piece, sheet)
table = _finish(grid, piece, sheet)
report(0.55 + 0.45 * number / len(pieces))
if table is None:
dropped.append(
@@ -125,41 +135,28 @@ def split_tables(
return tables, dropped
def _normalize(grid: list[list[Any]]) -> list[list[Any]]:
"""去掉字符串两端空白,空字符串记为 None,补齐成同样宽度"""
width = max((len(row) for row in grid), default=0)
rows = []
for row in grid:
cells = [_cell(value) for value in row]
cells.extend([None] * (width - len(cells)))
rows.append(cells)
return rows
def _blocks(filled: np.ndarray) -> list[tuple[int, int]]:
"""连续的非空行组成的块 [start, stop)"""
edges = np.flatnonzero(np.diff(np.concatenate(([0], filled, [0])).astype(np.int8)))
return list(zip(edges[0::2].tolist(), edges[1::2].tolist()))
def _cell(value: Any) -> Any:
if isinstance(value, str):
value = value.strip()
return value or None
return value
def _blocks(rows: list[list[Any]]) -> list[list[int]]:
blocks, current = [], []
for index, row in enumerate(rows):
if any(value is not None for value in row):
current.append(index)
elif current:
blocks.append(current)
current = []
if current:
blocks.append(current)
return blocks
def _used_columns(
present: list[np.ndarray], blocks: list[tuple[int, int]]
) -> list[list[int]]:
"""每块里有值的列;块之间只隔着空行,按块的起点分段求或即可"""
if not blocks:
return []
starts = np.array([start for start, _ in blocks])
used = np.array([np.logical_or.reduceat(column, starts) for column in present])
return [np.flatnonzero(used[:, index]).tolist() for index in range(len(blocks))]
def _read_block(
rows: list[list[Any]], block: list[int], columns: list[int], sheet: str | None
grid: Grid, start: int, stop: int, columns: list[int], sheet: str | None
) -> _Piece:
cells = [[rows[r][c] for c in columns] for r in block]
# 表头最多在第 HEADER_SEARCH_ROWS 行,两级表头还要看它下面两行
cells = grid.rows(start, min(stop, start + HEADER_SEARCH_ROWS + 2), columns)
need = max(1, math.ceil(len(columns) / 2))
dense = [sum(value is not None for value in row) >= need for row in cells]
search = range(min(HEADER_SEARCH_ROWS, len(cells)))
@@ -173,7 +170,7 @@ def _read_block(
warnings.append(
ParseWarning(
code="header_uncertain",
message=f"{where(sheet)}没有找到明确的表头,用第 {block[header_at] + 1} 行作表头,可能识别错误",
message=f"{where(sheet)}没有找到明确的表头,用第 {start + header_at + 1} 行作表头,可能识别错误",
)
)
if header_at > 0:
@@ -183,7 +180,7 @@ def _read_block(
warnings.append(
ParseWarning(
code="title_skipped",
message=f"{where(sheet)}{_row_range(block[0], block[header_at - 1])}是标题,已跳过:{skipped}",
message=f"{where(sheet)}{_row_range(start, start + header_at - 1)}是标题,已跳过:{skipped}",
)
)
@@ -191,7 +188,7 @@ def _read_block(
data_at = header_at + 1
if (
confident
and header_at + 2 < len(cells)
and header_at + 2 < stop - start
and _is_subheader(cells[header_at], cells[header_at + 1], cells[header_at + 2])
):
header = _merge_header(cells[header_at], cells[header_at + 1])
@@ -199,20 +196,28 @@ def _read_block(
warnings.append(
ParseWarning(
code="header_merged",
message=f"{where(sheet)}{_row_range(block[header_at], block[header_at + 1])}合并为表头",
message=f"{where(sheet)}{_row_range(start + header_at, start + header_at + 1)}合并为表头",
)
)
return _Piece(columns, block[header_at], header, block[data_at:], warnings)
return _Piece(
columns, start + header_at, header, [(start + data_at, stop)], warnings
)
def _finish(rows: list[list[Any]], piece: _Piece, sheet: str | None) -> RawTable | None:
data = [(index, [rows[index][c] for c in piece.columns]) for index in piece.data]
def _finish(grid: Grid, piece: _Piece, sheet: str | None) -> RawTable | None:
ranges = [np.arange(start, stop) for start, stop in piece.data if start < stop]
rows = np.concatenate(ranges) if ranges else np.zeros(0, np.int64)
warnings = piece.warnings
# 表格下方只有一个文字单元格的行是说明,例如"注:数据来源于……"
notes = []
while data and len(piece.columns) >= 3 and _is_note(data[-1][1]):
notes.append(data.pop())
while len(rows) and len(piece.columns) >= 3:
last = int(rows[-1])
row = grid.rows(last, last + 1, piece.columns)[0]
if not _is_note(row):
break
notes.append((last, row))
rows = rows[:-1]
if notes:
notes.reverse()
shown = "".join(_quote(row) for _, row in notes)
@@ -224,22 +229,59 @@ def _finish(rows: list[list[Any]], piece: _Piece, sheet: str | None) -> RawTable
)
# 合计、小计行不是明细数据,留着会让求和重复计算
totals = [index for index, row in data if _is_total(row)]
if totals:
data = [(index, row) for index, row in data if not _is_total(row)]
listed = brief([str(index + 1) for index in totals])
totals = _totals(grid, piece.columns, rows)
if totals.any():
listed = brief([str(index + 1) for index in rows[totals].tolist()])
warnings.append(
ParseWarning(
code="total_rows_removed",
message=f"{where(sheet)}已排除 {len(totals)} 行合计、小计(第 {listed} 行)",
message=f"{where(sheet)}已排除 {int(totals.sum())} 行合计、小计(第 {listed} 行)",
)
)
rows = rows[~totals]
if not data:
if not len(rows):
return None
if len(data) > MAX_TABLE_ROWS:
if len(rows) > MAX_TABLE_ROWS:
raise rows_exceeded(sheet)
return RawTable(sheet, None, piece.header, [row for _, row in data], warnings)
columns = [grid.columns[c].take(rows) for c in piece.columns]
return RawTable(sheet, None, piece.header, columns, warnings)
def _totals(grid: Grid, columns: list[int], rows: np.ndarray) -> np.ndarray:
"""rows 里哪些是合计行:每行只看前两个非空单元格,其中有合计字样的文字就是"""
if not len(rows):
return np.zeros(0, bool)
low, high = int(rows[0]), int(rows[-1]) + 1
seen = np.zeros(high - low, np.int64)
total = np.zeros(high - low, bool)
for c in columns:
column = grid.columns[c].slice(low, high)
present = column.present()
open_ = present & (seen < 2)
if open_.any():
total |= open_ & column.matches(_is_total_label)
seen += present
if (seen >= 2).all():
break
return total[rows - low]
def _is_total_label(strings: pa.Array) -> pa.Array:
"""去掉空白和冒号、转小写后是合计字样;合计字样都含 _TOTAL_HINT 里的字符,先用它筛掉绝大多数值"""
found = np.zeros(len(strings), bool)
rows = np.flatnonzero(mask(pc.match_substring_regex(strings, _TOTAL_HINT)))
if len(rows):
key = pc.utf8_lower(
pc.replace_substring_regex(strings.take(rows), _TOTAL_STRIP, "")
)
found[rows] = mask(
pc.or_(
pc.is_in(key, value_set=_TOTAL_SET),
pc.match_substring_regex(key, _TOTAL_SUFFIX),
)
)
return pa.array(found)
def _is_label(value: Any) -> bool:
@@ -276,21 +318,6 @@ def _is_note(row: list[Any]) -> bool:
return len(values) == 1 and _is_label(values[0])
def _is_total(row: list[Any]) -> bool:
checked = 0
for value in row:
if value is None:
continue
if isinstance(value, str):
key = re.sub(r"[\s::]", "", value).lower()
if key in TOTAL_LABELS or key.endswith(TOTAL_SUFFIXES):
return True
checked += 1
if checked == 2:
return False
return False
def _quote(row: list[Any]) -> str:
"""一行里第一个非空单元格,截断后加引号"""
value = next((value for value in row if value is not None), None)