使用redis+lua实现SQL中的select intersect的效果

afanti

2020-05-09

1.需求

业务中需要实现在两个集合中搜索数据，并返回交集。
用SQL的伪代码可以描述如下：

select key from set1 where sorted_key between min and max
INTERSECT
select key from set2 where sorted_key between min and max

2.现有存储格式

业务使用了redis的有序集合(sorted set)来存储数据：

set1 ->
  field1 -> value1
  field2 -> value2
set2 ->
  field1 -> value1
  field2 -> value2

常规的思路是：

# 分页拉取KEY1，然后分页拉取KEY2，然后计算交集
ZRANGEBYSCORE set1 ${min} ${max} WITHSCORES LIMIT 0 1000
ZRANGEBYSCORE set1 ${min} ${max} WITHSCORES LIMIT 1000 1000

3.lua实现

看了一下redis的lua脚本功能，可以完全在redis服务器端完成：

--like sql: insert ...select， 实现类似SQL中的insert select语句
local function select_and_insert(from_key, min_value, max_value, to_key)
    local cnt = redis.call("ZCOUNT", from_key, min_value, max_value)
    local start = 0
    local limit = 1000
    while(start<cnt)
    do
        local value = redis.call("ZRANGEBYSCORE", from_key, min_value, max_value, "WITHSCORES", "LIMIT", start, limit)
        start = start + (#value)/2
        local idx = 1
        while( idx<=#value )
        do
            value[idx], value[idx+1] = value[idx+1], value[idx]
            idx = idx + 2
        end
        local ret = redis.call("ZADD", to_key, unpack(value))
        if(ret==false)
        then
            return false, "zadd fail"
        end
    end
    return true,"success"
end

--like select intersect，实现类似SQL中的select intersect
local function intersect(key1, min1, max1, key2, min2, max2, to_key)
    local temp_key_1 = "___temp_1"
    redis.call("DEL", temp_key_1)
    local ret,msg = select_and_insert(key1, min1, max1, temp_key_1)
    if(ret==false)
    then
        return false, key1.." fail:"..msg
    end
    --
    local temp_key_2 = "___temp_2"
    redis.call("DEL", temp_key_2)
    ret,msg = select_and_insert(key2, min2, max2, temp_key_2)
    if(ret==false)
    then
        redis.call("DEL", temp_key_1)
        return false, key2.." fail:"..msg
    end 
    --
    ret = redis.call("ZINTERSTORE", to_key, 2, temp_key_1, temp_key_2)
    redis.call("DEL", temp_key_1)
    redis.call("DEL", temp_key_2)
    if(ret==false)
    then
        return false, "ZINTERSTORE fail"
    end
    return true, string.format("intersect count is %d", ret)
end

local function main()
    return intersect("set1", 10, 100, "set2", 200, 400, "my_result")
end

return main()

在命令行测试一下：
redis-cli -h 192.168.0.2 -p 6379 -a "my_password" --eval test_intersect.lua
达到了预期的效果。

4.做一个script load工具

按照网上的方法这样导入脚本，始终不成功：

redis-cli -h 192.168.0.2 -p 6379 -a "my_password" SCRIPT LOAD "`cat test_intersect.lua`"

于是用golang基于gin框架来做一个脚本保存功能：

上传表单：

//注册
my_gin.GET("/redis_script_form", redisScriptForm)

func redisScriptForm(c *gin.Context){
    c.Data(200, "text/html; charset=utf-8", []byte(`
        <html>
        <body>
        <form method="POST" enctype="application/x-www-form-urlencoded" action="/redis_script_load">
        <textarea name="lua" style="width:100%; height:300px"></textarea><br/>
        <input type="submit" value="upload script"/>
        </form>
        </body>
        </html>
    `))
}

保存接口

//
import "github.com/go-redis/redis/v7"
//
my_gin.POST("/redis_script_load", redisScriptLoad)
func redisScriptLoad(c *gin.Context){
    code := c.PostForm("lua")
    redis := utils.GetRedisClient()
    val, err := redis.ScriptLoad(code).Result()
    if err!=nil{
        c.Data(200, "text/plain", []byte(fmt.Sprintf("redis.ScriptLoad fail:%s", err.Error() )))
        return
    }
    c.Data(200, "text/plain", []byte(fmt.Sprintf("redis.ScriptLoad success:sha=%s", val )))
}

保存成功会返回代码的SHA hash值。

5.用golang调用redis中的lua脚本

lua的main()改一改

首先要再修改之前lua代码中的main()，不要写死参数：

local function main()
    local set1_param={KEYS[1], ARGV[1], ARGV[2]}
    local set2_param={KEYS[2], ARGV[3], ARGV[4]}
    local temp_key = "__temp_3"
    --
    redis.call("DEL", temp_key)
    local ret,msg = intersect(set1_param[1], set1_param[2], set1_param[3], set2_param[1], cpu_param[2], set2_param[3], temp_key)
    if(ret==false)
    then
        return msg
    end

    local values = redis.call("ZRANGE", temp_key, "0", "100000000000")
    redis.call("DEL", temp_key)
    return values
end

命令行测试

再次保存代码，得到SHA值。
用命令行进行测试：
EVALSHA f2f7d1b5439b8bb4c8320a7dce4b54c133a3d47d 2 "set1" "set2" "10" "1000" "100" "3000"
得到了预想的结果。

golang代码测试：

//注册
my_gin.GET("/redis_eval_script", redisEvalScript)

func redisEvalScript(c *gin.Context){
    redis := utils.GetRedisClient()
    val, err := redis.Do("EVALSHA",
        "f2f7d1b5439b8bb4c8320a7dce4b54c133a3d47d",
        2, "set1", "set2", 10, 1000,
        100, 3000).Result()
    if err!=nil{
        ResponseError(c, err.Error())
        return
    }
    j,err := json.Marshal(val)
    if err!=nil{
        c.Data(200, "text/plain", []byte(fmt.Sprintf("EVALSHA:%s", err.Error() )))
        return
    }
    c.Data(200, "text/plain; charset=utf-8",
        []byte(fmt.Sprintf("json:\n%s", string(j))))
}

have fun! ??

select redis