tantiny 0.3.3 → 0.4.0

data/src/query.rs

@@ -1,260 +1,303 @@
-use std::str::FromStr;
+use levenshtein_automata::{Distance, LevenshteinAutomatonBuilder};
+use magnus::{Error, Module, Object, RArray, RModule, Ruby, TryConvert, Value};
 use std::ops::Bound::Included;
-use rutie::{methods, Object, AnyObject, Integer, Float, Array, RString};
-use tantivy::{Term, DateTime};
-use tantivy::schema::{IndexRecordOption, Facet, Type, FieldType};
 use tantivy::query::*;
+use tantivy::schema::{Facet, FieldType, IndexRecordOption};
+use tantivy::Term;
+use time::OffsetDateTime;
 
-use crate::helpers::{try_unwrap_params, scaffold, TryUnwrap};
-use crate::index::{unwrap_index, RTantinyIndex};
+use crate::index::Index;
 
-pub struct TantinyQuery(pub(crate) Box<dyn Query>);
+#[magnus::wrap(class = "Tantiny::Query", free_immediately, size)]
+pub struct Query(Box<dyn tantivy::query::Query>);
 
-scaffold!(RTantinyQuery, TantinyQuery, "Query");
-
-fn wrap_query(query: Box<dyn Query>) -> RTantinyQuery {
-    klass().wrap_data(
-        TantinyQuery(query),
-        &*TANTINY_QUERY_WRAPPER
-    )
-}
-
-pub(crate) fn unwrap_query(query: &RTantinyQuery) -> &Box<dyn Query> {
-    &query.get_data(&*TANTINY_QUERY_WRAPPER).0
-}
-
-#[rustfmt::skip::macros(methods)]
-methods!(
-    RTantinyQuery,
-    _itself,
+impl Query {
+    pub fn get_query(&self) -> &dyn tantivy::query::Query {
+        self.0.as_ref()
+    }
 
-    fn new_all_query() -> RTantinyQuery {
-        wrap_query(Box::new(AllQuery))
+    fn new_all() -> Self {
+        Query(Box::new(AllQuery))
     }
 
-    fn new_empty_query() -> RTantinyQuery {
-        wrap_query(Box::new(EmptyQuery))
+    fn new_empty() -> Self {
+        Query(Box::new(EmptyQuery))
     }
 
-    fn new_term_query(
-        index: RTantinyIndex,
-        field: RString,
-        term: RString
-    ) -> RTantinyQuery {
-        try_unwrap_params!(
-            index,
-            field: String,
-            term: String
-        );
-
-        let schema = &unwrap_index(&index).schema;
-        let field = schema.get_field(&field).try_unwrap();
+    fn new_term(index: &Index, field: String, term: String) -> Result<Self, Error> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let field = index.schema.get_field(&field).map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Field not found: {}", e),
+            )
+        })?;
         let term = Term::from_field_text(field, &term);
         let query = TermQuery::new(term, IndexRecordOption::Basic);
-
-        wrap_query(Box::new(query))
+        Ok(Query(Box::new(query)))
     }
 
-    fn new_fuzzy_term_query(
-        index: RTantinyIndex,
-        field: RString,
-        term: RString,
-        distance: Integer
-    ) -> RTantinyQuery {
-        try_unwrap_params!(
-            index,
-            field: String,
-            term: String,
-            distance: i64
-        );
-
-        let schema = &unwrap_index(&index).schema;
-        let field = schema.get_field(&field).try_unwrap();
+    fn new_fuzzy_term(
+        index: &Index,
+        field: String,
+        term: String,
+        distance: i64,
+    ) -> Result<Self, Error> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let field = index.schema.get_field(&field).map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Field not found: {}", e),
+            )
+        })?;
         let term = Term::from_field_text(field, &term);
         let query = FuzzyTermQuery::new(term, distance as u8, true);
-
-        wrap_query(Box::new(query))
+        Ok(Query(Box::new(query)))
     }
 
