diff --git a/rocksdb/_rocksdb.pyx b/rocksdb/_rocksdb.pyx index 04e7fd3..134cdb3 100644 --- a/rocksdb/_rocksdb.pyx +++ b/rocksdb/_rocksdb.pyx @@ -500,6 +500,7 @@ cdef class Options(object): cdef PyFilterPolicy py_filter_policy cdef PyCache py_block_cache cdef PyCache py_block_cache_compressed + cdef PySliceTransform py_prefix_extractor def __cinit__(self): self.opts = new options.Options() @@ -513,6 +514,7 @@ cdef class Options(object): self.py_filter_policy = None self.py_block_cache = None self.py_block_cache_compressed = None + self.py_prefix_extractor = None for key, value in kwargs.items(): setattr(self, key, value) @@ -955,6 +957,16 @@ cdef class Options(object): self.opts.filter_policy = self.py_filter_policy.get_policy() + property prefix_extractor: + def __get__(self): + if self.py_prefix_extractor is None: + return None + return self.py_prefix_extractor.get_ob() + + def __set__(self, value): + self.py_prefix_extractor = PySliceTransform(value) + self.opts.prefix_extractor = self.py_prefix_extractor.get_transformer() + property block_cache: def __get__(self): if self.py_block_cache is None: @@ -1188,8 +1200,12 @@ cdef class DB(object): def iterkeys(self, prefix=None, *args, **kwargs): cdef options.ReadOptions opts cdef KeysIterator it + opts = self.build_read_opts(self.__parse_read_opts(*args, **kwargs)) + it = KeysIterator(self) + it.set_prefix(opts, prefix) + with nogil: it.ptr = self.db.NewIterator(opts) return it @@ -1197,8 +1213,12 @@ cdef class DB(object): def itervalues(self, prefix=None, *args, **kwargs): cdef options.ReadOptions opts cdef ValuesIterator it + opts = self.build_read_opts(self.__parse_read_opts(*args, **kwargs)) + it = ValuesIterator(self) + it.set_prefix(opts, prefix) + with nogil: it.ptr = self.db.NewIterator(opts) return it @@ -1206,8 +1226,12 @@ cdef class DB(object): def iteritems(self, prefix=None, *args, **kwargs): cdef options.ReadOptions opts cdef ItemsIterator it + opts = self.build_read_opts(self.__parse_read_opts(*args, **kwargs)) + it = ItemsIterator(self) + it.set_prefix(opts, prefix) + with nogil: it.ptr = self.db.NewIterator(opts) return it @@ -1300,6 +1324,9 @@ cdef class Snapshot(object): cdef class BaseIterator(object): cdef iterator.Iterator* ptr cdef DB db + # To keep a reference to the prefix + cdef object prefix + cdef Slice c_prefix def __cinit__(self, DB db): self.db = db @@ -1325,6 +1352,14 @@ cdef class BaseIterator(object): def __reversed__(self): return ReversedIterator(self) + cdef set_prefix(self, options.ReadOptions& opts, object prefix=None): + if prefix is None: + return + + self.c_prefix = bytes_to_slice(prefix) + self.prefix = prefix + opts.prefix = cython.address(self.c_prefix) + cpdef seek_to_first(self): with nogil: self.ptr.SeekToFirst() diff --git a/rocksdb/options.pxd b/rocksdb/options.pxd index 4bc30f8..73294dd 100644 --- a/rocksdb/options.pxd +++ b/rocksdb/options.pxd @@ -10,6 +10,7 @@ from cache cimport Cache from logger cimport Logger from slice_ cimport Slice from snapshot cimport Snapshot +from slice_transform cimport SliceTransform cdef extern from "rocksdb/options.h" namespace "rocksdb": ctypedef enum CompressionType: @@ -44,7 +45,7 @@ cdef extern from "rocksdb/options.h" namespace "rocksdb": CompressionType compression # TODO: compression_per_level # TODO: compression_opts - # TODO: prefix_extractor + SliceTransform* prefix_extractor cpp_bool whole_key_filtering int num_levels int level0_file_num_compaction_trigger diff --git a/rocksdb/tests/test_db.py b/rocksdb/tests/test_db.py index e00868a..8f5508c 100644 --- a/rocksdb/tests/test_db.py +++ b/rocksdb/tests/test_db.py @@ -274,3 +274,42 @@ class TestComparator(unittest.TestCase, TestHelper): self.db.put(int_to_bytes(x), int_to_bytes(x)) self.assertEqual(b'300', self.db.get(b'300')) + +class StaticPrefix(rocksdb.interfaces.SliceTransform): + def name(self): + return b'static' + + def transform(self, src): + return (0, 5) + + def in_domain(self, src): + return len(src) >= 5 + + def in_range(self, dst): + return len(dst) == 5 + +class TestPrefixExtractor(unittest.TestCase, TestHelper): + def setUp(self): + opts = rocksdb.Options(create_if_missing=True) + opts.prefix_extractor = StaticPrefix() + self._clean() + self.db = rocksdb.DB('/tmp/test', opts) + + def test_prefix(self): + for x in range(3000): + keyx = b'%s.x' % hex(x)[2:].zfill(5).encode('utf8') + keyy = b'%s.y' % hex(x)[2:].zfill(5).encode('utf8') + keyz = b'%s.z' % hex(x)[2:].zfill(5).encode('utf8') + self.db.put(keyx, b'x') + self.db.put(keyy, b'y') + self.db.put(keyz, b'z') + + self.assertEqual('x', self.db.get(b'00001.x')) + self.assertEqual('y', self.db.get(b'00001.y')) + self.assertEqual('z', self.db.get(b'00001.z')) + + it = self.db.iterkeys(prefix=b'00002') + it.seek(b'00002') + + ref = ['00002.x', '00002.y', '00002.z'] + self.assertEqual(ref, list(it))