166 lines
6.5 KiB
Python
166 lines
6.5 KiB
Python
"""
|
||
聚合函数处理模块
|
||
|
||
功能：处理统计类查询的聚合函数提取和结果检查
|
||
"""
|
||
|
||
import re
|
||
from typing import Dict, List, Any
|
||
|
||
|
||
def extract_aggregate_functions(cypher: str) -> Dict[str, str]:
|
||
"""
|
||
从Cypher查询中提取聚合函数及其别名
|
||
"""
|
||
if not cypher:
|
||
return {}
|
||
|
||
aggregate_map = {}
|
||
|
||
# 匹配RETURN子句（不区分大小写）
|
||
# 匹配到 ORDER BY, SKIP, LIMIT 或字符串结尾
|
||
return_match = re.search(r'\bRETURN\s+(.+?)(?:\s+ORDER\s+BY|\s+SKIP|\s+LIMIT|$)', cypher, re.IGNORECASE | re.DOTALL)
|
||
if not return_match:
|
||
return {}
|
||
|
||
return_clause = return_match.group(1).strip()
|
||
|
||
# 支持的聚合函数（不区分大小写）
|
||
aggregate_functions = ['COUNT', 'SUM', 'AVG', 'MAX', 'MIN', 'COLLECT', 'COLLECT_DISTINCT']
|
||
|
||
# 匹配聚合函数模式，支持以下格式：
|
||
# - COUNT(n) AS count
|
||
# - COUNT(n) AS `count`
|
||
# - COUNT(n) count
|
||
# - COUNT(DISTINCT n) AS count
|
||
# - SUM(n.value) AS total
|
||
# - COUNT(n) (没有别名，使用函数名)
|
||
|
||
# 先尝试匹配带AS的格式：COUNT(...) AS alias
|
||
as_pattern = r'\b(' + '|'.join(aggregate_functions) + r')\s*\([^)]*\)\s+AS\s+([`"]?)(\w+)\2'
|
||
as_matches = re.finditer(as_pattern, return_clause, re.IGNORECASE)
|
||
for match in as_matches:
|
||
func_name = match.group(1).upper()
|
||
alias = match.group(3)
|
||
aggregate_map[alias.lower()] = func_name
|
||
|
||
# 再尝试匹配不带AS但后面跟别名的格式：COUNT(...) alias
|
||
# 需要确保alias不是下一个聚合函数
|
||
no_as_pattern = r'\b(' + '|'.join(aggregate_functions) + r')\s*\([^)]*\)\s+([`"]?)(\w+)\2(?!\s*\()'
|
||
no_as_matches = re.finditer(no_as_pattern, return_clause, re.IGNORECASE)
|
||
for match in no_as_matches:
|
||
func_name = match.group(1).upper()
|
||
alias = match.group(3)
|
||
# 检查alias不是聚合函数名
|
||
if alias.upper() not in aggregate_functions:
|
||
aggregate_map[alias.lower()] = func_name
|
||
|
||
# 如果还没有找到，尝试匹配没有别名的格式：COUNT(...)
|
||
# 这种情况下，使用函数名本身作为别名
|
||
if not aggregate_map:
|
||
simple_pattern = r'\b(' + '|'.join(aggregate_functions) + r')\s*\([^)]*\)'
|
||
simple_matches = re.finditer(simple_pattern, return_clause, re.IGNORECASE)
|
||
for match in simple_matches:
|
||
func_name = match.group(1).upper()
|
||
# 检查后续是否有AS子句
|
||
after_func = return_clause[match.end():]
|
||
as_match = re.search(r'\bAS\s+([`"]?)(\w+)\1', after_func, re.IGNORECASE)
|
||
if as_match:
|
||
alias = as_match.group(2)
|
||
aggregate_map[alias.lower()] = func_name
|
||
else:
|
||
# 如果没有AS，使用函数名作为别名
|
||
aggregate_map[func_name.lower()] = func_name
|
||
|
||
return aggregate_map
|
||
|
||
|
||
def is_aggregate_result_zero(aggregate_results: List[Dict[str, Any]], cypher: str = None) -> bool:
|
||
"""
|
||
检查统计类查询结果是否为0
|
||
"""
|
||
if not aggregate_results:
|
||
return True
|
||
|
||
# 检查第一条记录的所有值
|
||
first_result = aggregate_results[0]
|
||
if not isinstance(first_result, dict):
|
||
return False
|
||
|
||
# 从Cypher中提取聚合函数信息
|
||
aggregate_map = {}
|
||
if cypher:
|
||
aggregate_map = extract_aggregate_functions(cypher)
|
||
|
||
# 如果没有提取到聚合函数信息，使用通用检查（兼容旧逻辑）
|
||
if not aggregate_map:
|
||
# 通用检查：任何值为None或0都认为是0结果
|
||
for key, value in first_result.items():
|
||
if value is None:
|
||
return True
|
||
if isinstance(value, (int, float)) and value == 0:
|
||
return True
|
||
if isinstance(value, str) and value.strip() == "0":
|
||
return True
|
||
if isinstance(value, list) and len(value) == 0:
|
||
return True
|
||
return False
|
||
|
||
# 基于提取的聚合函数进行检查
|
||
has_zero_aggregate = False
|
||
has_non_zero_value = False
|
||
|
||
for key, value in first_result.items():
|
||
key_lower = key.lower()
|
||
|
||
# 检查该字段是否对应某个聚合函数
|
||
func_name = aggregate_map.get(key_lower)
|
||
|
||
if func_name:
|
||
# 根据聚合函数类型进行判断
|
||
if func_name == 'COUNT':
|
||
# COUNT结果为0表示没有数据
|
||
if value is None or (isinstance(value, (int, float)) and value == 0):
|
||
has_zero_aggregate = True
|
||
elif isinstance(value, (int, float)) and value > 0:
|
||
has_non_zero_value = True
|
||
|
||
elif func_name == 'SUM':
|
||
# SUM结果为0或null表示没有数据
|
||
if value is None or (isinstance(value, (int, float)) and value == 0):
|
||
has_zero_aggregate = True
|
||
elif isinstance(value, (int, float)) and value != 0:
|
||
has_non_zero_value = True
|
||
|
||
elif func_name in ['AVG', 'MAX', 'MIN']:
|
||
# AVG/MAX/MIN结果为null表示没有数据
|
||
if value is None:
|
||
has_zero_aggregate = True
|
||
elif isinstance(value, (int, float)) and value != 0:
|
||
has_non_zero_value = True
|
||
# 注意：AVG为0可能是有效值，但通常null更常见
|
||
|
||
elif func_name in ['COLLECT', 'COLLECT_DISTINCT']:
|
||
# COLLECT结果为空列表表示没有数据
|
||
if value is None or (isinstance(value, list) and len(value) == 0):
|
||
has_zero_aggregate = True
|
||
elif isinstance(value, list) and len(value) > 0:
|
||
has_non_zero_value = True
|
||
else:
|
||
# 如果字段不在聚合函数映射中，但值不为0/null，说明可能有其他有效数据
|
||
if value is not None:
|
||
if isinstance(value, (int, float)) and value != 0:
|
||
has_non_zero_value = True
|
||
elif isinstance(value, str) and value.strip() and value.strip() != "0":
|
||
has_non_zero_value = True
|
||
elif isinstance(value, list) and len(value) > 0:
|
||
has_non_zero_value = True
|
||
|
||
# 如果有聚合函数结果为0/null，且没有其他非0值，认为是0结果
|
||
if has_zero_aggregate and not has_non_zero_value:
|
||
return True
|
||
|
||
# 如果所有聚合函数值都不为0或null，返回False
|
||
return False
|
||
|