需求背景

从Elasticsearch(以下简称ES)拉产品需要的数据并生成.csv文件, 由于数据量可能会很大, 采用循环的方式执行. 我之前一直用Java来实现这个功能, 但是Java开发成本大, 部署比较麻烦, 趁着有空来研究下Python的实现. 下面把我这个Python初学者的实现过程分享给诸位, 经历了很多坑, 文末会贴出完整的代码.

代码实现

连接ES平台

Python连接ES需要使用官方数据包elasticsearch, 如果没有可以使用pip安装. 这里有个坑: 建议指定一个高版本安装, 低版本在连接ES时出现了问题. 指定版本安装方式如下:

pip install elasticsearch==6.3.1

#-*- coding=utf-8'*-
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
import csv
import os
import sys

es = Elasticsearch(["http://**.**.**.**:9210"])
# 查看连接是否成功
print(es.info())

**换成公司的ES服务IP地址, 这里只连接了测试环境的IP, 我们公司的测试环境不需要用户身份信息. 如果还想做的健全一点, 可以做个环境切换, 条件由入参决定. 若连接成功, print会打印出ES注册和版本相关信息, 这里便不展示了

查询条件

本文采用原生的ES搜索条件, 未进行任何包装, 所以查询条件看上去会比较绕. 这里会有个坑: ES分页查询支持的偏移量最多只有10000, 如果超过一万会报错:

