tantiny-in-memory 1.0.0

data/sig/tantiny/query.rbs

@@ -0,0 +1,135 @@
+module Tantiny
+  class Query
+    TYPES: Array[Symbol]
+    DEFAULT_BOOST: Float
+    DEFAULT_FUZZY_DISTANCE: Integer
+
+    def self.disjunction: (*Query queries) -> Query
+    def self.conjunction: (*Query queries) -> Query
+
+    def self.all_query: (?Index _index) -> Query
+    def self.empty_query: (?Index _index) -> Query
+
+    def self.term_query: (
+      Index index,
+      fields fields,
+      String term,
+      **untyped options
+    ) -> Query
+
+    def self.fuzzy_term_query: (
+      Index index,
+      fields fields,
+      String term,
+      ?Integer distance,
+      **untyped options
+    ) -> Query
+
+    def self.phrase_query: (
+      Index index,
+      fields fields,
+      String phrase,
+      **untyped options
+    ) -> Query
+
+    def self.regex_query: (
+      Index index,
+      fields fields,
+      String regex,
+      **untyped options
+    ) -> Query
+
+    def self.prefix_query: (
+      Index index,
+      fields fields,
+      String prefix,
+      **untyped options
+    ) -> Query
+
+    def self.facet_query: (
+      Index index,
+      Symbol field,
+      String path,
+      **untyped options
+    ) -> Query
+
+    def self.range_query: (
+      Index index,
+      fields fields,
+      Range[numeric | date] range,
+      **untyped options
+    ) -> Query
+
+    def self.smart_query: (
+      Index index,
+      fields fields,
+      String query_string,
+      **untyped options
+    ) -> Query
+
+    def self.__new_all_query: () -> Query
+    def self.__new_empty_query: () -> Query
+
+    def self.__new_term_query: (
+      Index index,
+      String field,
+      String term
+    ) -> Query
+
+    def self.__new_fuzzy_term_query: (
+      Index index,
+      String field,
+      String term,
+      Integer distance
+    ) -> Query
+
+    def self.__new_regex_query: (
+      Index index,
+      String field,
+      String regex
+    ) -> Query
+
+    def self.__new_range_query: (
+      Index index,
+      String field,
+      untyped from,
+      untyped to
+    ) -> Query
+
+    def self.__new_phrase_query: (
+      Index index,
+      String field,
+      Array[String] terms
+    ) -> Query
+
+    def self.__new_facet_query: (
+      Index index,
+      String field,
+      String path
+    ) -> Query
+
+    def self.__disjunction: (Array[Query] queries) -> Query
+    def self.__conjunction: (Array[Query] queries) -> Query
+
+    def |: (Query query) -> Query
+    def &: (Query query) -> Query
+    def !: () -> Query
+    def boost: (numeric boost_factor) -> Query
+
+    def __negation: () -> Query
+    def __boost: (Float boost_factor) -> Query
+
+    private
+    
+    def self.construct_query: (
+      Index index,
+      Symbol query_type,
+      Array[Symbol] allowed_fields,
+      fields fields,
+      Array[untyped] params,
+      **untyped options
+    ) -> Query
+
+    def self.text_and_strings: (Index index) -> Array[Symbol]
+  end
+end

data/sig/tantiny/schema.rbs

@@ -0,0 +1,26 @@
+module Tantiny
+  class Schema
+    attr_reader default_tokenizer: Tokenizer
+    attr_reader id_field: Symbol
+    attr_reader text_fields: Array[Symbol]
+    attr_reader string_fields: Array[Symbol]
+    attr_reader integer_fields: Array[Symbol]
+    attr_reader double_fields: Array[Symbol]
+    attr_reader date_fields: Array[Symbol]
+    attr_reader facet_fields: Array[Symbol]
+    attr_reader field_tokenizers: Hash[Symbol, Tokenizer]
+
+    def initialize: (Tokenizer tokenizer) { (*untyped) -> void } -> void
+    def tokenizer_for: (Symbol field) -> Tokenizer
+
+    private
+
+    def id: (Symbol key) -> void
+    def text: (Symbol key, ?tokenizer: Tokenizer) -> void
+    def string: (Symbol key) -> void
+    def integer: (Symbol key) -> void
+    def double: (Symbol key) -> void
+    def date: (Symbol key) -> void
+    def facet: (Symbol key) -> void
+  end
+end

