实习工作小结-拆分大文件功能和Elasticsearch文档添加新字段

前言

距离上次的博客已经有1个多月了，这个月里我主要精力放在了《计算理论》的期末考试上，28日终于考完了，考完感觉还不错，长舒一口气。

上个月，实习单位里的事情也没干多少，跟领导也解释了，领导挺开明的，也没有多怪我。上个月没工作几天，就只完成了两个功能：拆分大文件功能和Elasticsearch文档添加新字段。

大文件的拆分功能

这个文件一共90来G是读不进内存的，所以我们需要对其拆分。

具体是这样的，有一个大文件上面呢有如下字段：pmid, span_start, span_end, entity string, entity type, entity id，示例文件是这样的：

你会发现，不仅有我们需要的6个字段（分隔符是制表符），还会出现一大段的abstruct，但这个是我们不需要的内容。可以发现这个功能一点都不难，就一个函数就可以解决，下面给出代码：

def split_sample(address):
    # root_url = 'https://ftp.ncbi.nlm.nih.gov/pub/lu/PubTatorCentral/'
    # root_url = 's3://meta-adhoc/pubtator_annotation/2022-03-17/'
    dir_name = "big_files"
    # wget.download(f"{root_url}{address}", out=f"./{dir_name}/{address}")
    # os.system(f"aws s3 cp {root_url}{address} ./{dir_name}/{address}")
    with open(f"./{dir_name}/{address}", 'r') as f:
        # 获取数据

        id_list =[]
        new_file_list = []
        for data in f:
            tmp_data = data.split('\t')
            # 如果第一个元素长度为8，说明这个元素应该是id
            if len(tmp_data[0]) == 8:
                # 由于相同的id都是紧挨的，如果id不在id_list之中，说明换新id了，要重新开一个文件
                if tmp_data[0] not in id_list:

                    # 在此之前，先把原来的new_file_list里的文件先写入到新文件中
                    # 第一此写入时，new_file_list为空，故需要判断一下
                    if len(new_file_list) > 0:
                        try:
                            file_handle = open(f"./result_files/{new_file_list[0][0]}.txt", 'w+')
                            for new_file in new_file_list:
                                for element in new_file:
                                    if element.endswith("\n"):
                                        file_handle.writelines(element)
                                    else:
                                        file_handle.writelines(element+'\t')
                            file_handle.close()
                        except Exception:

                    # 添加新id到id_list,说明这个id已经遍历过了或者正在遍历
                    id_list.append(tmp_data[0])
                    # 清空已经写入的new_file_list
                    new_file_list.clear()
                    # 当前列写入new_file_list
                    if tmp_data is not None:
                        new_file_list.append(tmp_data)
                # 如果id在id_list之中，说明正在遍历此id，
                else:
                    # 将其加入到new_file_list之中
                    if tmp_data is not None:
                        new_file_list.append(tmp_data)
            else:
                pass
                            
       f.close()


if __name__ == "__main__":
    split_sample("bioconcepts2pubtatorcentral.offset.sample")

split_sample是主要的功能函数，只要传入文件名或者说文件地址即可。具体逻辑呢，就是先把文件读入，然后按行读取。

首先要判断读入的是需要的数据还是abstract。如果是abstruct的话，第0号元素长度必然会很长，而不是8，其余时候第0号元素是pmid长度为8，通过判断长度即可确定是abstract还是需要的数据。

然后为什么要搞一个id_list列表呢？是因为我们这个拆分功能是根据pmid来拆分的，每同一个pmid的数据都放在同一个文件中，名字为“pmid的值.txt”。然后后续update的时候只要根据pmid来操作就可以了。new_file_list是要写入新的“pmid的值.txt”文件中的数据，如果是第一次遍历该pmid值，那么new_file_list必然是空的，那我们需要将pmid值加入id_list，标志这个pmid的相关数据正要遍历。同时，new_file_list为空说明什么？上个文件刚刚遍历完，所有数据都加入了new_file_list，我们需要将new_file_list的数据一个个写入新文件中，名字为“pmid的值.txt”。反之不为空，则id_list里面有了这个pmid那说明该pmid正在遍历（因为相同pmid数据是连一起的），如果该pmid正在遍历，就把数据按行加入new_file_list。

