ElasticSearch 向量检索插件开发


0X00、前言

Elasticsearch是一个基于Lucene库的搜索引擎，它提供了一个分布式、支持多租户的全文搜索引擎。随着业务的飞速发展，对于搜索的需求也会增加，比如：搜索图片、相似向量等。我们可以利用 ElasticSearch 良好的插件规范、丰富的查询函数、分布式可扩展的能力开发一个脚本插件使其支持向量检索。本教程主要参考StaySense的开源项目(见参考1)。

本教程演示环境配置：

  • Python: 3.6.4
  • Java: 1.8
  • Maven
  • Docker&Elaseticsearch: 6.7.0

通过Docker部署Elasticsearch:6.7.0参考：Elasticsearch安装使用

0X01、插件开发

项目地址：https://github.com/liheyou/ImageSimilarityPlugin

项目整体目录：

图片

1、pom.xml

主要配置一些项目环境、添加依赖、打包方式等，完整配置查看项目文件。

2、plugin.xml

由于 Elasticsearch 要求自定义插件需要打包成 zip 文件，我们可以配置 Maven Assembly 插件使其自动生成，完整配置查看项目文件。

3、plugin-descriptor.properties

根据 Elasticsearch 要求，所有的插件必须包含一个名为 plugin-descriptor.properties 的插件描述文件，对其内容有要求且必须放置在 elasticsearch 目录下。我们在 src/main/resources 目录下创建 plugin-descriptor.properties 并添加内容如下：

1
2
3
4
5
6
7
8
name=${elasticsearch.plugin.name}
description=${elasticsearch.plugin.description}
version=${project.version}
classname=${elasticsearch.plugin.classname}
java.version=${maven.compiler.target}
elasticsearch.version=${elasticsearch.version
# extended.plugins=${extendedPlugins}
# has.native.controller=${hasNativeController}

除了 extended.plugins和has.native.controller 都是必须参数，其中 classname 一定是插件运行的入口文件。具体含义参考官方说明：Help for plugin authors

4、代码

通过查看官网文档，脚本插件必须继承Plugin类，通过"ScriptEngine"来实现的，为了开发一个自定义的插件，我们需要实现"ScriptEngine"接口，并通过getScriptEngine()这个方法来加载我们的插件，该插件以base64类型读取ES数据，特征向量的相似算法采用欧式距离，具体代码查看项目文件。

5、打包

运行打包命令：

1
mvn clean package

在target目录下会生成zip


图片

0X02、插件安装

通过dcoker cp 命令把zip文件复制到docker 容器中 /usr/share/elasticsearch

进入Elasticsearch:6.7.0的docker容器中

1
docker exec -it es /bin/bash

安装自定义插件

1
elasticsearch-plugin install file:///usr/share/elasticsearch/ImageSimilarity-plugin.zip

安装成功后，会在plugin目录下发现安装文件，每次安装更新插件都需要重启ES

1
docker restart es

docker 启动 elasticsearch的话，日志默认没有输出到文档，默认被终端接收，可以使用 docker logs -f es 查看。

0X03、插件使用

通过Python程序向ES集群写入100w条数据，要注意索引的mappings设置，feature是ES存储特性向量数据的字段，以base64形式存储。不同编程语言List & base64转换程序见参考2。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
import random
import base64
import numpy as np
from elasticsearch import Elasticsearch, helpers

dbig = np.dtype('>f8')

es = Elasticsearch()

body = {
    "mappings": {
        "image_search": {
            "properties": {
                "id": {
                    "type": "keyword"
                },
                "feature": {
                    "type": "binary",
                    "doc_values": True
                }
            }
        }
    }
}

index = 'test'
es.indices.delete(index=index, ignore=404)
es.indices.create(index=index, ignore=400, body=body)


def decode_float_list(base64_string):
    """
    base64 转 list
    :param base64_string:
    :return:
    """
    bytes_ = base64.b64decode(base64_string)
    return np.frombuffer(bytes_, dtype=dbig).tolist()


def encode_array(arr):
    """
    List 转 base64
    :param arr:
    :return:
    """
    base64_str = base64.b64encode(np.array(arr).astype(dbig)).decode("utf-8")
    return base64_str


def generator():
    i = 0
    while True:
        yield {

                'id': i,
                'feature': encode_array([random.random(), random.random()])
            }
        i += 1
        if i >= 1000000:
            break


# 批量插入100w数据到es

helpers.bulk(es, generator(), index=index, doc_type='image_search')

查询程序，注意source和lang要和插件里一致。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import time
import json
import base64
import numpy as np
from elasticsearch import Elasticsearch, helpers

dbig = np.dtype('>f8')

es = Elasticsearch()

body = {
    "from": 0,
    "size": 5,
    "_source": {
        "excludes": ""
    },
    "sort": {
        "_score": {
            "order": "asc"
        }
    },
    "query": {
        "function_score": {
            "query": {
                "match_all": {}
            },
            "functions": [
                {
                    "script_score": {
                        "script": {
                            "source": "DebugWorld",
                            "lang": "ImageSimilarity",
                            "params": {
                                "field": "feature",
                                "feature": [0.01, 0.03]
                            }
                        }
                    }
                }
            ]
        }
    }
}


def decode_float_list(base64_string):
    """
    base64 转 list
    :param base64_string:
    :return:
    """
    bytes_ = base64.b64decode(base64_string)
    return np.frombuffer(bytes_, dtype=dbig).tolist()


time_list = list()
for i in range(1):
    start_time = time.time()
    result = es.search(index='test', doc_type='image_search', body=body)
    for hit in result['hits']['hits']:
        hit['_source']['feature'] = decode_float_list(hit['_source']['feature'])
    time_list.append(time.time() - start_time)
    print(json.dumps(result, indent=4))

print(sum(time_list)/len(time_list), max(time_list), min(time_list))

0X04、插件性能

测试服务器配置：2核8G，20G磁盘

数据: 100w数据量，160维的特征向量

1000次请求响应情况：

平均时间：0.348s

最慢时间：1.085s

最快时间：0.215s

呃。。。百万数据量之内还是坑得住的，但是阿里云开发了一个有点吊的插件aliyun-knn，但是未开源，有兴趣的可以看看文档。

0X05、注意事项

  • Elasticsearch的score不能为负数
  • 浏览器插件Elasticsearch Head 列表页不支持展示数组类型数据
  • Elasticsearch版本更新快，不兼容情况严重，实践请注意ES版本号
  • 从7.2版本开始，Elasticsearch提供了实验性的向量检索功能

0X06、参考

  1. Elasticsearch: Elasticsearch Plugins and Integrations [6.7]
  2. Github：StaySense/fast-cosine-similarity
  3. Github: elastic/elasticsearch
  4. 公众号：蚂蚁金服 ZSearch 在向量检索上的探索