""" 聚合函数处理模块 功能:处理统计类查询的聚合函数提取和结果检查 """ 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