class AddInitialGrantFlag(object): def __init__(self, load_config): self.load_config = load_config self.grant_num_groups = {} self.data_loader_utils = DataLoaderUtils(SERVER, INDEX, TYPE) self.bulk_data_size = 300000 self.docs_processed = 0 self.total_doc_count = 0 def process_doc(self, _id, doc): if 'grant_num' in doc: grant_num = doc['grant_num'] grant_num_comps = grant_num.split('-') if grant_num_comps[0] not in self.grant_num_groups: self.grant_num_groups[grant_num_comps[0]] = {} fy = None if 'fy' in doc: fy = doc['fy'] if len(fy) > 0: fy = int(fy) else: fy = None if fy is not None: self.grant_num_groups[grant_num_comps[0]][_id] = { 'id': _id, 'fy': fy, 'grant_num': grant_num } updated_doc = {} updated_doc['initial_grant'] = False return updated_doc def process_grant_num_groups(self): print 'Processing', len(self.grant_num_groups), 'grant_num groups' bulk_data = '' total_grant_num_groups = len(self.grant_num_groups) count = 0 for grant_num in self.grant_num_groups: count += 1 progress = ((count / float(total_grant_num_groups)) * 100) print 'Pass 2: progress', count, '/', total_grant_num_groups, progress, '%' grant_num_group = self.grant_num_groups[grant_num] # Find doc with lowest fy lowest_item = None for _id in grant_num_group: fy_data = grant_num_group[_id] if lowest_item is None: lowest_item = fy_data else: fy = fy_data['fy'] lowest_item_fy = lowest_item['fy'] if fy < lowest_item_fy: lowest_item = fy_data if lowest_item is not None: _id = lowest_item['id'] doc = {} doc['initial_grant'] = True bulk_data += self.data_loader_utils.bulk_update_header(_id) bulk_data += '\n' doc = {'doc': doc} bulk_data += json.dumps(doc) bulk_data += '\n' if len(bulk_data) >= self.bulk_data_size: self.load_bulk_data(bulk_data) bulk_data = '' if len(bulk_data) > 0: self.load_bulk_data(bulk_data) def load_bulk_data(self, bulk_data): self.data_loader_utils.load_bulk_data(bulk_data) # pass def run(self): doc_ids = get_doc_ids( server=self.load_config.server, src_index=self.load_config.index, src_type=self.load_config.type, dest_dir=self.load_config.other_files_directory(), dest_file_name="INITIAL_GRANT_ALL_IRDB_IDS.json") doc_ids = doc_ids.keys() self.total_doc_count = len(doc_ids) data_utils = DataUtils() data_utils.batch_fetch_docs_for_ids(base_url=self.load_config.server, ids=doc_ids, index=self.load_config.index, type=self.load_config.type, docs_fetched=self.docs_fetched) self.process_grant_num_groups() def docs_fetched(self, docs, index, type): docs_to_process = {} self.docs_processed += len(docs) progress = ((self.docs_processed / float(self.total_doc_count)) * 100) print 'Pass 1: progress', self.docs_processed, '/', self.total_doc_count, progress, '%' # print 'Docs fetched', len(docs) for doc in docs: _id = doc['_id'] if '_source' in doc: existing_doc = doc['_source'] docs_to_process[_id] = existing_doc self.process_docs(docs_to_process) # print 'Processed docs', self.processed_docs, 'Pubmed relations', len(self.pubmed_relations) def process_docs(self, docs): bulk_data = '' for _id in docs: doc = docs[_id] processed_doc = self.process_doc(_id, doc) if processed_doc is not None: bulk_data += self.data_loader_utils.bulk_update_header(_id) bulk_data += '\n' updated_doc = {'doc': processed_doc} bulk_data += json.dumps(updated_doc) bulk_data += '\n' if len(bulk_data) >= self.bulk_data_size: # print 'loading bulk data...' self.load_bulk_data(bulk_data) bulk_data = '' if len(bulk_data) > 0: # print 'loading bulk data...' self.load_bulk_data(bulk_data)
class ProcessIndex(object): def __init__(self, server, src_index, src_type, process_doc_method): self.server = server self.index = src_index self.type = src_type self.process_doc_method = process_doc_method self.batch_size = 5000 self.process_count = 2 self.process_spawn_delay = 0.15 self.bulk_data_size = 300000 self.data_loader_utils = DataLoaderUtils(self.server, self.index, self.type) self.data_utils = DataUtils() def run(self): # doc_ids = export_doc_ids( self.server, self.index, # self.type, self.index + '_' + self.type , 'doc_ids.json') doc_ids = file_utils.load_file(self.index, self.index + '_ids.json') if len(doc_ids) == 0: doc_ids = export_doc_ids.export_doc_ids(self.server, self.index, self.type) doc_ids = doc_ids.keys() batch_doc_processor = BatchDocProcessor(doc_ids, self.process_batch, self.batch_size, self.process_count, self.process_spawn_delay) batch_doc_processor.run() def docs_fetched(self, docs, index, type): docs_to_process = {} print 'Docs fetched', len(docs) for doc in docs: _id = doc['_id'] if '_source' in doc: existing_doc = doc['_source'] docs_to_process[_id] = existing_doc self.process_docs(docs_to_process) def process_docs(self, docs): bulk_data = '' for _id in docs: doc = docs[_id] processed_doc = self.process_doc_method(_id, doc) if processed_doc is not None: bulk_data += self.data_loader_utils.bulk_update_header(_id) bulk_data += '\n' updated_doc = {'doc': processed_doc} bulk_data += json.dumps(updated_doc) bulk_data += '\n' if len(bulk_data) >= self.bulk_data_size: # print 'loading bulk data...' self.load_bulk_data(bulk_data) bulk_data = '' if len(bulk_data) > 0: # print 'loading bulk data...' self.load_bulk_data(bulk_data) def load_bulk_data(self, bulk_data): self.data_loader_utils.load_bulk_data(bulk_data) # pass def process_batch(self, doc_ids): self.data_utils.batch_fetch_docs_for_ids( base_url=self.server, ids=doc_ids, index=self.index, type=self.type, docs_fetched=self.docs_fetched)
class CopyGrants(object): def __init__(self, src_server, dest_server, src_index, src_type, dst_index, dst_type, username, password): self.src_data_loader_utils = DataLoaderUtils(src_server, src_index, src_type) self.dest_data_loader_utils = DataLoaderUtils(dest_server, dst_index, dst_type) self.data_utils = DataUtils() self.username = username self.password = password file_utils.make_directory(TEMP_DIR) def run(self): self.process_batches() def process_batches(self): batch_file_names = [] for batch_file_name in os.listdir(TEMP_DIR): file_path = os.path.join(TEMP_DIR, batch_file_name) if os.path.isfile(file_path) and batch_file_name.startswith('batch_'): batch_file_names.append(batch_file_name) print "Generated ", len(batch_file_names), 'batch file names' batch_file_names.sort() if len(batch_file_names) == 0: batch_file_names = self.split_to_batches() print len(batch_file_names) raw_input('Continue?') processed_batches = file_utils.load_file(TEMP_DIR, 'processed_pubmed2018_docs_with_grants_batches.json') for batch_file_name in batch_file_names: if batch_file_name not in processed_batches: print 'Loading batch', batch_file_name batch = file_utils.load_file(TEMP_DIR, batch_file_name) self.copy_docs_batch(batch) processed_batches[batch_file_name] = 0 file_utils.save_file(TEMP_DIR, 'processed_pubmed2018_docs_with_grants_batches.json', processed_batches) def split_to_batches(self): server = self.src_data_loader_utils.server src_index = self.src_data_loader_utils.index src_type = self.src_data_loader_utils.type print 'Fetching doc ids for', src_index, src_type query = { "nested": { "path": "grants", "query": { "bool": { "must": [ { "exists": { "field": "grants" } } ] } } } } all_pubmed_ids = export_doc_ids.get_doc_ids(server, src_index, src_type, TEMP_DIR, 'pubmed2018_docs_with_grants.json', query=query) # all_pubmed_ids = all_pubmed_ids.keys() # all_pubmed_ids.sort() self.total_doc_count = len(all_pubmed_ids) max_batch_count = 5000 batch_file_names = [] batch_index = 0 batch_ids = [] # Splitting into batches for _id in all_pubmed_ids: batch_ids.append(_id) if len(batch_ids) >= max_batch_count: print 'Writing batch:', batch_index batch_file_name = 'batch_' + str(batch_index) + '.json' batch_file_names.append(batch_file_name) file_utils.save_file(TEMP_DIR, batch_file_name, batch_ids) batch_ids = [] batch_index += 1 if len(batch_ids) > 0: print 'Writing batch:', batch_index batch_file_name = 'batch_' + str(batch_index) + '.json' batch_file_names.append(batch_file_name) file_utils.save_file(TEMP_DIR, batch_file_name, batch_ids) batch_index += 1 return batch_file_names def copy_docs_batch(self, doc_ids): print 'Fetching docs' self.data_utils.batch_fetch_docs_for_ids(base_url=self.src_data_loader_utils.server, ids=doc_ids, index=self.src_data_loader_utils.index, type=self.src_data_loader_utils.type, docs_fetched=self.docs_fetched, batch_size=500) def docs_fetched(self, docs, index, type): print 'Docs fetched', len(docs) docs_to_copy = {} # print 'Docs fetched', len(docs) for doc in docs: _id = doc['_id'] if '_source' in doc: existing_doc = doc['_source'] docs_to_copy[_id] = existing_doc self.copy_relations(docs_to_copy) def load_bulk_data(self, bulk_data): print 'Bulk data size', len(bulk_data), 'loading...' response = self.dest_data_loader_utils.load_bulk_data(bulk_data) if response: pass # print 'Done loading bulk data, saving response' else: print 'Bulk data load failed' def copy_relations(self, src_docs): bulk_data = '' count = 0 # Copy relations for _id in src_docs: src_doc = src_docs[_id] doc = {} if 'grants' in src_doc: doc['grants'] = src_doc['grants'] count += 1 if len(doc) > 0: bulk_data += self.dest_data_loader_utils.bulk_update_header(_id) bulk_data += '\n' doc = { 'doc': doc } bulk_data += json.dumps(doc) bulk_data += '\n' # if count % 1000 == 0: # print 'Processed', 1000, 'docs' if len(bulk_data) >= 150000: print _id self.load_bulk_data(bulk_data) # print 'Copied', count, 'docs' bulk_data = '' if len(bulk_data) > 0: self.load_bulk_data(bulk_data) pass
class CopyRelationships(object): def __init__(self, src_server, dest_server, src_index, src_type, dst_index, dst_type, username, password): self.src_data_loader_utils = DataLoaderUtils(src_server, src_index, src_type) self.dest_data_loader_utils = DataLoaderUtils(dest_server, dst_index, dst_type) self.processed_doc_count = 0 self.total_doc_count = 0 self.data_utils = DataUtils() self.relations_to_exclude = [] self.missing_destination_ids = [] self.username = username self.password = password self.last_time_stamp = 0 self.diff_average = 0 def run(self): self.processed_doc_count = 0 self.total_doc_count = self.get_total_doc_count() print 'Total doc count', self.total_doc_count # self.create_destination_index(mapping=None) self.export_doc_ids(server=self.src_data_loader_utils.server, src_index=self.src_data_loader_utils.index, src_type=self.src_data_loader_utils.type) print 'saving missing docs' file_utils.save_file('/data/data_loading/pubmed_2019', 'missing_docs_pubmed2019.json', self.missing_destination_ids) def run_for_ids(self, doc_ids, mapping=None): self.processed_doc_count = 0 self.total_doc_count = len(doc_ids) print 'Total doc count', self.total_doc_count print 'Fetching docs from source index' batch_doc_processor = BatchDocProcessor(doc_ids, self.copy_docs_batch, 1000, 1, 0) batch_doc_processor.run() file_utils.save_file('/data/data_loading/pubmed_2019', 'missing_docs_pubmed2019.json', self.missing_destination_ids) def export_doc_ids(self, server, src_index, src_type): print 'Fetching doc ids for', src_index, src_type query = {"match_all": {}} self.data_utils.batch_fetch_ids_for_query(base_url=server, index=src_index, type=src_type, query=query, ids_fetched=self.ids_fetched, batch_size=10000) # print 'Done, fetched', len(documents_ids), 'doc ids' def ids_fetched(self, ids, index, type): print 'Ids fetched', len(ids) self.copy_docs_batch(ids) def copy_docs_batch(self, doc_ids): print 'Fetching docs' self.data_utils.batch_fetch_docs_for_ids( base_url=self.src_data_loader_utils.server, ids=doc_ids, index=self.src_data_loader_utils.index, type=self.src_data_loader_utils.type, docs_fetched=self.docs_fetched, batch_size=500) def docs_fetched(self, docs, index, type): print 'Docs fetched', len(docs) docs_to_copy = {} # print 'Docs fetched', len(docs) for doc in docs: _id = doc['_id'] if '_source' in doc: existing_doc = doc['_source'] docs_to_copy[_id] = existing_doc self.copy_relations(docs_to_copy) # Update progress self.processed_doc_count += len(docs) progress = ((self.processed_doc_count / float(self.total_doc_count)) * 100) current_time_stamp = time.time() diff = current_time_stamp - self.last_time_stamp self.diff_average = float(diff + self.diff_average) / 2 time_remaining = diff * (float(self.total_doc_count) / len(docs)) self.last_time_stamp = current_time_stamp print '---------------------------------------------------------------------------------------------' print 'Progress', self.processed_doc_count, '/', self.total_doc_count, progress, '%', time_remaining, 'secs' print '---------------------------------------------------------------------------------------------' def get_src_relations(self, src_doc, relationship_type): src_relations = [] if relationship_type in src_doc: relations = src_doc[relationship_type] for relation_item in relations: exclude_relation_item = False for relation_to_exclude in self.relations_to_exclude: if relation_to_exclude['source'] == relation_item[ 'source'] and relation_to_exclude[ 'index_id'] == relation_item['index_id']: exclude_relation_item = True break if not exclude_relation_item: src_relations.append(relation_item) return src_relations def get_dest_relations(self, dest_doc, relationship_type): dest_relations = [] if relationship_type in dest_doc: dest_relations = dest_doc[relationship_type] return dest_relations def add_relations(self, append_ids, relation, relations_list): relation_found = False for existing_relation in relations_list: # print existing_relation['source'], relation['source'], existing_relation['index_id'], relation['index_id'] if existing_relation['source'] == relation[ 'source'] and existing_relation['index_id'] == relation[ 'index_id']: existing_relation_ids = existing_relation['ids'] if append_ids: relation_ids = relation['ids'] for _id in relation_ids: if _id not in existing_relation_ids: existing_relation_ids.append(_id) existing_relation['ids'] = existing_relation_ids relation_found = True break if not relation_found: relations_list.append(relation) return relations_list def merge_relations(self, src_doc, dest_doc, relationship_type): dest_relations = self.get_dest_relations(dest_doc, relationship_type) src_relations = self.get_src_relations(src_doc, relationship_type) # print 'src_relations', len(src_relations) # print 'dest_relations', len(dest_relations) combined_relations = [] for relation in dest_relations: combined_relations = self.add_relations(True, relation, combined_relations) for relation in src_relations: combined_relations = self.add_relations(True, relation, combined_relations) return combined_relations def copy_relations(self, src_docs): bulk_data = '' count = 0 # Fetch destination docs destination_ids = src_docs.keys() destination_docs_array = self.data_utils.fetch_docs_for_ids( base_url=self.dest_data_loader_utils.server, ids=destination_ids, index=self.dest_data_loader_utils.index, type=self.dest_data_loader_utils.type, username=self.username, password=self.password) # Create destination doc dict destination_docs = {} for doc in destination_docs_array: _id = doc['_id'] if '_source' in doc: destination_docs[_id] = doc['_source'] # Find missing destination docs for _id in destination_ids: if _id not in destination_docs: self.missing_destination_ids.append(_id) print 'Missing ids', len(self.missing_destination_ids) # print 'dest ids', len() # Copy relations for _id in destination_docs: dest_doc = destination_docs[_id] src_doc = src_docs[_id] dest_relations = {} dest_relations[RELATIONSHIP_TYPE_CITATIONS] = self.merge_relations( src_doc, dest_doc, RELATIONSHIP_TYPE_CITATIONS) dest_relations[RELATIONSHIP_TYPE_CITED_BYS] = self.merge_relations( src_doc, dest_doc, RELATIONSHIP_TYPE_CITED_BYS) dest_relations[RELATIONSHIP_TYPE_RELATIONS] = self.merge_relations( src_doc, dest_doc, RELATIONSHIP_TYPE_RELATIONS) doc = {} if len(dest_relations[RELATIONSHIP_TYPE_CITATIONS]) > 0: doc[RELATIONSHIP_TYPE_CITATIONS] = dest_relations[ RELATIONSHIP_TYPE_CITATIONS] if len(dest_relations[RELATIONSHIP_TYPE_CITED_BYS]) > 0: doc[RELATIONSHIP_TYPE_CITED_BYS] = dest_relations[ RELATIONSHIP_TYPE_CITED_BYS] if len(dest_relations[RELATIONSHIP_TYPE_RELATIONS]) > 0: doc[RELATIONSHIP_TYPE_RELATIONS] = dest_relations[ RELATIONSHIP_TYPE_RELATIONS] # if len(dest_relations[RELATIONSHIP_TYPE_CITATIONS]) >= 2: # print _id count += 1 # doc = docs_to_copy[es_id] bulk_data += self.dest_data_loader_utils.bulk_update_header(_id) bulk_data += '\n' doc = {'doc': doc} bulk_data += json.dumps(doc) bulk_data += '\n' # if count % 1000 == 0: # print 'Processed', 1000, 'docs' if len(bulk_data) >= 150000: print _id self.load_bulk_data(bulk_data) # print 'Copied', count, 'docs' bulk_data = '' if len(bulk_data) > 0: self.load_bulk_data(bulk_data) pass # print 'Copied', count, 'docs' # def create_destination_index(self, mapping=None): # if mapping is None: # # Get mapping from src index # mapping = self.src_data_loader_utils.get_mapping_from_server() # if not self.dest_data_loader_utils.index_exists(): # print 'Creating index' # self.dest_data_loader_utils.put_mapping(mapping) # # migrate_index(self.dest_data_loader_utils.index) # else: # print self.dest_data_loader_utils.index, 'exists' def load_bulk_data(self, bulk_data): print 'Bulk data size', len(bulk_data), 'loading...' response = self.dest_data_loader_utils.load_bulk_data(bulk_data) if response: pass # print 'Done loading bulk data, saving response' else: print 'Bulk data load failed' def get_total_doc_count(self): return self.data_utils.get_total_doc_count( base_url=self.src_data_loader_utils.server, index=self.src_data_loader_utils.index, type=self.src_data_loader_utils.type) # src_server = 'http://localhost:9200' # src_index = 'pubmed2018_v5' # src_type = 'article' # dest_server = 'http://localhost:9200' # dest_index = 'pubmed2019' # dest_type = 'article' # copy_relations = CopyRelationships(src_server=src_server, # dest_server=dest_server, # src_index=src_index, # src_type=src_type, # dst_index=dest_index, # dst_type=dest_type, # username='', # password='') # copy_relations.relations_to_exclude.append({ # "source": "", # "index_id": ID_PUBMED # }) # copy_relations.run() # copy_relations.run_for_ids([12620793])