my-v0-server/app/datafiles/profile.py

"""列概况和预览：上传时在 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)