tantiny 0.3.3 → 0.4.1

data/src/query.rs

@@ -1,260 +1,323 @@
-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;
+use crate::tokenizer::Tokenizer;
 
-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,
+        query_string: String,
+        fuzzy_distance: i64,
+        tokenizer: &Tokenizer,
+        last_term_min_length_prefix_match: i64,
+    ) -> Result<String, Error> {
+        use tantivy::tokenizer::TokenStream;
+        let mut analyzer = tokenizer.get_analyzer();
+
+        // Tokenizer the query string
+        let query_tokens = {
+            let mut query_string_token_stream = analyzer.token_stream(&query_string);
+            let mut tokens = Vec::new();
+            while query_string_token_stream.advance() {
+                let token = query_string_token_stream.token();
+                tokens.push(token.text.clone());
+            }
+            tokens
+        };
+
+        // Tokenize the input text
+        let mut input_text_token_stream = analyzer.token_stream(&text);
+        let mut tokens = Vec::new();
+        while input_text_token_stream.advance() {
+            let token = input_text_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<_> = query_tokens
+            .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 fuzzy_match = query_tokens.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(_))
+            });
+
+            // Check if this token is a prefix match for the last query term
+            let prefix_match = token_text.len() > last_term_min_length_prefix_match as usize
+                && query_tokens
+                    .last()
+                    .map(|last_token| token_text.starts_with(last_token))
+                    .unwrap_or(false);
+            let should_highlight = fuzzy_match || prefix_match;
+
+            // 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;
+        }
 
-        let query = BoostQuery::new(unwrap_query(&_itself).box_clone(), score as f32);
+        // Add any remaining text after the last token
+        result.push_str(&text[last_pos..]);
 
-        wrap_query(Box::new(query))
+        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, 5))?;
+
+    Ok(())
+}