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])