【代码分享】用redis+lua实现多个集合取交集并过滤，类似于: select key from set2 where key in (select key from set1) and value>=xxx

redis中的zset结构可以看成一个个包含数值的集合，或者认为是一个关系数据库中用列存储方式存储的一列。

需求

假设我有这样一个数据筛选需求，用SQL表示为：

select key
from set3
where value>${v3} and
  key in (
    select key
    from set2
    where value>${v2} and
      key in (
        $key1, $key2, $key3 ...
      )
  )

总结起来就是：

• 输入：
  • key的列表：key1, key2, key3... 任意多个
  • 每个集合及其需要筛选的下限： set2中值大于v2的key, set3中值大于v3的key...等等多个集合
• 计算过程：
  • 取所有集合的交集，并在每个集合上用下限值进行过滤
• 输出：筛选后剩下的集合

redis lua代码

下面是实现这一目的的lua代码：

-- User: ahfuzhang
-- Date: 2020/5/22
-- Time: 16:20

--给定一个集合from_key
--与集合join_key取交集
--然后过滤掉小于need的数据
--然后作为一个全新的集合存入to_key
local function join_and_filter(from_key, join_key, need, to_key)
    local temp_key = "__temp_join_and_filter"
    redis.call("DEL", temp_key)
    redis.call("ZINTERSTORE", temp_key, 2, from_key, join_key)
    redis.call("ZREMRANGEBYSCORE", temp_key, "-inf", "("..need)
    local values = redis.call("ZRANGEBYSCORE", temp_key, "-inf", "+inf", "WITHSCORES")
    if (#values==0)
    then
        return false
    end
    for idx=1,#values,2 do
        values[idx], values[idx+1] = 0, values[idx]
    end
    redis.call("DEL", to_key)
    redis.call("ZADD", to_key, "NX", unpack(values))
    return true
end

local function main()
    local argc = tonumber(ARGV[1])  --代表输入列表的数量
    local keys_list = {}
    for i=2,argc+2,1 do
        table.insert(keys_list, 0)
        table.insert(keys_list, ARGV[i])
    end
    local temp_key_list = "__temp_key_list"
    local temp_key = "__temp_middle_result"
    redis.call("DEL", temp_key_list)
    redis.call("ZADD", temp_key_list, "NX", unpack(keys_list))
    --
    local filter_count = tonumber(ARGV[argc+2]) --每个集合上的过滤条件的数量

    local from_key = temp_key_list
    for filter_idx=1, filter_count, 1 do
        local ret = join_and_filter(from_key, KEYS[filter_idx], ARGV[filter_idx+argc+2], temp_key)
        if (ret==false)
        then
            redis.call("DEL", temp_key)
            return {1, "no data after key "..KEYS[filter_idx], {}}
        end
        from_key = temp_key
    end
    --
    local values = redis.call("ZRANGEBYSCORE", temp_key, "-inf", "+inf")
    redis.call("DEL", temp_key)
    return {0, "success", values}
end

return main()

调用命令行

redis-cli -h 192.168.0.5 -p 6379 -a test123 \
  --eval redis_script_join_and_filter.lua \
  "set1" "set2" "set3" \   #这里是要逐个过滤的几个集合，其实就是redis里面zset结构的key
   , \   # 这个逗号非常重要，曾经在这里采坑，这是一个分隔符，前面是KEYS，后面是ARGV。 注意，内容是： 空格 逗号 空格，不能和前后连起来
  "4" "user1" "user2" "user3" "user4" \  #本行的第一个字段4代表了有四个输入的二级KEY，是最初要过滤的二级KEY
  "3" "393" "20" "800"   #本行的第一个字段3代表了有三个过滤值。这里的数量必须和前面KEYS的数量一致。分别对应了每个KEY下面的过滤最小值

P.S 如果要调试lua脚本，可以酱紫：

redis-cli -h 192.168.0.5 -p 6379 -a test123 --ldb --eval redis_script_join_and_filter.lua xxxx

have fun!