Fast Index Tarantool Iterator

fast_iterator.lua

local ffi = require 'ffi'
local fun = require 'fun'
local C = ffi.C


local function generator(iterator, rtuple)
	local res = C.box_iterator_next(iterator, rtuple)

	if res == -1 then
		box.error()
	end

	if rtuple[0] == nil then
		return nil
	end

	-- no refs/unrefs => can't forward tuple anywhere
	return rtuple, ffi.cast('const struct tuple &', rtuple[0])
end


return function(index, iter_type, tuple)
	assert(index, "Usage: fast_iterator(box.space.myspace.index.myindex, box.index.EQ, tuple):each(...)")

	local space_id = assert(index.space_id, "no space_id")
	if type(iter_type) == 'string' then
		iter_type = box.index[iter_type]
	end

	assert(type(iter_type) == 'number', "Invalid iter_type")

	if not tuple then
		tuple = box.tuple.new()
	end

	local iterator = C.box_index_iterator(space_id, index.id, iter_type, box.tuple.encode(tuple))

	if iterator == nil then
		return box.error() -- throws last exception
	end

	return fun.iter(generator, ffi.gc(iterator, C.box_iterator_free), ffi.new('box_tuple_t *[1]'))
end

fast_iterator_3.0.lua

local ffi = require 'ffi'
local fun = require 'fun'
local builtin = ffi.C
local buffer = require 'buffer'
local cord_ibuf_take = buffer.internal.cord_ibuf_take
local cord_ibuf_put = buffer.internal.cord_ibuf_put

local tuple_encode = box.internal.tuple.encode
local check_pairs_opts = box.internal.check_pairs_opts
local check_index_arg = box.internal.check_index_arg

-- pointer to iterator position used by select(), pairs() and tuple_pos()
local iterator_pos = ffi.new('const char *[1]')
local iterator_pos_end = ffi.new('const char *[1]')
local iterator_t = ffi.typeof('struct iterator')

-- a static box_tuple_t ** instance for calling box_index_* API
local ptuple = ffi.new('box_tuple_t *[1]')
local const_tuple_ref_t = ffi.typeof('box_tuple_t&')

--
-- Sets iterator_pos and iterator_pos_end to a user-supplied position.
--
-- The input position may be nil, string, table, or tuple. If the input
-- position is given as string, iterator_pos is set to point to its data,
-- otherwise the iterator_pos data is allocated from the fiber region.
--
-- The ibuf is used to encode a position given as table or tuple.
--
-- Returns true on success. On failure, sets box.error and returns false.
--
local function iterator_pos_set(index, pos, ibuf)
    if pos == nil then
        iterator_pos[0] = nil
        iterator_pos_end[0] = nil
        return true
    elseif type(pos) == 'string' then
        iterator_pos[0] = pos
        iterator_pos_end[0] = iterator_pos[0] + #pos
        return true
    else
        ibuf:consume(ibuf.wpos - ibuf.rpos)
        local tuple, tuple_end = tuple_encode(ibuf, pos)
        return builtin.box_index_tuple_position(
                index.space_id, index.id, tuple, tuple_end,
                iterator_pos, iterator_pos_end) == 0
    end
end

local function iterator_gen(param, state)
    if not ffi.istype(iterator_t, state) then
        error('usage: next(param, state)')
    end
    -- next() modifies state in-place
    if builtin.box_iterator_next(state, ptuple) ~= 0 then
        return box.error() -- error
    elseif ptuple[0] ~= nil then
		return state, ffi.cast(const_tuple_ref_t, ptuple[0])
    else
        return nil
    end
end


function _G.iterator(index, key, opts)
	check_index_arg(index, 'pairs')
	local ibuf = cord_ibuf_take()
	local pkey, pkey_end = tuple_encode(ibuf, key)
	local svp = builtin.box_region_used()
	local itype, after = check_pairs_opts(opts, pkey + 1 >= pkey_end)
	local ok = iterator_pos_set(index, after, ibuf)
	local keybuf = ffi.string(pkey, pkey_end - pkey)
	cord_ibuf_put(ibuf)
	local cdata
	if ok then
		local pkeybuf = ffi.cast('const char *', keybuf)
		cdata = builtin.box_index_iterator_after(
				index.space_id, index.id, itype, pkeybuf, pkeybuf + #keybuf,
				iterator_pos[0], iterator_pos_end[0])
	end
	builtin.box_region_truncate(svp)
	if cdata == nil then
		box.error()
	end
	return fun.wrap(iterator_gen, keybuf,
		ffi.gc(cdata, builtin.box_iterator_free))
end

tarantool> _TARANTOOL
---
- 3.0.0-0-g6ba34da
...

tarantool> jit.status()
---
- true
- fold
- cse
- dce
- fwd
- dse
- narrow
- loop
- abc
- sink
- fuse
...

tarantool> fiber.extend_slice({warn=30, err=120}) return clock.bench(function() return box.space.number:pairs():length() end)
---
- - 6.581532
  - 20000000
...

tarantool> fiber.extend_slice({warn=30, err=120}) return clock.bench(function() return iterator(box.space.number.index.primary,{},{}):length() end)
---
- - 0.404926
  - 20000000
...