114 lines
4.8 KiB
Python
114 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
从达梦数据库导出三元组数据到 JSON 文件
|
||
三元组来源：RELATION_INSTANCE 表 (HEAD_ID, RELATION_TYPE, TAIL_ID)
|
||
支持可选清洗：去重、过滤悬空引用、移除空值、规范化关系类型
|
||
用法: python scripts/export_triplets.py
|
||
需设置环境变量 DM_PASSWORD，且 config 中 source=db
|
||
"""
|
||
import json
|
||
import os
|
||
import sys
|
||
|
||
_script_dir = os.path.dirname(os.path.abspath(__file__))
|
||
_project_root = os.path.dirname(_script_dir)
|
||
sys.path.insert(0, _project_root)
|
||
|
||
|
||
def _json_serializer(obj):
|
||
"""处理 Timestamp、datetime、numpy 等不可直接 JSON 序列化的类型"""
|
||
if hasattr(obj, 'isoformat'):
|
||
return obj.isoformat()
|
||
if hasattr(obj, 'item'):
|
||
return obj.item()
|
||
if isinstance(obj, float) and (obj != obj or obj == float('inf') or obj == float('-inf')):
|
||
return None
|
||
raise TypeError(f'Object of type {type(obj).__name__} is not JSON serializable')
|
||
|
||
|
||
def load_config(project_root: str) -> dict:
|
||
"""加载 config"""
|
||
config_dir = os.path.join(project_root, 'config')
|
||
json_path = os.path.join(config_dir, 'config.json')
|
||
yaml_path = os.path.join(config_dir, 'config.yaml')
|
||
if os.path.exists(json_path):
|
||
with open(json_path, 'r', encoding='utf-8') as f:
|
||
return json.load(f)
|
||
if os.path.exists(yaml_path):
|
||
try:
|
||
import yaml
|
||
with open(yaml_path, 'r', encoding='utf-8') as f:
|
||
return yaml.safe_load(f)
|
||
except ImportError:
|
||
pass
|
||
return {'source': 'json', 'db': {}, 'triplet': {}}
|
||
|
||
|
||
def main():
|
||
from src.data_loader import load_triplets_from_dm, load_from_dm
|
||
from src.triplet_cleaner import clean_triplets, transform_triplet_output_format
|
||
|
||
config = load_config(_project_root)
|
||
if config.get('source') != 'db':
|
||
print("错误：三元组导出仅支持达梦模式，请将 config 中 source 改为 db")
|
||
sys.exit(1)
|
||
|
||
tc = config.get('triplet', {})
|
||
output_file = tc.get('output_file', 'data/triples.json')
|
||
report_file = tc.get('report_file', 'report/triplet_cleaning_report.json')
|
||
output_path = output_file if os.path.isabs(output_file) else os.path.join(_project_root, output_file)
|
||
report_path = report_file if os.path.isabs(report_file) else os.path.join(_project_root, report_file)
|
||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||
os.makedirs(os.path.dirname(report_path), exist_ok=True)
|
||
|
||
# 清洗配置（clean: false 则跳过；否则按子项配置）
|
||
clean_cfg = tc.get('clean', True)
|
||
do_clean = False
|
||
if clean_cfg is True:
|
||
do_clean = {'deduplicate': True, 'filter_orphans': False, 'remove_null_core': True, 'normalize_relation_type': True}
|
||
elif isinstance(clean_cfg, dict):
|
||
do_clean = clean_cfg
|
||
|
||
deduplicate = do_clean.get('deduplicate', True) if do_clean else False
|
||
filter_orphans = do_clean.get('filter_orphans', False) if do_clean else False
|
||
remove_null_core = do_clean.get('remove_null_core', True) if do_clean else False
|
||
normalize_relation_type = do_clean.get('normalize_relation_type', True) if do_clean else False
|
||
|
||
print("正在从达梦读取三元组 (RELATION_INSTANCE)...")
|
||
triplets = load_triplets_from_dm(config, _project_root)
|
||
print(f"读取到 {len(triplets)} 条三元组")
|
||
|
||
valid_ids = None
|
||
if filter_orphans:
|
||
print("正在加载实体表以校验悬空引用...")
|
||
entities = load_from_dm(config, _project_root)
|
||
valid_ids = {str(e.get('TARGET_ID', '')).strip() for e in entities if e.get('TARGET_ID')}
|
||
print(f"有效实体数: {len(valid_ids)}")
|
||
|
||
if do_clean and (deduplicate or filter_orphans or remove_null_core or normalize_relation_type):
|
||
print("正在执行三元组清洗...")
|
||
triplets, report = clean_triplets(
|
||
triplets,
|
||
valid_ids=valid_ids,
|
||
deduplicate=deduplicate,
|
||
filter_orphans=filter_orphans,
|
||
remove_null_core=remove_null_core,
|
||
normalize_relation_type=normalize_relation_type,
|
||
)
|
||
with open(report_path, 'w', encoding='utf-8') as f:
|
||
json.dump(report, f, ensure_ascii=False, indent=2, default=_json_serializer)
|
||
print(f"清洗完成：{report['original_count']} -> {report['final_count']}，报告: {report_path}")
|
||
|
||
# 转为输出格式：head, head_type, relation, tail, tail_type
|
||
entities = load_from_dm(config, _project_root)
|
||
id_to_role = {str(e.get('TARGET_ID', '')).strip(): str(e.get('ROLE_ID', '')).strip() for e in entities if e.get('TARGET_ID')}
|
||
triplets = transform_triplet_output_format(triplets, id_to_role)
|
||
|
||
with open(output_path, 'w', encoding='utf-8') as f:
|
||
json.dump(triplets, f, ensure_ascii=False, indent=2, default=_json_serializer)
|
||
print(f"已导出至 {output_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|