262 lines
9.4 KiB
Python
262 lines
9.4 KiB
Python
"""列名清理和类型推断，输出 Arrow 列
|
||
|
||
单元格的类型规则在 cells.py；这里按列汇总：一列的非空值必须全部符合同一类型才转换，否则保留文本，
|
||
宁可让模型在 SQL 里 TRY_CAST，也不悄悄丢掉值。各列互不相关，用线程并行推断（Arrow 的计算不占 GIL）
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
from collections.abc import Sequence
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from dataclasses import dataclass
|
||
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
|
||
|
||
NUMERIC_KINDS = (INT, FLOAT, PERCENT)
|
||
DATE_KINDS = (DATE, DATETIME)
|
||
# 同时推断的列数
|
||
INFER_THREADS = min(8, os.cpu_count() or 1)
|
||
|
||
|
||
@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*", " ", as_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
|
||
return names, ParseWarning(
|
||
code="column_renamed", message=f"{where(sheet)}{brief(changes, '；')}"
|
||
)
|
||
|
||
|
||
def build_columns(
|
||
names: list[str],
|
||
columns: list[Column],
|
||
sheet: str | None,
|
||
report: Reporter = no_report,
|
||
) -> tuple[TypedColumns, ParseWarning | None]:
|
||
"""按列推断类型；有列因少数无法解析的值保留为文本时，返回一条 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)
|
||
warning = ParseWarning(
|
||
code="kept_as_text",
|
||
message=f"{where(sheet)}{listed}有无法解析的值，整列按文本处理",
|
||
)
|
||
return TypedColumns(
|
||
described, 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, 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}
|
||
|
||
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_()
|
||
)
|
||
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):
|
||
array = build(
|
||
lambda part: _pick(part, "moments", DATE_KINDS, pa.timestamp("us")),
|
||
pa.timestamp("us"),
|
||
)
|
||
return DataColumn(name=name, type="datetime"), array, None
|
||
|
||
array = build(_texts, pa.string())
|
||
return DataColumn(name=name, type="text"), array, _text_note(counts, parts)
|
||
|
||
|
||
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 _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 _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 _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 _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(counts: np.ndarray, parts: list[Classified]) -> str | None:
|
||
"""大部分值是数值或日期、少数无法解析时，说明原因，模型可以在 SQL 里用 TRY_CAST"""
|
||
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 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}"
|