142 lines
5.0 KiB
Python
142 lines
5.0 KiB
Python
"""把结构化的聚合请求编译成 DuckDB 的 SELECT：前端按列分组、汇总，不用自己写 SQL，也不用关心引号和类型
|
|
|
|
列名、别名一律通过 quote_identifier 引用，请求里的字符串不会拼进 SQL 的其他位置
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
from app.config import settings
|
|
from app.core.exceptions import BizException
|
|
from app.datafiles.query import quote_identifier, run_query
|
|
from app.schemas.data_file import (
|
|
AggregateRequest,
|
|
AggregateResult,
|
|
DataColumn,
|
|
GroupBy,
|
|
Metric,
|
|
)
|
|
|
|
# 只能用于数值列的聚合；其余聚合对列类型没有要求
|
|
NUMBER_ONLY = {"sum", "avg"}
|
|
TEMPLATES = {
|
|
"sum": "sum({})",
|
|
"avg": "avg({})",
|
|
"min": "min({})",
|
|
"max": "max({})",
|
|
"count": "count({})",
|
|
"count_distinct": "count(DISTINCT {})",
|
|
}
|
|
TIME_TYPES = {"date", "datetime"}
|
|
|
|
|
|
def run_aggregate(
|
|
path: Path, request: AggregateRequest, columns: list[DataColumn]
|
|
) -> AggregateResult:
|
|
"""对一张表的 Parquet 执行汇总；接口和上传后的分析共用
|
|
|
|
请求不合法时抛 422；执行失败（包括文件已被删除）抛 QueryError
|
|
"""
|
|
sql, names = _build_sql(request, columns)
|
|
result = run_query(
|
|
sql,
|
|
{"t": path},
|
|
timeout=settings.data_query_timeout,
|
|
max_rows=request.limit,
|
|
)
|
|
return AggregateResult(
|
|
columns=names,
|
|
rows=[dict(zip(names, row)) for row in result.rows],
|
|
truncated=result.truncated,
|
|
)
|
|
|
|
|
|
def _build_sql(
|
|
request: AggregateRequest, columns: list[DataColumn]
|
|
) -> tuple[str, list[str]]:
|
|
"""返回 SQL（表名固定为 t）和结果列名；请求不合法时抛 422
|
|
|
|
SQL 的 LIMIT 比请求多 1，调用方据此判断结果是否被截断
|
|
"""
|
|
by_name = {column.name: column for column in columns}
|
|
names: list[str] = []
|
|
selects: list[str] = []
|
|
|
|
for item in request.group_by:
|
|
column = _column(by_name, item.column)
|
|
selects.append(
|
|
f"{_group_expression(item, column)} AS {quote_identifier(item.column)}"
|
|
)
|
|
names.append(item.column)
|
|
|
|
for metric in request.metrics:
|
|
alias = metric.alias or _default_alias(metric)
|
|
selects.append(
|
|
f"{_metric_expression(metric, by_name)} AS {quote_identifier(alias)}"
|
|
)
|
|
names.append(alias)
|
|
|
|
duplicated = sorted({name for name in names if names.count(name) > 1})
|
|
if duplicated:
|
|
raise BizException(
|
|
422, f"结果列名重复：{'、'.join(duplicated)}，请给指标指定不同的 alias"
|
|
)
|
|
|
|
sql = f"SELECT {', '.join(selects)} FROM t"
|
|
if request.group_by:
|
|
positions = ", ".join(str(i) for i in range(1, len(request.group_by) + 1))
|
|
sql += f" GROUP BY {positions}"
|
|
sql += f" ORDER BY {_order(request, names)} LIMIT {request.limit + 1}"
|
|
return sql, names
|
|
|
|
|
|
def _column(by_name: dict[str, DataColumn], name: str) -> DataColumn:
|
|
if name not in by_name:
|
|
raise BizException(422, f"列不存在：{name}")
|
|
return by_name[name]
|
|
|
|
|
|
def _group_expression(item: GroupBy, column: DataColumn) -> str:
|
|
identifier = quote_identifier(item.column)
|
|
if item.grain is None:
|
|
return identifier
|
|
if column.type not in TIME_TYPES:
|
|
raise BizException(422, f"只有日期列才能按粒度分组：{item.column}")
|
|
# 都转成字符串，方便直接做图表的分类轴；字符串顺序和时间顺序一致
|
|
if item.grain == "quarter":
|
|
return (
|
|
f"strftime({identifier}, '%Y') || '-Q' || "
|
|
f"CAST(quarter({identifier}) AS VARCHAR)"
|
|
)
|
|
pattern = {"year": "%Y", "month": "%Y-%m", "day": "%Y-%m-%d"}[item.grain]
|
|
return f"strftime({identifier}, '{pattern}')"
|
|
|
|
|
|
def _metric_expression(metric: Metric, by_name: dict[str, DataColumn]) -> str:
|
|
if metric.column is None:
|
|
if metric.agg != "count":
|
|
raise BizException(422, f"{metric.agg} 需要指定 column")
|
|
return "count(*)"
|
|
column = _column(by_name, metric.column)
|
|
if metric.agg in NUMBER_ONLY and column.type != "number":
|
|
raise BizException(422, f"{metric.agg} 只能用于数值列：{metric.column}")
|
|
return TEMPLATES[metric.agg].format(quote_identifier(metric.column))
|
|
|
|
|
|
def _default_alias(metric: Metric) -> str:
|
|
"""汇总列默认沿用列名，直接对应图表物料按列名引用的写法；计数类加后缀避免和原列混淆"""
|
|
if metric.agg == "count":
|
|
return f"{metric.column}_count" if metric.column else "count"
|
|
if metric.agg == "count_distinct":
|
|
return f"{metric.column}_distinct"
|
|
return metric.column or metric.agg
|
|
|
|
|
|
def _order(request: AggregateRequest, names: list[str]) -> str:
|
|
if request.order_by is None:
|
|
# 有分组时按第一个分组列升序，结果顺序稳定；只有汇总时只有一行
|
|
return "1 ASC"
|
|
if request.order_by.by not in names:
|
|
raise BizException(422, f"排序列必须是结果里的列：{request.order_by.by}")
|
|
position = names.index(request.order_by.by) + 1
|
|
return f"{position} {request.order_by.direction.upper()} NULLS LAST"
|