feat:优化增量执行、工具调用容错与评测体系
This commit is contained in:
@@ -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
@@ -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
@@ -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 = []
|
||||
|
||||
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user