Initial attempt at efficient range queries across multiple Set objects.

Author: davidblewett

fstwrapper/src/set.rs

@@ -146,13 +146,28 @@ pub extern "C" fn fst_set_make_opbuilder(ptr: *mut Set) -> *mut set::OpBuilder<'
 }
 make_free_fn!(fst_set_opbuilder_free, *mut set::OpBuilder);
 
+#[no_mangle]
+pub extern "C" fn fst_set_make_opstreambuilder(ptr: *mut set::Stream) -> *mut set::OpBuilder<'static> {
+    let stream = ref_from_ptr!(ptr);
+    let ob = set::OpBuilder::new().add(stream);
+    to_raw_ptr(ob)
+}
+make_free_fn!(fst_set_opstreambuilder_free, *mut set::OpBuilder);
+
 #[no_mangle]
 pub extern "C" fn fst_set_opbuilder_push(ptr: *mut set::OpBuilder, set_ptr: *mut Set) {
     let set = ref_from_ptr!(set_ptr);
     let ob = mutref_from_ptr!(ptr);
     ob.push(set);
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_opbuilder_push_stream(ptr: *mut set::OpBuilder, stream_ptr: *mut set::Stream) {
+    let stream = ref_from_ptr!(stream_ptr);
+    let ob = mutref_from_ptr!(ptr);
+    ob.push(stream);
+}
+
 #[no_mangle]
 pub extern "C" fn fst_set_opbuilder_union(ptr: *mut set::OpBuilder)
                                           -> *mut set::Union {

rust_fst/_build_ffi.py

@@ -45,6 +45,7 @@
     typedef struct SetLevStream SetLevStream;
     typedef struct SetRegexStream SetRegexStream;
     typedef struct SetOpBuilder SetOpBuilder;
+    typedef struct SetOpStreamBuilder SetOpStreamBuilder;
     typedef struct SetUnion SetUnion;
     typedef struct SetIntersection SetIntersection;
     typedef struct SetDifference SetDifference;
@@ -80,6 +81,9 @@
     char* fst_set_regexstream_next(SetRegexStream*);
     void fst_set_regexstream_free(SetRegexStream*);
 
+    void fst_set_opbuilder_push_stream(SetOpStreamBuilder*, SetStream*);
+    void fst_set_opstreambuilder_free(SetOpStreamBuilder*);
+
     void fst_set_opbuilder_push(SetOpBuilder*, Set*);
     void fst_set_opbuilder_free(SetOpBuilder*);
     SetUnion* fst_set_opbuilder_union(SetOpBuilder*);

rust_fst/set.py

@@ -56,13 +56,20 @@ def get_set(self):
 
 
 class OpBuilder(object):
-    def __init__(self, set_ptr):
+    def __init__(self, ptr, as_stream=False):
         # NOTE: No need for `ffi.gc`, since the struct will be free'd
         #       once we call union/intersection/difference
-        self._ptr = lib.fst_set_make_opbuilder(set_ptr)
+        self.as_stream = as_stream
+        if self.as_stream:
+            self._ptr = lib.fst_set_make_opstreambuilder(ptr)
+        else:
+            self._ptr = lib.fst_set_make_opbuilder(ptr)
 
-    def push(self, set_ptr):
-        lib.fst_set_opbuilder_push(self._ptr, set_ptr)
+    def push(self, ptr):
+        if self.as_stream:
+            lib.fst_set_opbuilder_push_stream(self._ptr, ptr)
+        else:
+            lib.fst_set_opbuilder_push(self._ptr, ptr)
 
     def union(self):
         stream_ptr = lib.fst_set_opbuilder_union(self._ptr)
@@ -220,6 +227,28 @@ def _make_opbuilder(self, *others):
             opbuilder.push(oth._ptr)
         return opbuilder
 
+    def _make_opstreambuilder(self, *streams):
+        if not streams:
+            raise ValueError("Must pass at least one stream")
+        opbuilder = OpBuilder(streams[0]._ptr, as_stream=True)
+        for item in streams[1:]:
+            opbuilder.push(item._ptr)
+        return opbuilder
+
+    def range_union(self, str_slice, *others):
+        """ Get an iterator over a range of keys in the union of this set and others.
+
+        :param others:  List of :py:class:`Set` objects
+        :returns:       Iterator over all keys in all sets in lexicographical
+                        order
+        """
+        items = [self] + list(others)
+        streams = [
+            item[str_slice]
+            for item in items
+        ]
+        return self._make_opstreambuilder(*streams).union()
+
     def union(self, *others):
         """ Get an iterator over the keys in the union of this set and others.