TransportError(500, u'search_phase_execution_exception', u'Result window is too large, from + size must be less than or equal to: [10000] but was [50000]

解决办法是: 可对某查询条件A进行排序, 然后每次查询都带上上次查询的最后值A.X, 用A条件>A.X即可, 这里只简单查个'companyId > 200'并且按照'companyId升序'的分页公司列表, 不作其他条件展开.

def buildQueryBody(pageNo, pageSize):
    # 由于ES限制, 当偏移量>10000时需要特殊处理, 此处略
    offset = (pageNo -1) * pageSize
    body = {
        "query":{
            "bool":{
                "must":[
                    {
                        "range":{
                            "companyId":{
                                "gt":"200"
                            }
                        }
                    }
                ],
                "must_not":[

                ],
                "should":[

                ]
            }
        },
        "from":offset,
        "size":pageSize,
        "sort":[
            {"companyId":{ "order": "asc" }}
        ],
        "aggs":{

        }
    }
    return body

查询结果提取

连接上ES, 创建好查询条件后, 即可执行查询:

# 创建查询条件
body = buildQueryBody(initPageNo, initPageSize)
# 执行查询
res = es.search(index=indexName, body=body)

这里的indexName为索引名, body为查询条件. 查询结果进行判断:

# seccessful == 1则成功
ifSuccess = res['_shards']['successful']

查询结果列表和总数:

resultList = res['hits']['hits']
total = res['hits']['total']

输出文件

定义一个文件输出地址, 这里采用的方案是: 如果文件存在, 则追加; 如果文件不存在, 则新建.

# 文件是否存在
ifExist = os.path.exists(file)

if (ifExist):
    # 存在文件, 在后面追加
    print ("file already exist:", file)
    csvFile = open(file,'a+')
else:
    # 新建文件
    csvFile = open(file,'wb')

将数据写入到csv文件用的是这行代码, 这里要注意的是: 不要将每个元素拼接成一个字符串"e1,e2,e3,e4", 否则会出现每个字符占据一个单元格的情况, 应该是(e1,e2,e3,e4)这样的:

writer.writerow((e1, e2, e3, e4))

如果出现编码异常的话, 可在全局定义下编码格式'utf-8':

import sys

reload(sys)
sys.setdefaultencoding('utf-8')

对文件的处理需要注意的是: 1对异常的捕获处理, 2对文件对象资源的关闭. 我这里需要打印出四列数据.

try:
    ...
    writer=csv.writer(csvFile)
    writer.writerow((companyId, name, address, telephone))
except IOError as e:
    print ("IOError happen in:", e)
except Exception as e:
    print ("Error happen in:", e)
finally:
    csvFile.close()

循环处理

为什么会做循环处理, 是因为考虑到数据量可能会很大, 进行分批处理. 循环处理的逻辑是: 不断以'下一页'的形式进行查询, 当某次查询的结果为空或小于分页值时(在成功的前提下), 认为查询结束. 听起来很容易就想到了do...while循环, 可惜Python没有这种, 于是就用'while True'来实现, 当达到'结束'的条件时, break跳出循环. 对于单次查询失败的情况, 我未做重试机制, 理论上应该是要加的.

注意: 无论是Java还是Python执行循环, 都要避免逻辑上的死循环的出现, 这里也做了一个兜底(循环1000次).

initPageNo = 1
initPageSize = 100

# 定义一个兜底循环控制, 防止出现死循环. loopMaxCount值视情况决定
loopCount = 0
loopMaxCount = 1000
while True:
    # 创建查询条件
    body = buildQueryBody(initPageNo, initPageSize)
    # 获取查询结果
    res = es.search(index=indexName, body=body)
    # seccessful = 1则成功
    ifSuccess = res['_shards']['successful']
    if (ifSuccess == 1):
        # 此次查询的数量, 用于判断是否继续查询
        resultSize = len(res['hits']['hits'])
        ...
        if (resultSize < initPageSize):
            print ("查询结束")
            break
        else:
            print ("下一页, 继续")
            initPageNo += 1
    else:
        print ("执行查询失败!", res)
        
    # 防止每次查询都失败造成的死循环    
    loopCount += 1
    if (loopCount >= loopMaxCount):
        break

以上是我实现的全过程, 下面是完整版代码, 部分关键信息已做处理, 基本上可以直接改改就用了. 如果补充上文中提到的一些可优化点, 应该是一个非常健全的用Python从Elasticsearch获取数据并输出.csv方案了.

完整代码

#-*- coding=utf-8'*-
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
import csv
import os
import sys

# 解决写入文件时中文编码问题
reload(sys)
sys.setdefaultencoding('utf-8')

# 此处替换IP地址
es = Elasticsearch(["http://**.**.**.**:9210"])
# 查看连接是否成功
print(es.info())

# 索引名
indexName = "company_info"

# 输出文件路径+文件名, 路径选择当前项目路径
fileName = "es_out.csv"
filePath = os.getcwd()
print("current file path:",filePath)
file = filePath + "/" + fileName

def printCsv(file, dataList):
    # 防止出现空数组
    if (len(dataList) == 0):
        return
    
    csvFile = ""
    writer = ""

    # 文件是否存在
    ifExist = os.path.exists(file)

    if (ifExist):
        # 文件存在, 则在后面追加
        print ("file already exist:", file)
        csvFile = open(file,'a+')
    else:
        # 文件不存在, 则新建文件
        csvFile = open(file,'wb')

    try:
        writer=csv.writer(csvFile)
        if (ifExist == False):
            # 如果是新建的表, 此处定义表的列名信息
            writer.writerow(('公司Id','公司名','地址','手机'))

        # 此处循环写入表数据
        for data in dataList:
            companyId = data['_source']['companyId'] if (data['_source'].has_key('companyId')) else 0
            name = data['_source']['name'] if (data['_source'].has_key('name')) else '--'
            address = data['_source']['address'] if (data['_source'].has_key('address')) else '--'
            telephone = data['_source']['telephone'] if (data['_source'].has_key('telephone')) else '--'

            # dataStr = "%d,%s,%s,%s" % (companyId, name, address, telephone)
            # print (dataStr)
            writer.writerow((companyId, name, address, telephone))

    except IOError as e:
        print ("IOError happen in:", e)
    except Exception as e:
        print ("Error happen in:", e)
    finally:
        csvFile.close()


def buildQueryBody(pageNo, pageSize):
    # 由于ES限制, 当偏移量>10000时需要特殊处理, 此处略
    offset = (pageNo -1) * pageSize
    body = {
        "query":{
            "bool":{
                "must":[
                    {
                        "range":{
                            "companyId":{
                                "gt":"200"
                            }
                        }
                    }
                ],
                "must_not":[

                ],
                "should":[

                ]
            }
        },
        "from":offset,
        "size":pageSize,
        "sort":[
            {"companyId":{ "order": "asc" }}
        ],
        "aggs":{

        }
    }
    return body


def main():
    print ("start...")

    initPageNo = 1
    initPageSize = 50

    # 定义一个兜底循环控制, 防止出现死循环. loopMaxCount值视情况决定
    loopCount = 0
    loopMaxCount = 1000
    while True:
        # 创建查询条件
        body = buildQueryBody(initPageNo, initPageSize)
        # 获取查询结果
        res = es.search(index=indexName, body=body)
        # seccessful == 1则成功
        ifSuccess = res['_shards']['successful']
        if (ifSuccess == 1):
            print ("success query %d times!" % initPageNo)
            resultList = res['hits']['hits']
            total = res['hits']['total']
            # 查询的数量, 用于判断是否继续查询
            resultSize = len(resultList)
            print ("total: %d, get result: %d" % (total, resultSize))

            # 打印输出结果
            printCsv(file, resultList)
            
            if (resultSize < initPageSize):
                print ("查询结束")
                break
            else:
                print ("下一页, 继续")
                initPageNo += 1
        else:
            print ("执行查询失败!", res)

        loopCount += 1
        if (loopCount >= loopMaxCount):
            break

    print ("end...")
    
if __name__ == '__main__':
    main()

觉得实用的话点个赞呗~.~