大文件拆分功能就是这么多，柴拆完文件，我们需要把这些数据导入到elasticsearch中去。

给Elasticsearch文档添加新字段

在前几篇的博文中，我将数据导入到了elasticsearch去了，但是那个xml文件里的数据并不是全部的数据，现在拆完的这个大文件还有一些其他的字段需要加入到elasticsearch的每个文档中去。

首先先展示下update后的文档的样子：

 可以看到多了一个"entities"字段，"entities"字段中每一个元素就是一个“pmid的值.txt”中的一行数据，最终我们目的是实现通过检索疾病等类型来获得文章的abstract和title、author等信息。

首先我们需要通过pmid检索出对应的文档id，以方便下一步操作：

def get_id_from_pmid(file_name):
    body1 = {
        "query": {
            "match": {
                "pmid": file_name
            }
        },
            "track_total_hits": "true",
    }

    size = 10
    es = Elasticsearch(hosts=['http://localhost:9200'], request_timeout=3600)

    index_name = "pubmed-paper-index-2"
    result = es.search(index=index_name, body=body1)
    total_value = result['hits']['total']['value']

    for i in range(total_value):
        body = {
            "query": {
                "match": {
                    "pmid": file_name
                }
            },
            "track_total_hits": "true",
            "from": i * size,
            "size": size
        }
        result = es.search(index=index_name, body=body)
        tmps = result['hits']['hits']
        for tmp in tmps:
            # print("outside", tmp['_id'], tmp['_source']['pmid'])
            update_file(tmp['_id'], tmp['_source']['pmid'])

可以看到，这里使用了分页操作。通过传给update_file函数以tmp['_id']和tmp['_source']['pmid']即id和pmid来实现，不是很难，就一个检索和分页操作而已。

接着来看update_file函数：

def update_file(id, pmid):
    '''
    id是随机数，是es中doc的id
    pmid是确定的，是文章的pmid
    '''

    es = Elasticsearch(hosts=['http://localhost:9200'], request_timeout=3600)
    try:
        datas = open(f"result_files/{pmid}.txt").readlines()
        print(id, pmid)
        logging.basicConfig()
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.INFO)
        handler = logging.FileHandler("log_es_update.txt")
        logger.addHandler(handler)
        logger.info(pmid+".txt")
        # pmid, span_start, span_end, entity string, entity type, entity id
        update_body = {
            "doc": {
                "entities": []
            }
        }
        for data in datas:
            data = data.replace("\n", "")
            data_list = data.split("\t")
            # print(data_list)

            if len(data_list) == 6:

                body = {
                    "span": [data_list[1], data_list[2]],
                    "entity_string": data_list[3],
                    "entity_type": data_list[4],
                    "entity_id": data_list[5],
                }
                update_body['doc']['entities'].append(body)
            # print(update_body)
        es.update(index="pubmed-paper-index-2", body=update_body, id=id)

        # 去重操作
        number_list = []
        with open('log_es_update.txt', 'r') as file1:
            for data in file1:
                if data not in number_list:
                    number_list.append(data)
                    # print(data)
        file1.close()
        with open('log_es_update.txt', 'w') as file1:
            for data in number_list:
                file1.write(data)
        file1.close()
    except Exception as e:
        print(e)
        pass

doc的id是随机数，应为我们一开始没有决定用哪个来当id所以干脆让elasticsearch自己生成随机值。而一个pmid可能有多个doc所以得通过pmid检索id然后再通过每个doc的id来更新doc。这里有个去重操作，原因是是我们使用了log_es_update.txt文件来记录已经update过的pmid，这样下次更新就可以接着上次没更新的继续更新。为什么要去重？应为每次更新一个doc就会将这个doc的pmid写入log_es_update.txt，但是有多个doc的pmid是一样的，所以能需要去重。

接下来就是增量更新的功能，具体来说就是看看哪个更新了，那个没更新过。更新过了的就不更新了，没更新的再去更新。代码如下：