-    fn new_phrase_query(
-        index: RTantinyIndex,
-        field: RString,
-        terms: Array
-    ) -> RTantinyQuery {
-        try_unwrap_params!(
-            index,
-            field: String,
-            terms: Vec<String>
-        );
-
-        let schema = &unwrap_index(&index).schema;
-        let field = schema.get_field(&field).try_unwrap();
-
-        let terms: Vec<Term> = terms.into_iter().map(|term| {
-            Term::from_field_text(field, &term)
-        }).collect();
+    fn new_phrase(index: &Index, field: String, terms: Vec<String>) -> Result<Self, Error> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let field = index.schema.get_field(&field).map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Field not found: {}", e),
+            )
+        })?;
+
+        let terms: Vec<Term> = terms
+            .into_iter()
+            .map(|term| Term::from_field_text(field, &term))
+            .collect();
         let query = PhraseQuery::new(terms);
-
-        wrap_query(Box::new(query))
+        Ok(Query(Box::new(query)))
     }
 
-    fn new_regex_query(
-        index: RTantinyIndex,
-        field: RString,
-        regex: RString
-    ) -> RTantinyQuery {
-        try_unwrap_params!(
-            index,
-            field: String,
-            regex: String
-        );
-
-        let schema = &unwrap_index(&index).schema;
-        let field = schema.get_field(&field).try_unwrap();
-        let query = RegexQuery::from_pattern(&regex, field).try_unwrap();
-
-        wrap_query(Box::new(query))
+    fn new_regex(index: &Index, field: String, regex: String) -> Result<Self, Error> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let field = index.schema.get_field(&field).map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Field not found: {}", e),
+            )
+        })?;
+        let query = RegexQuery::from_pattern(&regex, field).map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Invalid regex: {}", e),
+            )
+        })?;
+        Ok(Query(Box::new(query)))
     }
 
-    fn new_range_query(
-        index: RTantinyIndex,
-        field: RString,
-        from: AnyObject,
-        to: AnyObject
-    ) -> RTantinyQuery {
-        try_unwrap_params!(index, from, to, field: String);
-
-        let schema = &unwrap_index(&index).schema;
-        let field = schema.get_field(&field).try_unwrap();
-        let field_name = schema.get_field_name(field);
-        let field_type = schema.get_field_entry(field).field_type();
-
-        let range = match field_type {
+    fn new_range(index: &Index, field: String, from: Value, to: Value) -> Result<Self, Error> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let field_obj = index.schema.get_field(&field).map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Field not found: {}", e),
+            )
+        })?;
+        let field_name = index.schema.get_field_name(field_obj);
+        let field_type = index.schema.get_field_entry(field_obj).field_type();
+
+        let (left, right) = match field_type {
             FieldType::Date(_) => {
-                let from: String = from.try_unwrap();
-                let to: String = to.try_unwrap();
-                let from = DateTime::from_str(&from).try_unwrap();
-                let to = DateTime::from_str(&to).try_unwrap();
-
-                Ok((
-                    Type::Date,
-                    Included(Term::from_field_date(field, &from)),
-                    Included(Term::from_field_date(field, &to))
-                ))
-            },
+                let from_str: String = String::try_convert(from)?;
+                let to_str: String = String::try_convert(to)?;
+                let from_datetime = OffsetDateTime::parse(
+                    &from_str,
+                    &time::format_description::well_known::Rfc3339,
+                )
+                .map_err(|e| {
+                    Error::new(
+                        ruby.exception_runtime_error(),
+                        format!("Invalid date format: {}", e),
+                    )
+                })?;
+                let to_datetime =
+                    OffsetDateTime::parse(&to_str, &time::format_description::well_known::Rfc3339)
+                        .map_err(|e| {
+                            Error::new(
+                                ruby.exception_runtime_error(),
+                                format!("Invalid date format: {}", e),
+                            )
+                        })?;
+                let from_dt = tantivy::DateTime::from_timestamp_nanos(
+                    from_datetime.unix_timestamp_nanos() as i64,
+                );
+                let to_dt = tantivy::DateTime::from_timestamp_nanos(
+                    to_datetime.unix_timestamp_nanos() as i64,
+                );
+
+                (
+                    Term::from_field_date(field_obj, from_dt),
+                    Term::from_field_date(field_obj, to_dt),
+                )
+            }
             FieldType::I64(_) => {
-                let from: i64 = from.try_unwrap();
-                let to: i64 = to.try_unwrap();
-
-                Ok((
-                    Type::I64,
-                    Included(Term::from_field_i64(field, from)),
-                    Included(Term::from_field_i64(field, to))
-                ))
-            },
+                let from_val: i64 = i64::try_convert(from)?;
+                let to_val: i64 = i64::try_convert(to)?;
+                (
+                    Term::from_field_i64(field_obj, from_val),
+                    Term::from_field_i64(field_obj, to_val),
+                )
+            }
             FieldType::F64(_) => {
-                let from: f64 = from.try_unwrap();
-                let to: f64 = to.try_unwrap();
-
-                Ok((
-                    Type::F64,
-                    Included(Term::from_field_f64(field, from)),
-                    Included(Term::from_field_f64(field, to))
+                let from_val: f64 = f64::try_convert(from)?;
+                let to_val: f64 = f64::try_convert(to)?;
+                (
+                    Term::from_field_f64(field_obj, from_val),
+                    Term::from_field_f64(field_obj, to_val),
+                )
+            }
+            _ => {
+                return Err(Error::new(
+                    ruby.exception_runtime_error(),
+                    format!("Field '{}' is not supported by range query.", field_name),
                 ))
-            },
-            _ => { Err(format!("Field '{}' is not supported by range query.", field_name)) }
+            }
         };
 
