307 lines
11 KiB
Python
307 lines
11 KiB
Python
"""按列存放的单元格，读取、拆表和类型推断共用
|
||
|
||
每列是若干段 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())
|