data/sig/tantiny/tokenizer.rbs

@@ -0,0 +1,25 @@
+module Tantiny
+  class Tokenizer
+    def self.default: () -> Tokenizer
+
+    def self.new: (Symbol kind, **untyped options) -> Tokenizer
+
+    def self.__new_ngram_tokenizer: (
+      Integer min,
+      Integer max,
+      bool prefix_only
+    ) -> Tokenizer
+
+    def self.__new_stemmer_tokenizer: (
+      String locale_code
+    ) -> Tokenizer
+
+    def self.__new_simple_tokenizer: () -> Tokenizer
+
+    public
+
+    def terms: (String string) -> Array[String]
+
+    def __extract_terms: (String string) -> Array[String]
+  end
+end

data/sig/tantiny/version.rbs

@@ -0,0 +1,3 @@
+module Tantiny
+  VERSION: String
+end

data/sig/tantiny.rbs

@@ -0,0 +1,5 @@
+module Tantiny
+    type date = Date | DateTime
+    type numeric = Integer | Float
+    type fields = Array[Symbol] | Symbol
+end

data/src/helpers.rs

@@ -0,0 +1,202 @@
+use std::collections::HashMap;
+use rutie::{AnyException, Array, Exception, RString, Hash, Integer, Float, Boolean, Module};
+use tantivy::tokenizer::Language;
+
+// Macro dependencies:
+pub(super) use paste::paste;
+pub(super) use rutie::{class, wrappable_struct, AnyObject, VerifiedObject, VM, Object, Class};
+
+pub(crate) fn namespace() -> Module {
+    Module::from_existing("Tantiny")
+}
+
+pub(crate) struct LanguageWrapper(pub(crate) Language);
+
+impl std::str::FromStr for LanguageWrapper {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "en" => Ok(LanguageWrapper(Language::English)),
+            "ar" => Ok(LanguageWrapper(Language::Arabic)),
+            "da" => Ok(LanguageWrapper(Language::Danish)),
+            "nl" => Ok(LanguageWrapper(Language::Dutch)),
+            "fi" => Ok(LanguageWrapper(Language::Finnish)),
+            "fr" => Ok(LanguageWrapper(Language::French)),
+            "de" => Ok(LanguageWrapper(Language::German)),
+            "el" => Ok(LanguageWrapper(Language::Greek)),
+            "hu" => Ok(LanguageWrapper(Language::Hungarian)),
+            "it" => Ok(LanguageWrapper(Language::Italian)),
+            "no" => Ok(LanguageWrapper(Language::Norwegian)),
+            "pt" => Ok(LanguageWrapper(Language::Portuguese)),
+            "ro" => Ok(LanguageWrapper(Language::Romanian)),
+            "ru" => Ok(LanguageWrapper(Language::Russian)),
+            "es" => Ok(LanguageWrapper(Language::Spanish)),
+            "sv" => Ok(LanguageWrapper(Language::Swedish)),
+            "ta" => Ok(LanguageWrapper(Language::Tamil)),
+            "tr" => Ok(LanguageWrapper(Language::Turkish)),
+            _ => Err(format!("Language '{}' is not supported.", s)),
+        }
+    }
+}
+
+pub(crate) trait TryUnwrap<T> {
+    fn try_unwrap(self) -> T;
+}
+
+macro_rules! primitive_try_unwrap_impl {
+    ( $ruby_type:ty, $type:ty ) => {
+        paste! {
+            impl TryUnwrap<$type> for $ruby_type {
+                fn try_unwrap(self) -> $type {
+                    self.[<to_ $type:lower>]()
+                }
+            }
+
+            impl TryUnwrap<$type> for AnyObject {
+                fn try_unwrap(self) -> $type {
+                    self.try_convert_to::<$ruby_type>()
+                        .try_unwrap()
+                        .[<to_ $type:lower>]()
+                }
+            }
+        }
+    };
+}
+
+primitive_try_unwrap_impl!(RString, String);
+primitive_try_unwrap_impl!(Integer, i64);
+primitive_try_unwrap_impl!(Float, f64);
+primitive_try_unwrap_impl!(Boolean, bool);
+
+impl<T> TryUnwrap<Vec<T>> for Array where
+    AnyObject: TryUnwrap<T>
+{
+    fn try_unwrap(self) -> Vec<T> {
+        let mut vec = Vec::new();
+
+        for elem in self {
+            vec.push(elem.try_unwrap());
+        }
+
+        vec
+    }
+}
+
+impl<K, V> TryUnwrap<HashMap<K, V>> for Hash where
+    AnyObject: TryUnwrap<K> + TryUnwrap<V>,
+    K: Eq + std::hash::Hash
+{
+    fn try_unwrap(self) -> HashMap<K, V> {
+        let mut hashmap = HashMap::new();
+
+        self.each(|key, value| {
+            hashmap.insert(key.try_unwrap(), value.try_unwrap());
+        });
+
+        hashmap
+    }
+}
+
+impl<T, E> TryUnwrap<T> for Result<T, E>
+where
+    E: ToString,
+{
+    fn try_unwrap(self) -> T {
+        self.map_err(|e| {
+            VM::raise_ex(AnyException::new(
+                "Tantiny::TantivyError",
+                Some(&e.to_string()),
+            ))
+        })
+        .unwrap()
+    }
+}
+
+impl<T> TryUnwrap<T> for Option<T> {
+    fn try_unwrap(self) -> T {
+        if let Some(value) = self {
+            value
+        } else {
+            VM::raise_ex(AnyException::new(
+                "Tantiny::UnexpectedNone",
+                Some(&*format!("{}", std::any::type_name::<T>())))
+            );
+
+            self.unwrap()
+        }
+    }
+}
+
+macro_rules! try_unwrap_params {
+    (
+        $param:ident: $type:ty,
+        $( $rest:tt )*
+    ) => {
+        let _tmp = $param.map_err(|e| $crate::helpers::VM::raise_ex(e)).unwrap();
+        let $param = <_ as $crate::helpers::TryUnwrap<$type>>::try_unwrap(_tmp);
+
+        try_unwrap_params!($($rest)*)
+    };
+    (
+        $param:ident,
+        $( $rest:tt )*
+    ) => {
+        let $param = $param.map_err(|e| $crate::helpers::VM::raise_ex(e)).unwrap();
+
+        try_unwrap_params!($($rest)*)
+    };
+
+    // Handle optional trailing commas.
+    ( $param:ident: $type:ty ) => {
+        try_unwrap_params!($param: $type,)
+    };
+    ( $param:ident ) => {
+        try_unwrap_params!($param,)
+    };
+
+    () => {}
+}
+
+pub(crate) use try_unwrap_params;
+
+macro_rules! scaffold {
+    ( $ruby_type:ident, $type:ty, $klass:literal ) => {
+        $crate::helpers::class!($ruby_type);
+
+        // There is a bug in Rutie which prevents using this macro
+        // by resolving it by a full path, so the only workaround is:
+        use crate::helpers::wrappable_struct;
+         
+        $crate::helpers::paste! {
+            wrappable_struct!(
+                $type,
+                [<$type Wrapper>],
+                [<$type:snake:upper _WRAPPER>]
+            );
+        }
+
+        pub(crate) fn klass() -> $crate::helpers::Class {
+            $crate::helpers::namespace().get_nested_class($klass)
+        }
+
+        impl $crate::helpers::TryUnwrap<$ruby_type> for $crate::helpers::AnyObject {
+            fn try_unwrap(self) -> $ruby_type {
+                let result = self.try_convert_to::<$ruby_type>();
+                <_ as $crate::helpers::TryUnwrap<$ruby_type>>::try_unwrap(result)
+            }
+        }
+
+        impl $crate::helpers::VerifiedObject for $ruby_type {
+            fn is_correct_type<T: $crate::helpers::Object>(object: &T) -> bool {
+                object.class() == klass()
+            }
+
+            fn error_message() -> &'static str {
+                concat!("Error converting to ", stringify!($ruby_type), ".")
+            }
+        }
+    }
+}
+
+pub(crate) use scaffold;

