Redis Quadtrees in Hash - finally motivated because of http://stackoverflow.com/questions/32911604/intersection-of-two-or-more-sorted-sets

formal-benchmark.txt

foo@bar:~/src/quadtree$ redis-cli zcard x
(integer) 100000
foo@bar:~/src/quadtree$ redis-cli zcard y
(integer) 100000
foo@bar:~/src/quadtree$ redis-cli hlen qt
(integer) 56325
foo@bar:~/src/quadtree$ redis-cli script load "`cat rqtih.lua`"
"84f91fddf261e9a2ef14ff4b417443eed3dc4b2b"
foo@bar:~/src/quadtree$ redis-benchmark -n 100 -r 1000000 evalsha 84f91fddf261e9a2ef14ff4b417443eed3dc4b2b 3 x y tmp  zquery __rand_int__ __rand_int__ __rand_int__ __rand_int__
====== evalsha 84f91fddf261e9a2ef14ff4b417443eed3dc4b2b 3 x y tmp zquery __rand_int__ __rand_int__ __rand_int__ __rand_int__ ======
  100 requests completed in 41.50 seconds
  50 parallel clients
  3 bytes payload
  keep alive: 1

1.00% <= 404 milliseconds
2.00% <= 20128 milliseconds
3.00% <= 20523 milliseconds
8.00% <= 20524 milliseconds
52.00% <= 20966 milliseconds
100.00% <= 20966 milliseconds
2.41 requests per second

foo@bar:~/src/quadtree$ redis-benchmark -n 100 -r 1000000 evalsha 84f91fddf261e9a2ef14ff4b417443eed3dc4b2b 1 qt query __rand_int__ __rand_int__ __rand_int__ __rand_int__
====== evalsha 84f91fddf261e9a2ef14ff4b417443eed3dc4b2b 1 qt query __rand_int__ __rand_int__ __rand_int__ __rand_int__ ======
  100 requests completed in 1.36 seconds
  50 parallel clients
  3 bytes payload
  keep alive: 1

1.00% <= 37 milliseconds
2.00% <= 638 milliseconds
3.00% <= 665 milliseconds
9.00% <= 666 milliseconds
24.00% <= 667 milliseconds
43.00% <= 668 milliseconds
52.00% <= 679 milliseconds
59.00% <= 680 milliseconds
100.00% <= 680 milliseconds
73.69 requests per second

rqtih.lua

-- rqtih.lua - Redis Quadtree in Hash
--
-- TODO:
--  Extend to store boxes and polygons instead of just points
--  More dimensions - kd-trees
--  Get rid of object-oriented design - flatten everything to nested loops
--  JSON, really? msgpck or even better - a custom bit-level encoder (otoh, readable and implemented in C...)
--  node capacity should be configurable as well
--  Currently requires declaring the qt's min/max dimensions a priori
--  Use HMGET and HMSET
--  Missing QT functionality: removePoint, exists, rebuild, ...
--  Add member uniqueness and then support union, intersect & diff on QTs
--  ...
--  Do the entire thing in Redis core and add this as an internal data type - win!
--
--  Copyright Itamar Haber & Redis Labs. Provided under the 3-Clause BSD licene.

-- Helper functions
-- add a call to _trace in strategic code paths, open a MONITOR session to view
-- output interleaved with the actual calls to redis.call + timing
local function _trace(...)
  local msg = "_trace:"
  for i, v in ipairs(arg) do
    msg = msg .. ' [' .. i .. ']: ' .. v
  end
  return redis.call('ECHO', msg)
end

-- returns true if two boxes intersect
local function _intersect(ax1,ay1,aw,ah, bx1,by1,bw,bh)
  
  local ax2,ay2,bx2,by2 = ax1 + aw, ay1 + ah, bx1 + bw, by1 + bh
  return not (bx1 > ax2 or bx2 < ax1 or by1 > ay2 or by2 < ay1)
end

-- returns true if a is contained in b
local function _contained(ax1,ay1,aw,ah, bx1,by1,bw,bh)
  
  local ax2,ay2,bx2,by2 = ax1 + aw, ay1 + ah, bx1 + bw, by1 + bh
  return bx1 <= ax1 and bx2 >= ax2 and by1 <= ay1 and by2 >= ay2
end

