【代碼分享】用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! 😃