dump_graph_data.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. '''
  2. 这个脚本是用来从postgre数据库中导出图谱数据到json文件的。
  3. '''
  4. import sys,os
  5. current_path = os.getcwd()
  6. sys.path.append(current_path)
  7. from sqlalchemy import text
  8. from sqlalchemy.orm import Session
  9. import json
  10. #这个是数据库的连接
  11. from db.session import SessionLocal
  12. #两个会话,分别是读取节点和属性的
  13. db = SessionLocal()
  14. prop = SessionLocal()
  15. #图谱id
  16. version = 'ER'
  17. def get_props(ref_id):
  18. props = {}
  19. sql = """select prop_name, prop_value,prop_title from kg_props where ref_id=:ref_id"""
  20. result = prop.execute(text(sql), {'ref_id':ref_id})
  21. for record in result:
  22. prop_name, prop_value,prop_title = record
  23. props[prop_name] = prop_title + ":" +prop_value
  24. return props
  25. def get_entities():
  26. COUNT_SQL = "select count(*) from kg_nodes where version=:version"
  27. result = db.execute(text(COUNT_SQL), {'version': version})
  28. count = result.scalar()
  29. print("total nodes: ", count)
  30. entities = []
  31. batch = 100
  32. start = 1
  33. while start < count:
  34. sql = """select id,name,category from kg_nodes where version=:version order by id limit :batch OFFSET :start"""
  35. result = db.execute(text(sql), {'start':start, 'batch':batch, 'version': version})
  36. #["发热",{"type":"症状","description":"发热是诊断的主要目的,用于明确发热病因。"}]
  37. row_count = 0
  38. for row in result:
  39. id,name,category = row
  40. props = get_props(id)
  41. #description = props.get('description', props.get('standard_description', ''))
  42. entities.append([name,{'type':category, **props}])
  43. row_count += 1
  44. if row_count == 0:
  45. break
  46. start = start + row_count
  47. print("start: ", start, "row_count: ", row_count)
  48. with open(current_path+"\\web\\cached_data\\entities_med.json", "w", encoding="utf-8") as f:
  49. f.write(json.dumps(entities, ensure_ascii=False,indent=4))
  50. def get_names(src_id, dest_id):
  51. sql = """select name from kg_nodes where id = :src_id"""
  52. result = db.execute(text(sql), {'src_id':src_id})
  53. src_name = result.scalar()
  54. result = db.execute(text(sql), {'src_id':dest_id})
  55. dest_name = result.scalar()
  56. return (src_name, dest_name)
  57. def get_relationships():
  58. COUNT_SQL = "select count(*) from kg_edges where version=:version"
  59. result = db.execute(text(COUNT_SQL), {'version': version})
  60. count = result.scalar()
  61. print("total edges: ", count)
  62. edges = []
  63. batch = 1000
  64. start = 1
  65. file_index = 1
  66. while start < count:
  67. sql = """select id,name,category,src_id,dest_id from kg_edges where version=:version order by id limit :batch OFFSET :start"""
  68. result = db.execute(text(sql), {'start':start, 'batch':batch, 'version': version})
  69. #["发热",{"type":"症状","description":"发热是诊断的主要目的,用于明确发热病因。"}]
  70. row_count = 0
  71. for row in result:
  72. id,name,category,src_id,dest_id = row
  73. props = get_props(id)
  74. src_name, dest_name = get_names(src_id, dest_id)
  75. edges.append([src_name, dest_name,{'type':category, **props}])
  76. row_count += 1
  77. if row_count == 0:
  78. break
  79. start = start + row_count
  80. print("start: ", start, "row_count: ", row_count)
  81. if len(edges) > 10000:
  82. with open(current_path+f"\\web\\cached_data\\relationship_med_{file_index}.json", "w", encoding="utf-8") as f:
  83. f.write(json.dumps(edges, ensure_ascii=False,indent=4))
  84. edges = []
  85. file_index += 1
  86. with open(current_path+"\\web\\cached_data\\relationship_med.json", "w", encoding="utf-8") as f:
  87. f.write(json.dumps(edges, ensure_ascii=False,indent=4))
  88. #导出节点数据
  89. get_entities()
  90. #导出关系数据
  91. get_relationships()