"""列概况和预览:上传时在 Parquet 上用 DuckDB 计算,存进文件记录,供 describe_data、上传后的分析和接口返回""" import math import os from pathlib import Path from typing import Any import duckdb 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 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 # 算列概况的线程数:在解析子进程里、只读自己刚写的文件,不受模型查询的线程限制; # 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}, threads=PROFILE_THREADS) try: total, statistics = _statistics(connection, columns) head = connection.execute(f"FROM t LIMIT {EXAMPLE_SCAN_ROWS}").fetchall() profiled = [ _profile_column( connection, column, stats, total, head, position, notes.get(column.name), ) for position, (column, stats) in enumerate(zip(columns, statistics)) ] 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 _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": 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"): # 和预览里的时间写法一致 item = { key: None if stats[key] is None else stats[key].isoformat() 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 return ColumnProfile( **column.model_dump(), null_ratio=round(1 - stats["count"] / total, 4), # approx_unique 是估算值,小表上可能比非空值的个数还大 distinct=min(stats["approx_unique"], stats["count"]), examples=examples, note=note, **item, ) def _number(raw: Any) -> int | float | None: """整数不带小数点,其余保留 6 位小数;Parquet 里的 inf、NaN 返回空""" if raw is None: return None value = float(raw) if not math.isfinite(value): return None if value.is_integer() and abs(value) < EXACT_FLOAT_INT: return int(value) return round(value, 6)