Make 'prefix_extractor' active

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()

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

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))