-- Point class
-- An arbitrary point in 2D space
local Point = {}
Point.__index = Point
Point.__tostring = function(cls)
    return '{ "Point" : { "x" : "' .. (cls.x or "nil") .. '", "y" : "' .. (cls.y or "nil") .. '", "data" :"' .. (cls.data or nil) .. '" } }'
end

setmetatable(Point, {
    __call = function(cls, ...)
      return cls.new(...)
    end,
  })    

function Point.new(x, y, data)
  local self = setmetatable({}, Point)

  self.x, self.y, self.data = x, y, data
  return self
end

-- Quadtree class
local Quadtree = {}
Quadtree.__index = Quadtree
Quadtree.__tostring = function(cls)
  local reply =
  '{ "Quadtree" :' ..
  ' { "x" = "' .. (cls.x or "nil") .. 
  '", "y" = "' .. (cls.y or "nil") ..
  '", "w" = "' .. (cls.w or "nil") ..
  '", "h" = "' .. (cls.h or "nil") ..
  '", "points" = [ '

  local spoints
  for _, v in pairs(cls.points) do
    if spoints then
      spoints = spoints .. ", " .. tostring(v)
    else
      spoints = tostring(v)
    end
  end  

  reply = reply .. (spoints or "") .. " ] } }"

  return reply
end

setmetatable(Quadtree, {
    __call = function (cls, ...)
      return cls.new(...)
    end,
  })