def update_data():
    # update_file(id, pmid)
    # 读取已经加载的列表
    already_loaded = []
    with open("log_es_update.txt") as f:
        files = f.readlines()
        for file in files:
            # print(file.replace("\n", ""))
            already_loaded.append(file.replace("\n", ""))

    # 检测本地已经下载的文件名，并生成列表
    local_address = []
    tmp_daily = glob.glob('result_files/*.txt')
    for file in tmp_daily:
        local_address.append(file.replace("result_files\\", ""))

    logging.basicConfig()
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    handler = logging.FileHandler("log_es_update.txt")
    logger.addHandler(handler)

    for file in local_address:
        if file not in already_loaded:
            get_id_from_pmid(file.replace(".txt", ""))

这个函数功能很简单，就不再赘述了。基本想法就是通过比较已更新的列表和全部更新的列表，更新没更新的那一部分。

全部代码如下：

import glob
from elasticsearch import Elasticsearch
import logging


def update_data():
    # update_file(id, pmid)
    # 读取已经加载的列表
    already_loaded = []
    with open("log_es_update.txt") as f:
        files = f.readlines()
        for file in files:
            # print(file.replace("\n", ""))
            already_loaded.append(file.replace("\n", ""))

    # 检测本地已经下载的文件名，并生成列表
    local_address = []
    tmp_daily = glob.glob('result_files/*.txt')
    for file in tmp_daily:
        local_address.append(file.replace("result_files\\", ""))

    logging.basicConfig()
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    handler = logging.FileHandler("log_es_update.txt")
    logger.addHandler(handler)

    for file in local_address:
        if file not in already_loaded:
            get_id_from_pmid(file.replace(".txt", ""))


def get_id_from_pmid(file_name):
    body1 = {
        "query": {
            "match": {
                "pmid": file_name
            }
        },
            "track_total_hits": "true",
    }

    size = 10
    es = Elasticsearch(hosts=['http://localhost:9200'], request_timeout=3600)

    index_name = "pubmed-paper-index-2"
    result = es.search(index=index_name, body=body1)
    total_value = result['hits']['total']['value']

    for i in range(total_value):
        body = {
            "query": {
                "match": {
                    "pmid": file_name
                }
            },
            "track_total_hits": "true",
            "from": i * size,
            "size": size
        }
        result = es.search(index=index_name, body=body)
        tmps = result['hits']['hits']
        for tmp in tmps:
            # print("outside", tmp['_id'], tmp['_source']['pmid'])
            update_file(tmp['_id'], tmp['_source']['pmid'])


def update_file(id, pmid):
    '''
    id是随机数，是es中doc的id
    pmid是确定的，是文章的pmid
    '''

    es = Elasticsearch(hosts=['http://localhost:9200'], request_timeout=3600)
    try:
        datas = open(f"result_files/{pmid}.txt").readlines()
        print(id, pmid)
        logging.basicConfig()
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.INFO)
        handler = logging.FileHandler("log_es_update.txt")
        logger.addHandler(handler)
        logger.info(pmid+".txt")
        # pmid, span_start, span_end, entity string, entity type, entity id
        update_body = {
            "doc": {
                "entities": []
            }
        }
        for data in datas:
            data = data.replace("\n", "")
            data_list = data.split("\t")
            # print(data_list)

            if len(data_list) == 6:

                body = {
                    "span": [data_list[1], data_list[2]],
                    "entity_string": data_list[3],
                    "entity_type": data_list[4],
                    "entity_id": data_list[5],
                }
                update_body['doc']['entities'].append(body)
            # print(update_body)
        es.update(index="pubmed-paper-index-2", body=update_body, id=id)

        # 去重操作
        number_list = []
        with open('log_es_update.txt', 'r') as file1:
            for data in file1:
                if data not in number_list:
                    number_list.append(data)
                    # print(data)
        file1.close()
        with open('log_es_update.txt', 'w') as file1:
            for data in number_list:
                file1.write(data)
        file1.close()
    except Exception as e:
        print(e)
        pass


if __name__ == "__main__":

    update_data()
    # test
    # update_file('4z4PnYABkau2ZVG6RHpl', '35174437')

END