350 lines
12 KiB
Python
350 lines
12 KiB
Python
"""列名清理和类型推断，输出 Arrow 列
|
||
|
||
CSV 的单元格全是字符串；Excel 的数值、日期单元格直接采用，文本单元格再按规则推断
|
||
一列的非空值必须全部符合同一类型才转换，否则保留文本：宁可让模型在 SQL 里 TRY_CAST，也不悄悄丢掉值
|
||
"""
|
||
|
||
import json
|
||
import math
|
||
import re
|
||
from collections import Counter
|
||
from collections.abc import Sequence
|
||
from dataclasses import dataclass
|
||
from datetime import UTC, date, datetime, time
|
||
from typing import Any
|
||
|
||
import pyarrow as pa
|
||
|
||
from app.datafiles.base import where
|
||
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")
|
||
|
||
|
||
@dataclass
|
||
class TypedColumns:
|
||
columns: list[DataColumn]
|
||
table: pa.Table
|
||
# 列名 → 写进列概况的说明，例如保留为文本的原因
|
||
notes: dict[str, str]
|
||
|
||
|
||
def clean_names(
|
||
raw: Sequence[Any], sheet: str | None
|
||
) -> tuple[list[str], ParseWarning | None]:
|
||
"""去掉空白和换行；空列名按位置命名，重复列名加后缀"""
|
||
names, changes, seen = [], [], set()
|
||
for position, value in enumerate(raw, start=1):
|
||
name = re.sub(r"\s*[\r\n]+\s*", " ", label_text(value)).strip()
|
||
if not name:
|
||
name = f"列{position}"
|
||
changes.append(f"第 {position} 列没有列名，命名为「{name}」")
|
||
if name in seen:
|
||
suffix = 2
|
||
while f"{name}_{suffix}" in seen:
|
||
suffix += 1
|
||
changes.append(f"「{name}」重名，第 {position} 列改为「{name}_{suffix}」")
|
||
name = f"{name}_{suffix}"
|
||
seen.add(name)
|
||
names.append(name)
|
||
if not changes:
|
||
return names, None
|
||
shown = "；".join(changes[:5]) + ("等" if len(changes) > 5 else "")
|
||
return names, ParseWarning(code="column_renamed", message=f"{where(sheet)}{shown}")
|
||
|
||
|
||
def label_text(value: Any) -> str:
|
||
"""表头单元格转成列名；Excel 里年份等数字表头读出来是 2025.0"""
|
||
if value is None:
|
||
return ""
|
||
if isinstance(value, float) and value.is_integer():
|
||
return str(int(value))
|
||
if isinstance(value, (date, time)):
|
||
return value.isoformat()
|
||
return str(value)
|
||
|
||
|
||
def build_columns(
|
||
names: list[str], values: list[Sequence[Any]], sheet: str | None
|
||
) -> tuple[TypedColumns, ParseWarning | None]:
|
||
"""按列推断类型；有列因少数无法解析的值保留为文本时，返回一条 kept_as_text 警告"""
|
||
columns, arrays, notes = [], [], {}
|
||
for name, column_values in zip(names, values):
|
||
column, array, note = infer_column(name, column_values)
|
||
columns.append(column)
|
||
arrays.append(array)
|
||
if note:
|
||
notes[name] = note
|
||
warning = None
|
||
if notes:
|
||
listed = "、".join(f"「{name}」" for name in notes)
|
||
warning = ParseWarning(
|
||
code="kept_as_text",
|
||
message=f"{where(sheet)}{listed}有无法解析的值，整列按文本处理",
|
||
)
|
||
return TypedColumns(
|
||
columns, pa.Table.from_arrays(arrays, names=names), notes
|
||
), warning
|
||
|
||
|
||
def arrow_columns(table: pa.Table) -> tuple[TypedColumns, ParseWarning | None]:
|
||
"""Parquet 自带类型，只映射到接口类型；列表、嵌套结构等存成 JSON 字符串"""
|
||
names, warning = clean_names(table.column_names, None)
|
||
columns, arrays = [], []
|
||
for name, array in zip(names, table.columns):
|
||
kind = _arrow_kind(array.type)
|
||
if kind is None:
|
||
values = array.to_pylist()
|
||
array = pa.array(
|
||
[None if v is None else as_text(v) for v in values], pa.string()
|
||
)
|
||
kind = "text"
|
||
elif kind == "text" and pa.types.is_dictionary(array.type):
|
||
array = array.cast(pa.string())
|
||
columns.append(DataColumn(name=name, type=kind))
|
||
arrays.append(array)
|
||
return TypedColumns(columns, pa.Table.from_arrays(arrays, names=names), {}), warning
|
||
|
||
|
||
def _arrow_kind(arrow_type: pa.DataType) -> str | None:
|
||
types = pa.types
|
||
if (
|
||
types.is_integer(arrow_type)
|
||
or types.is_floating(arrow_type)
|
||
or types.is_decimal(arrow_type)
|
||
):
|
||
return "number"
|
||
if types.is_boolean(arrow_type):
|
||
return "boolean"
|
||
if types.is_date(arrow_type):
|
||
return "date"
|
||
if types.is_timestamp(arrow_type):
|
||
return "datetime"
|
||
if (
|
||
types.is_string(arrow_type)
|
||
or types.is_large_string(arrow_type)
|
||
or types.is_string_view(arrow_type)
|
||
):
|
||
return "text"
|
||
if types.is_dictionary(arrow_type) and _arrow_kind(arrow_type.value_type) == "text":
|
||
return "text"
|
||
return 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)
|
||
|
||
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,
|
||
)
|
||
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
|
||
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,
|
||
)
|
||
|
||
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
|
||
|
||
|
||
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 parse_number(text: str) -> int | float | None:
|
||
"""千分位和货币符号可以有；有前导零或 16 位以上的数字串是编号，不算数值"""
|
||
# 最常见的纯数字走快速路径，不用正则
|
||
if text.isdigit():
|
||
if len(text) >= MIN_ID_DIGITS or (len(text) > 1 and text[0] == "0"):
|
||
return None
|
||
return int(text)
|
||
if text and text[0] in CURRENCY_SYMBOLS:
|
||
text = text[1:].lstrip()
|
||
if not NUMBER.match(text):
|
||
return None
|
||
digits = text.lstrip("+-")
|
||
if len(digits) > 1 and digits[0] == "0" and digits[1].isdigit():
|
||
return None
|
||
if digits.isdigit() and len(digits) >= MIN_ID_DIGITS:
|
||
return None
|
||
text = text.replace(",", "")
|
||
if INTEGER.match(text):
|
||
return int(text)
|
||
return float(text)
|
||
|
||
|
||
def parse_datetime(text: str) -> date | datetime | None:
|
||
if len(text) < 6 or not text[:4].isdigit():
|
||
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 as_text(value: Any) -> str:
|
||
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)):
|
||
return json.dumps(value, ensure_ascii=False)
|
||
return str(value)
|
||
|
||
|
||
def _naive(value: datetime) -> datetime:
|
||
# 带时区的时间统一换成 UTC 再去掉时区，同一列里的值才能比较
|
||
if value.tzinfo is None:
|
||
return value
|
||
return value.astimezone(UTC).replace(tzinfo=None)
|
||
|
||
|
||
def _text_note(
|
||
texts: list[str | None], value_kinds: list[str], kinds: Counter[str]
|
||
) -> 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)
|
||
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)
|
||
if len(examples) == 2:
|
||
break
|
||
shown = "、".join(f"「{example}」" for example in examples)
|
||
label = "数值" if numeric >= dates else "日期"
|
||
return f"看起来是{label}，有 {failed} 个值无法解析，例如{shown}"
|