149 lines
4.9 KiB
Python
149 lines
4.9 KiB
Python
"""列概况和预览：上传时在 Parquet 上用 DuckDB 计算，存进文件记录，供 describe_data、上传后的分析和接口返回"""
|
|
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import duckdb
|
|
|
|
from app.datafiles.columns import EXACT_FLOAT_INT
|
|
from app.datafiles.query import connect, json_value, quote_identifier
|
|
from app.schemas.data_file import ColumnProfile, DataColumn, TableProfile, TopValue
|
|
|
|
PREVIEW_ROWS = 20
|
|
SAMPLE_ROWS = 5
|
|
EXAMPLES = 3
|
|
TOP_VALUES = 5
|
|
# 去重值不超过这么多的文本、布尔列才给出现次数最多的值
|
|
TOP_VALUES_MAX_DISTINCT = 50
|
|
# 样例值从开头这么多行里取
|
|
EXAMPLE_SCAN_ROWS = 1000
|
|
PROFILE_CELL_CHARS = 50
|
|
PREVIEW_CELL_CHARS = 200
|
|
|
|
|
|
def profile_table(
|
|
path: Path, columns: list[DataColumn], notes: dict[str, str]
|
|
) -> tuple[TableProfile, list[list[Any]]]:
|
|
"""返回列概况和接口用的预览（前 PREVIEW_ROWS 行）"""
|
|
connection = connect({"t": path})
|
|
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()}
|
|
head = connection.execute(f"FROM t LIMIT {EXAMPLE_SCAN_ROWS}").fetchall()
|
|
profiled = [
|
|
_profile_column(
|
|
connection,
|
|
column,
|
|
summary[column.name],
|
|
head,
|
|
position,
|
|
notes.get(column.name),
|
|
)
|
|
for position, column in enumerate(columns)
|
|
]
|
|
sample = connection.execute(
|
|
f"FROM t USING SAMPLE reservoir({SAMPLE_ROWS} ROWS) REPEATABLE (42)"
|
|
).fetchall()
|
|
finally:
|
|
connection.close()
|
|
profile = TableProfile(
|
|
columns=profiled,
|
|
sample=[
|
|
[json_value(value, PROFILE_CELL_CHARS) for value in row] for row in sample
|
|
],
|
|
)
|
|
preview = [
|
|
[json_value(value, PREVIEW_CELL_CHARS) for value in row]
|
|
for row in head[:PREVIEW_ROWS]
|
|
]
|
|
return profile, preview
|
|
|
|
|
|
def read_rows(
|
|
path: Path,
|
|
offset: int,
|
|
limit: int,
|
|
sort_by: str | None = None,
|
|
descending: bool = False,
|
|
) -> list[list[Any]]:
|
|
"""读一页数据，单元格截断规则与预览相同；sort_by 必须是已校验过的列名，空值排最后"""
|
|
order = (
|
|
f" ORDER BY {quote_identifier(sort_by)} {'DESC' if descending else 'ASC'} NULLS LAST"
|
|
if sort_by
|
|
else ""
|
|
)
|
|
connection = connect({"t": path})
|
|
try:
|
|
rows = connection.execute(
|
|
f"FROM t{order} LIMIT {int(limit)} OFFSET {int(offset)}"
|
|
).fetchall()
|
|
finally:
|
|
connection.close()
|
|
return [[json_value(value, PREVIEW_CELL_CHARS) for value in row] for row in rows]
|
|
|
|
|
|
def _profile_column(
|
|
connection: duckdb.DuckDBPyConnection,
|
|
column: DataColumn,
|
|
stats: dict[str, Any],
|
|
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")
|
|
}
|
|
elif column.type in ("date", "datetime"):
|
|
# SUMMARIZE 的时间用空格分隔，和预览里的 ISO 写法统一
|
|
item = {
|
|
key: None if stats[key] is None else stats[key].replace(" ", "T", 1)
|
|
for key in ("min", "max")
|
|
}
|
|
if (
|
|
column.type in ("text", "boolean")
|
|
and stats["approx_unique"] <= TOP_VALUES_MAX_DISTINCT
|
|
):
|
|
name = quote_identifier(column.name)
|
|
counts = connection.execute(
|
|
f"SELECT {name}, count(*) AS n FROM t WHERE {name} IS NOT NULL "
|
|
f"GROUP BY 1 ORDER BY n DESC, 1 LIMIT {TOP_VALUES}"
|
|
).fetchall()
|
|
item["top_values"] = [
|
|
TopValue(value=json_value(value, PROFILE_CELL_CHARS), count=count)
|
|
for value, count in counts
|
|
]
|
|
examples: list[Any] = []
|
|
for row in head:
|
|
value = json_value(row[position], PROFILE_CELL_CHARS)
|
|
if value is not None and value not in examples:
|
|
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))),
|
|
examples=examples,
|
|
note=note,
|
|
**item,
|
|
)
|
|
|
|
|
|
def _number(text: str | None) -> int | float | None:
|
|
"""SUMMARIZE 的统计值是字符串；整数不带小数点，其余保留 6 位小数；Parquet 里的 inf 返回空"""
|
|
if text is None:
|
|
return None
|
|
value = float(text)
|
|
if not math.isfinite(value):
|
|
return None
|
|
if value.is_integer() and abs(value) < EXACT_FLOAT_INT:
|
|
return int(value)
|
|
return round(value, 6)
|