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/DebugWorld-1024/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 并添加内容如下:

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、打包

运行打包命令:

mvn clean package

在target目录下会生成zip


图片

0X02、插件安装

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

进入Elasticsearch:6.7.0的docker容器中

docker exec -it es /bin/bash

安装自定义插件

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

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

docker restart es

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

0X03、插件使用

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

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要和插件里一致。

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 在向量检索上的探索

你可能感兴趣的:(ElasticSearch 向量检索插件开发)