-        let (value_type, left, right) = range.try_unwrap();
-
-        let query = RangeQuery::new_term_bounds(field, value_type, &left, &right);
-
-        wrap_query(Box::new(query))
+        let query = RangeQuery::new(Included(left), Included(right));
+        Ok(Query(Box::new(query)))
     }
 
-    fn new_facet_query(
-        index: RTantinyIndex,
-        field: RString,
-        path: RString
-    ) -> RTantinyQuery {
-        try_unwrap_params!(
-            index,
-            field: String,
-            path: String
-        );
-
-        let schema = &unwrap_index(&index).schema;
-        let field = schema.get_field(&field).try_unwrap();
+    fn new_facet(index: &Index, field: String, path: String) -> Result<Self, Error> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let field = index.schema.get_field(&field).map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Field not found: {}", e),
+            )
+        })?;
         let facet = Facet::from(&path);
         let term = Term::from_facet(field, &facet);
         let query = TermQuery::new(term, IndexRecordOption::Basic);
-
-        wrap_query(Box::new(query))
+        Ok(Query(Box::new(query)))
     }
 
-    fn disjunction(queries: Array) -> RTantinyQuery {
-        try_unwrap_params!(queries);
-
+    fn disjunction(queries: RArray) -> Result<Self, Error> {
         let mut query_vec = Vec::new();
 
-        for query in queries {
-            let query: RTantinyQuery = query.try_unwrap();
-            query_vec.push((Occur::Should, unwrap_query(&query).box_clone()));
+        for item in queries.into_iter() {
+            let query: &Query = <&Query>::try_convert(item)?;
+            query_vec.push((Occur::Should, query.0.box_clone()));
         }
 
-        let disjunction_query = BooleanQuery::from(query_vec);
-
-        wrap_query(Box::new(disjunction_query))
+        Ok(Query(Box::new(BooleanQuery::from(query_vec))))
     }
 
-    fn conjunction(queries: Array) -> RTantinyQuery {
-        try_unwrap_params!(queries);
-
+    fn conjunction(queries: RArray) -> Result<Self, Error> {
         let mut query_vec = Vec::new();
 
-        for query in queries {
-            let query: RTantinyQuery = query.try_unwrap();
-            query_vec.push((Occur::Must, unwrap_query(&query).box_clone()));
+        for item in queries.into_iter() {
+            let query: &Query = <&Query>::try_convert(item)?;
+            query_vec.push((Occur::Must, query.0.box_clone()));
         }
 
-        let conjunction_query = BooleanQuery::from(query_vec);
-
-        wrap_query(Box::new(conjunction_query))
+        Ok(Query(Box::new(BooleanQuery::from(query_vec))))
     }
 
-    fn negation() -> RTantinyQuery {
-        // See: https://github.com/quickwit-oss/tantivy/issues/1153
-        let all_query: Box<dyn Query> = Box::new(AllQuery);
+    fn negation(&self) -> Self {
+        let all_query: Box<dyn tantivy::query::Query> = Box::new(AllQuery);
         let negation_query = BooleanQuery::from(vec![
             (Occur::Must, all_query.box_clone()),
-            (Occur::MustNot, unwrap_query(&_itself).box_clone()),
+            (Occur::MustNot, self.0.box_clone()),
         ]);
 
-        wrap_query(Box::new(negation_query))
+        Query(Box::new(negation_query))
+    }
+
+    fn boost(&self, score: f64) -> Self {
+        let query = BoostQuery::new(self.0.box_clone(), score as f32);
+        Query(Box::new(query))
     }
 
-    fn boost(score: Float) -> RTantinyQuery {
-        try_unwrap_params!(score: f64);
+    fn highlight(text: String, terms: Vec<String>, fuzzy_distance: i64) -> Result<String, Error> {
+        use tantivy::tokenizer::{LowerCaser, SimpleTokenizer, TextAnalyzer, TokenStream};
+
+        // Create a simple tokenizer for highlighting
+        let mut analyzer = TextAnalyzer::builder(SimpleTokenizer::default())
+            .filter(LowerCaser)
+            .build();
 
-        let query = BoostQuery::new(unwrap_query(&_itself).box_clone(), score as f32);
+        // Tokenize the input text
+        let mut token_stream = analyzer.token_stream(&text);
 
-        wrap_query(Box::new(query))
+        // Collect all tokens with their positions
+        let mut tokens = Vec::new();
+        while token_stream.advance() {
+            let token = token_stream.token();
+            tokens.push((token.text.clone(), token.offset_from, token.offset_to));
+        }
+
+        // Build Levenshtein automata for each term (same algorithm as Tantivy's FuzzyTermQuery)
+        let lev_builder = LevenshteinAutomatonBuilder::new(fuzzy_distance as u8, true);
+        let automata: Vec<_> = terms
+            .iter()
+            .map(|term| lev_builder.build_dfa(term))
+            .collect();
+
+        // Build the highlighted text
+        let mut result = String::new();
+        let mut last_pos = 0;
+
+        for (token_text, start, end) in tokens {
+            // Check if this token matches any of the query terms (exact or fuzzy)
+            let should_highlight = terms.iter().zip(&automata).any(|(term, dfa)| {
+                // Exact match
+                if token_text.eq_ignore_ascii_case(term) {
+                    return true;
+                }
+
+                // Fuzzy match using Levenshtein automaton (same as Tantivy's FuzzyTermQuery)
+                matches!(dfa.eval(&token_text), Distance::Exact(_))
+            });
+
+            // Add the text before the token
+            result.push_str(&text[last_pos..start]);
+
+            // Add the token, highlighted if it matches
+            if should_highlight {
+                result.push_str("<b>");
+                result.push_str(&text[start..end]);
+                result.push_str("</b>");
+            } else {
+                result.push_str(&text[start..end]);
+            }
+
+            last_pos = end;
+        }
+
+        // Add any remaining text after the last token
+        result.push_str(&text[last_pos..]);
+
+        Ok(result)
     }
-);
-
-pub(super) fn init() {
-    klass().define(|klass| {
-        klass.def_self("__new_all_query", new_all_query);
-        klass.def_self("__new_empty_query", new_empty_query);
-        klass.def_self("__new_term_query", new_term_query);
-        klass.def_self("__new_fuzzy_term_query", new_fuzzy_term_query);
-        klass.def_self("__new_regex_query", new_regex_query);
-        klass.def_self("__new_range_query", new_range_query);
-        klass.def_self("__new_phrase_query", new_phrase_query);
-        klass.def_self("__new_facet_query", new_facet_query);
-        klass.def_self("__disjunction", disjunction);
-        klass.def_self("__conjunction", conjunction);
-        klass.def("__negation", negation);
-        klass.def("__boost", boost);
-    });
-} 
+}
+
+pub fn init(ruby: &Ruby, module: RModule) -> Result<(), Error> {
+    let class = module.define_class("Query", ruby.class_object())?;
+
+    class.define_singleton_method("__new_all_query", magnus::function!(Query::new_all, 0))?;
+    class.define_singleton_method("__new_empty_query", magnus::function!(Query::new_empty, 0))?;
+    class.define_singleton_method("__new_term_query", magnus::function!(Query::new_term, 3))?;
+    class.define_singleton_method(
+        "__new_fuzzy_term_query",
+        magnus::function!(Query::new_fuzzy_term, 4),
+    )?;
+    class.define_singleton_method(
+        "__new_phrase_query",
+        magnus::function!(Query::new_phrase, 3),
+    )?;
+    class.define_singleton_method("__new_regex_query", magnus::function!(Query::new_regex, 3))?;
+    class.define_singleton_method("__new_range_query", magnus::function!(Query::new_range, 4))?;
+    class.define_singleton_method("__new_facet_query", magnus::function!(Query::new_facet, 3))?;
+    class.define_singleton_method("__disjunction", magnus::function!(Query::disjunction, 1))?;
+    class.define_singleton_method("__conjunction", magnus::function!(Query::conjunction, 1))?;
+    class.define_method("__negation", magnus::method!(Query::negation, 0))?;
+    class.define_method("__boost", magnus::method!(Query::boost, 1))?;
+    class.define_singleton_method("__highlight", magnus::function!(Query::highlight, 3))?;
+
+    Ok(())
+}

data/src/tokenizer.rs

@@ -1,94 +1,81 @@
+use magnus::{Error, Module, Object, RModule, Ruby};
+use tantivy::tokenizer::{
+    LowerCaser, NgramTokenizer, RemoveLongFilter, SimpleTokenizer, Stemmer, TextAnalyzer,
+};
 
-use rutie::{methods, Object, Array, RString, Integer, Boolean};
-use tantivy::tokenizer::{TextAnalyzer, SimpleTokenizer, RemoveLongFilter, LowerCaser, Stemmer, NgramTokenizer};
+use crate::helpers::LanguageWrapper;
 
-use crate::helpers::{try_unwrap_params, scaffold, TryUnwrap, LanguageWrapper};
+#[magnus::wrap(class = "Tantiny::Tokenizer", free_immediately, size)]
+pub struct Tokenizer(TextAnalyzer);
 
-pub struct TantinyTokenizer(pub(crate) TextAnalyzer);
-
-scaffold!(RTantinyTokenizer, TantinyTokenizer, "Tokenizer");
-
-fn wrap_tokenizer(tokenizer: TextAnalyzer) -> RTantinyTokenizer {
-    klass().wrap_data(
-        TantinyTokenizer(tokenizer),
-        &*TANTINY_TOKENIZER_WRAPPER
-    )
-}
-
-pub(crate) fn unwrap_tokenizer(tokenizer: &RTantinyTokenizer) -> &TextAnalyzer {
-    &tokenizer.get_data(&*TANTINY_TOKENIZER_WRAPPER).0
-}
-
-#[rustfmt::skip::macros(methods)]
-methods!(
-    RTantinyTokenizer,
-    _itself,
+impl Tokenizer {
+    pub fn get_analyzer(&self) -> TextAnalyzer {
+        self.0.clone()
+    }
 
-    fn new_simple_tokenizer() -> RTantinyTokenizer {
-        let tokenizer = TextAnalyzer::from(SimpleTokenizer)
+    fn new_simple() -> Result<Self, Error> {
+        let tokenizer = TextAnalyzer::builder(SimpleTokenizer::default())
             .filter(RemoveLongFilter::limit(40))
-            .filter(LowerCaser);
-
-        wrap_tokenizer(tokenizer)
+            .filter(LowerCaser)
+            .build();
+        Ok(Tokenizer(tokenizer))
     }
 
-    fn new_stemmer_tokenizer(locale_code: RString) -> RTantinyTokenizer {
-        try_unwrap_params!(locale_code: String);
-
-        let language: LanguageWrapper = locale_code.parse().try_unwrap();
-        let tokenizer = TextAnalyzer::from(SimpleTokenizer)
+    fn new_stemmer(language: String) -> Result<Self, Error> {
+        let lang_wrapper = LanguageWrapper::try_from(language)?;
+        let tokenizer = TextAnalyzer::builder(SimpleTokenizer::default())
             .filter(RemoveLongFilter::limit(40))
             .filter(LowerCaser)
-            .filter(Stemmer::new(language.0));
-
-        wrap_tokenizer(tokenizer)
+            .filter(Stemmer::new(lang_wrapper.0))
+            .build();
+        Ok(Tokenizer(tokenizer))
     }
 
-    fn new_ngram_tokenizer(
-        min_gram: Integer,
-        max_gram: Integer,
-        prefix_only: Boolean
-    ) -> RTantinyTokenizer {
-        try_unwrap_params!(
-            min_gram: i64,
-            max_gram: i64,
-            prefix_only: bool
-        );
-
-        let tokenizer = NgramTokenizer::new(
-            min_gram as usize,
-            max_gram as usize,
-            prefix_only
-        );
-
-        wrap_tokenizer(TextAnalyzer::from(tokenizer))
+    fn new_ngram(min_gram: i64, max_gram: i64, prefix_only: bool) -> Result<Self, Error> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let tokenizer = NgramTokenizer::new(min_gram as usize, max_gram as usize, prefix_only)
+            .map_err(|e| {
+                Error::new(
+                    ruby.exception_runtime_error(),
+                    format!("Failed to create ngram tokenizer: {}", e),
+                )
+            })?;
+
+        Ok(Tokenizer(TextAnalyzer::builder(tokenizer).build()))
     }
 
-    fn extract_terms(text: RString) -> Array {
-        try_unwrap_params!(text: String);
-
-        let mut token_stream = unwrap_tokenizer(&_itself).token_stream(&text);
-        let mut terms = vec![];
+    fn extract_terms(&self, text: String) -> Result<Vec<String>, Error> {
+        let mut cloned_analyzer = self.0.clone();
+        let mut token_stream = cloned_analyzer.token_stream(&text);
+        let mut terms = Vec::new();
 
         while token_stream.advance() {
-            terms.push(token_stream.token().clone().text);
+            terms.push(token_stream.token().text.clone());
         }
 
-        let mut array = Array::with_capacity(terms.len());
-
-        for term in terms {
-            array.push(RString::from(term));
-        }
-
-        array
+        Ok(terms)
     }
-);
+}
 
-pub(super) fn init() {
-    klass().define(|klass| {
-        klass.def_self("__new_simple_tokenizer", new_simple_tokenizer);
-        klass.def_self("__new_stemmer_tokenizer", new_stemmer_tokenizer);
-        klass.def_self("__new_ngram_tokenizer", new_ngram_tokenizer);
-        klass.def("__extract_terms", extract_terms);
-    });
-} 
+pub fn init(ruby: &Ruby, module: RModule) -> Result<(), Error> {
+    let class = module.define_class("Tokenizer", ruby.class_object())?;
+
+    class.define_singleton_method(
+        "__new_simple_tokenizer",
+        magnus::function!(Tokenizer::new_simple, 0),
+    )?;
+    class.define_singleton_method(
+        "__new_stemmer_tokenizer",
+        magnus::function!(Tokenizer::new_stemmer, 1),
+    )?;
+    class.define_singleton_method(
+        "__new_ngram_tokenizer",
+        magnus::function!(Tokenizer::new_ngram, 3),
+    )?;
+    class.define_method(
+        "__extract_terms",
+        magnus::method!(Tokenizer::extract_terms, 1),
+    )?;
+
+    Ok(())
+}