-- Initializes the quadtree. Call with dimension arguments to create a new one
-- or without them to load from storage
function Quadtree.new(x, y, w, h, name)
  local self = setmetatable({}, Quadtree)
  self.name = name or (KEYS[1])
  self.children = {}
  self.points = {}

  -- load the node from storage if it exists, otherwise init
  local storage = redis.call('HGET', KEYS[1], self.name)
  if storage then
    self.exists = true
    storage = cjson.decode(storage)
    self.x = storage.x
    self.y = storage.y
    self.w = storage.w
    self.h = storage.h
    if storage.points then
      for _, v in pairs(storage.points) do
        self.points[#(self.points)+1] = Point.new(v.x, v.y, v.data)
      end
    end
  else
    self.exists = false
    self.x, self.y, self.w, self.h = x or 0, y or 0, w or 0, h or 0
  end

  self.loadedChildren = false
  self.dirty = false
  return(self)
end

-- loads a node's children from storage
function Quadtree:loadChildren()
  local children = redis.call('HEXISTS', KEYS[1], self.name .. "/NW")
  if children ~= 0 then
    self.children[#(self.children)+1] = Quadtree.new(0, 0, 0, 0, self.name .. "/NW")
    self.children[#(self.children)+1] = Quadtree.new(0, 0, 0, 0, self.name .. "/NE")
    self.children[#(self.children)+1] = Quadtree.new(0, 0, 0, 0, self.name .. "/SW")
    self.children[#(self.children)+1] = Quadtree.new(0, 0, 0, 0, self.name .. "/SE")
  end
  self.loadedChildren = true
end

-- inserts a point to a quadtree node
function Quadtree:insertPoint(point)
  if not _contained(point.x, point.y, 0, 0, self.x, self.y, self.w, self.h) then
    return false
  end

  -- try to insert the point to this node
  if #(self.points) < 4 then -- node capacity
    self.points[#(self.points)+1] = point
    self.dirty = true
    return true
  end

  if not self.loadedChildren then
    self:loadChildren()
  end

  -- subdivide the node to 4 children
  if #(self.children) == 0 then
    local hw = self.w / 2.0
    local hh = self.h / 2.0
    self.children[#(self.children)+1] = Quadtree.new(self.x, self.y, hw, hh, self.name .. "/NW")
    self.children[#(self.children)+1] = Quadtree.new(self.x + hw, self.y, hw, hh, self.name .. "/NE")
    self.children[#(self.children)+1] = Quadtree.new(self.x, self.y + hh, hw, hh, self.name .. "/SW")
    self.children[#(self.children)+1] = Quadtree.new(self.x + hw, self.y + hh , hw, hh, self.name .. "/SE")
  end

  -- insert to relevant child
  for _, child in pairs(self.children) do
    if child:insertPoint(point) then
      return true
    end
  end

  return false
end

-- returns the points in the area specified
function Quadtree:query(x, y, w, h)
  local reply = {}

  if not _intersect(x, y, w, h, self.x, self.y, self.w, self.h) then
    return reply
  end

  for _, point in pairs(self.points) do
    if _contained(point.x, point.y, 0, 0, x, y, w, h) then
      reply[#reply + 1] = point
    end
  end
    
  if not self.loadedChildren then
    self:loadChildren()
  end

  if #(self.children) ~= 0 then
    for _, child in pairs(self.children) do
      local r = child:query(x, y, w, h)
      for _, p in pairs(r) do
        reply[#reply + 1] = p
      end
    end
  end

  return reply
end

-- converts a quadtree to a table
function Quadtree:toTable(dirtyOnly)
  local reply = {}
  
  if not dirtyOnly or dirtyOnly and self.dirty then
    reply[self.name] = cjson.encode({ x = self.x, y = self.y, h = self.h, w = self.w, points = self.points})
  end
  
  for _, child in pairs(self.children) do
    local cs = child:toTable(dirtyOnly)
    for k, v in pairs(cs) do
      reply[k] = v
    end
  end

  return reply
end

-- saves the quadtree to storage
function Quadtree:dump(dirtyOnly)
  local t = self:toTable(dirtyOnly)
  for k, v in pairs(t) do
    redis.call('HSET', KEYS[1], k, v)
  end
end

local function _debug_populatetree()
  redis.call('DEL', KEYS[1])
  local q = Quadtree.new(0, 0, 100, 100)
  q:insertPoint(Point.new(1,1, 'first'))
  q:insertPoint(Point.new(2,2, '2nd'))
  q:insertPoint(Point.new(4,4, 'last'))
  q:dump()
end

local function _debug_populatezsets(size, xmax, ymax)
  redis.call('DEL', KEYS[1], KEYS[2])
  
  for i = 1, size do
    redis.call('ZADD', KEYS[1], math.random(xmax), 'id:' .. i)
    redis.call('ZADD', KEYS[2], math.random(ymax), 'id:' .. i)
  end
end

local function fuzzy_test(items, queries)  
  local dim = 2
  local timings = {}
  local avgt = 0.0
  
  redis.call('DEL', KEYS[1])
  local q = Quadtree.new(1, 1, 1000, 1000)
  
  local id = 0
  local dataset = {}
  for i = 1, items do
    local vars = {}
    for j = 1, dim do
      table.insert(vars, math.random(1000))
    end
    q:insertPoint(Point.new(vars[1], vars[2], id))
    table.insert(vars, id)
    table.insert(dataset, vars)
    id = id + 1
  end
  q:dump()
  
  for i = 1, queries do
    local random = {}
    for j = 1, dim do
      local s = math.random(1000)
      local e = math.random(1000)
      if e < s then
        s, e = e, s
      end
      table.insert(random, { s, e })
    end
    
    local tq = Quadtree.new(1, 1, 1000, 1000)
    local start_t = redis.call('TIME')
    local res1 = tq:query(random[1][1], random[2][1], random[1][2] - random[1][1], random[2][2] - random[2][1])
    local end_t = redis.call('TIME')
    
    for i, v in ipairs(res1) do
      local t = {}
      res1[i] = {v.x, v.y, v.data}
    end
    
    start_t[1], start_t[2] = tonumber(start_t[1]), tonumber(start_t[2])
    end_t[1], end_t[2] = tonumber(end_t[1]), tonumber(end_t[2])
    if end_t[2] > start_t[2] then
      table.insert(timings, { end_t[1] - start_t[1], end_t[2] - start_t[2] })
    else
      table.insert(timings, { end_t[1] - start_t[1] - 1, math.abs(end_t[2] - start_t[2]) })
    end
    
    avgt = (avgt * (#timings - 1) + tonumber(string.format('%d.%06d', timings[#timings][1], timings[#timings][2]))) / #timings
    
    local res2 = {}
    for _, v in pairs(dataset) do
      local included = true
      for j = 1, dim do
        if v[j] < random[j][1] or v[j] > random[j][2] then
          included = false
        end
      end
      if included then
        table.insert(res2, v)
      end
    end
        
    -- table sorting is so much FUN!
    local function cmp(a, b, depth)
      depth = depth or 1
      
      if depth > dim + 1 then
        return false
      end
      if a[depth] < b[depth] then
        return true
      end
      if a[depth] == b[depth] then
        return cmp(a, b, depth + 1)
      end
      if a[depth] > b[depth] then
        return false
      end
    end
    
    table.sort(res1, cmp)
    table.sort(res2, cmp)
    
    for i1, r1 in ipairs(res1) do
      for i2 = 1, dim + 1 do
        if r1[i2] ~= res2[i1][i2] then
          error('ERROR ' .. r1[i2] .. ' ~= ' .. res2[i1][i2])
        end
      end
    end
  end

  -- housekeeping drop() can't be called unless replication mode is changed
  return({ 'fuzzily tested.', {dim, items, queries, timings}})
end


-- Example #1
--_debug_populatetree()
--local q = Quadtree.new()
--q:insertPoint(Point.new(15,5, 'for reals'))
--q:dump(true)
--return cjson.encode(q:query(3, 3, 20, 20))

-- Example 2
--_debug_populatezsets(100000, 1000000, 42000000)

local _USAGE = {
  'commands',
--  '  add     - adds a point, ARGV: x y data',
--  '  create  - creates a quatree, ARGV: x y w h',
  '  query   - returns points in range, KEYS: quadtree ARGV: x y w h',
  '  zcreate - creates a quadtree from the intersect of two sorted sets, KEYS: quadtree xzset yzset tempkeyname',
  '  zquery  - returns points in range w/o quadtree, KEYS: xzset yzset tempkeyname ARGV: x y w h',
  '  help    - this message',
  }

--- "Parser"
if #ARGV == 0 or #KEYS == 0 then
  return(_USAGE)
end

local cmd = ARGV[1]:lower()

if cmd == "help" then
  return _USAGE
elseif cmd == "query" then
  local q = Quadtree.new()
  return cjson.encode(q:query(tonumber(ARGV[2]), tonumber(ARGV[3]), tonumber(ARGV[4]), tonumber(ARGV[5])))
elseif cmd == "zcreate" then
  redis.call('ZINTERSTORE', KEYS[4], 2, KEYS[2], KEYS[3], 'WEIGHTS', 1, 0)
  local xmin = tonumber(redis.call('ZRANGE', KEYS[4], 0, 0, 'WITHSCORES')[2])
  local xmax = tonumber(redis.call('ZRANGE', KEYS[4], -1, -1, 'WITHSCORES')[2])
  redis.call('ZINTERSTORE', KEYS[4], 2, KEYS[2], KEYS[3], 'WEIGHTS', 0, 1)
  local ymin = tonumber(redis.call('ZRANGE', KEYS[4], 0, 0, 'WITHSCORES')[2])
  local ymax = tonumber(redis.call('ZRANGE', KEYS[4], -1, -1, 'WITHSCORES')[2])
  
  redis.call('DEL', KEYS[1])
  local q = Quadtree.new(xmin, ymin, xmax-xmin, ymax-ymin)
  local members = redis.call('ZRANGE', KEYS[4], 0, -1, 'WITHSCORES')
  for i = 1, #members/2 do
    local x = redis.call('ZSCORE', KEYS[2], members[i*2-1])
    q:insertPoint(Point.new(tonumber(x), tonumber(members[i*2]), members[i*2-1]))
  end
  
  redis.call('DEL', KEYS[4])
  q:dump()
elseif cmd == "zquery" then
  redis.call('ZUNIONSTORE', KEYS[3], 1, KEYS[1], 'WEIGHTS', 1)
  redis.call('ZREMRANGEBYSCORE', KEYS[3], '-inf', '(' .. ARGV[2])
  redis.call('ZREMRANGEBYSCORE', KEYS[3], '(' .. ARGV[2] + ARGV[4], '+inf')
  redis.call('ZINTERSTORE', KEYS[3], 2, KEYS[3], KEYS[2], 'WEIGHTS', 0, 1)
  return cjson.encode(redis.call('ZRANGEBYSCORE', KEYS[3], ARGV[3], ARGV[3] + ARGV[5]))
elseif cmd == "fuzzy_test" then
  return fuzzy_test(tonumber(ARGV[2]), tonumber(ARGV[3]))
else
  return "-ERR: Unknown command '" .. ARGV[1] .. "' - run with 'help' for usage"
end

test.sh

#!/bin/sh

SHA1=$(redis-cli SCRIPT LOAD "$(cat rqtih.lua)")

redis-cli EVALSHA $SHA1 2 h fuzzy_test