data/src/index.rs

@@ -0,0 +1,286 @@
+use std::collections::HashMap;
+use std::str::FromStr;
+use rutie::{methods, Object, AnyObject, Integer, NilClass, Array, RString, Hash};
+use tantivy::{doc, Document, Term, ReloadPolicy, Index, IndexWriter, IndexReader, DateTime};
+use tantivy::schema::{Schema, TextOptions, TextFieldIndexing, IndexRecordOption, FacetOptions, STRING, STORED, INDEXED, FAST};
+use tantivy::collector::TopDocs;
+
+use crate::helpers::{scaffold, try_unwrap_params, TryUnwrap};
+use crate::query::{unwrap_query, RTantinyQuery};
+use crate::tokenizer::{unwrap_tokenizer, RTantinyTokenizer};
+
+pub struct TantinyIndex {
+    pub(crate) schema: Schema,
+    pub(crate) index: Index,
+    pub(crate) index_writer: Option<IndexWriter>,
+    pub(crate) index_reader: IndexReader,
+}
+
+scaffold!(RTantinyIndex, TantinyIndex, "Index");
+
+pub(crate) fn unwrap_index(index: &RTantinyIndex) -> &TantinyIndex {
+    index.get_data(&*TANTINY_INDEX_WRAPPER)
+}
+
+pub(crate) fn unwrap_index_mut(index: &mut RTantinyIndex) -> &mut TantinyIndex {
+    index.get_data_mut(&*TANTINY_INDEX_WRAPPER)
+}
+
+#[rustfmt::skip::macros(methods)]
+methods!(
+    RTantinyIndex,
+    _itself,
+
+    fn new_index(
+        _path: RString,
+        default_tokenizer: AnyObject,
+        field_tokenizers: Hash,
+        text_fields: Array,
+        string_fields: Array,
+        integer_fields: Array,
+        double_fields: Array,
+        date_fields: Array,
+        facet_fields: Array
+    ) -> RTantinyIndex {
+        try_unwrap_params!(
+            _path: String,
+            default_tokenizer: RTantinyTokenizer,
+            field_tokenizers: HashMap<String, RTantinyTokenizer>,
+            text_fields: Vec<String>,
+            string_fields: Vec<String>,
+            integer_fields: Vec<String>,
+            double_fields: Vec<String>,
+            date_fields: Vec<String>,
+            facet_fields: Vec<String>
+        );
+
+        let mut schema_builder = Schema::builder();
+
+        schema_builder.add_text_field("id", STRING | STORED);
+
+        for field in text_fields {
+            let tokenizer_name =
+                if field_tokenizers.contains_key(&field) {
+                    &*field
+                } else {
+                    "default"
+                };
+            let indexing = TextFieldIndexing::default()
+                .set_tokenizer(tokenizer_name)
+                .set_index_option(IndexRecordOption::WithFreqsAndPositions);
+            let options = TextOptions::default()
+                .set_indexing_options(indexing);
+            schema_builder.add_text_field(&field, options);
+        }
+
+        for field in string_fields {
+            schema_builder.add_text_field(&field, STRING);
+        }
+
+        for field in integer_fields {
+            schema_builder.add_i64_field(&field, FAST | INDEXED);
+        }
+
+        for field in double_fields {
+            schema_builder.add_f64_field(&field, FAST | INDEXED);
+        }
+
+        for field in date_fields {
+            schema_builder.add_date_field(&field, FAST | INDEXED);
+        }
+
+        for field in facet_fields {
+            let options = FacetOptions::default().set_indexed();
+            schema_builder.add_facet_field(&field, options);
+        }
+
+        let schema = schema_builder.build();
+        let index = Index::create_in_ram(schema.clone());
+        let tokenizers = index.tokenizers();
+
+        tokenizers.register("default", unwrap_tokenizer(&default_tokenizer).clone());
+
+        for (field, tokenizer) in field_tokenizers {
+            tokenizers.register(&field, unwrap_tokenizer(&tokenizer).clone())
+        }
+
+        let index_writer = None;
+
+        let index_reader = index
+            .reader_builder()
+            .reload_policy(ReloadPolicy::Manual)
+            .try_into()
+            .try_unwrap();
+
+        klass().wrap_data(
+            TantinyIndex { index, index_writer, index_reader, schema },
+            &*TANTINY_INDEX_WRAPPER
+        )
+    }
+
+    fn add_document(
+        id: RString,
+        text_fields: Hash,
+        string_fields: Hash,
+        integer_fields: Hash,
+        double_fields: Hash,
+        date_fields: Hash,
+        facet_fields: Hash
+    ) -> NilClass {
+        try_unwrap_params!(
+            id: String,
+            text_fields: HashMap<String, String>,
+            string_fields: HashMap<String, String>,
+            integer_fields: HashMap<String, i64>,
+            double_fields: HashMap<String, f64>,
+            date_fields: HashMap<String, String>,
+            facet_fields: HashMap<String, String>
+        );
+
+        let internal = unwrap_index(&_itself);
+        let index_writer = internal.index_writer.as_ref().try_unwrap();
+        let schema = &internal.schema;
+
+        let mut doc = Document::default();
+
+        let id_field = schema.get_field("id").try_unwrap();
+        doc.add_text(id_field, &id);
+
+        for (key, value) in text_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_text(field, value);
+        }
+
+        for (key, value) in string_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_text(field, value);
+        }
+
+        for (key, &value) in integer_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_i64(field, value);
+        }
+
+        for (key, &value) in double_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_f64(field, value);
+        }
+
+        for (key, value) in date_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            let value = DateTime::from_str(value).try_unwrap();
+            doc.add_date(field, &value);
+        }
+
+        for (key, value) in facet_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_facet(field, &value);
+        }
+
+        let doc_id = Term::from_field_text(id_field, &id);
+        index_writer.delete_term(doc_id.clone());
+
+        index_writer.add_document(doc);
+
+        NilClass::new()
+    }
+
+    fn delete_document(id: RString) -> NilClass {
+        try_unwrap_params!(id: String);
+
+        let internal = unwrap_index(&_itself);
+        let index_writer = internal.index_writer.as_ref().unwrap();
+
+        let id_field = internal.schema.get_field("id").try_unwrap();
+        let doc_id = Term::from_field_text(id_field, &id);
+
+        index_writer.delete_term(doc_id.clone());
+
+        NilClass::new()
+    }
+
+    fn acquire_index_writer(
+        overall_memory: Integer
+    ) -> NilClass {
+        try_unwrap_params!(overall_memory: i64);
+
+        let internal = unwrap_index_mut(&mut _itself);
+
+        let mut index_writer = internal.index
+            .writer(overall_memory as usize)
+            .try_unwrap();
+
+        internal.index_writer = Some(index_writer);
+
+        NilClass::new()
+    }
+
+    fn release_index_writer() -> NilClass {
+        let internal = unwrap_index_mut(&mut _itself);
+
+        drop(internal.index_writer.as_ref().try_unwrap());
+        internal.index_writer = None;
+
+        NilClass::new()
+    }
+
+    fn commit() -> NilClass {
+        let internal = unwrap_index_mut(&mut _itself);
+        let index_writer = internal.index_writer.as_mut().try_unwrap();
+
+        index_writer.commit().try_unwrap();
+
+        NilClass::new()
+    }
+
+    fn reload() -> NilClass {
+        unwrap_index(&_itself).index_reader.reload().try_unwrap();
+
+        NilClass::new()
+    }
+
+    fn search(
+        query: AnyObject,
+        limit: Integer
+    ) -> Array {
+        try_unwrap_params!(
+            query: RTantinyQuery,
+            limit: i64
+        );
+
+        let internal = unwrap_index(&_itself);
+        let id_field = internal.schema.get_field("id").try_unwrap();
+        let searcher = internal.index_reader.searcher();
+        let query = unwrap_query(&query);
+
+        let top_docs = searcher
+            .search(query, &TopDocs::with_limit(limit as usize))
+            .try_unwrap();
+
+        let mut array = Array::with_capacity(top_docs.len());
+
+        for (_score, doc_address) in top_docs {
+            let doc = searcher.doc(doc_address).try_unwrap();
+            if let Some(value) = doc.get_first(id_field) {
+                if let Some(id) = (&*value).text() {
+                    array.push(RString::from(String::from(id)));
+                }
+            }
+        }
+
+        array
+    }
+);
+
+pub(super) fn init() {
+    klass().define(|klass| {
+        klass.def_self("__new", new_index);
+        klass.def("__add_document", add_document);
+        klass.def("__delete_document", delete_document);
+        klass.def("__acquire_index_writer", acquire_index_writer);
+        klass.def("__release_index_writer", release_index_writer);
+        klass.def("__commit", commit);
+        klass.def("__reload", reload);
+        klass.def("__search", search);
+    });
+}

data/src/lib.rs

@@ -0,0 +1,14 @@
+mod helpers;
+#[allow(improper_ctypes_definitions)]
+mod index;
+#[allow(improper_ctypes_definitions)]
+mod query;
+#[allow(improper_ctypes_definitions)]
+mod tokenizer;
+
+#[no_mangle]
+pub extern "C" fn Init_tantiny() {
+    index::init();
+    query::init();
+    tokenizer::init();
+}