linkrs_query/executor/expression/functions/
fulltext.rs

1//! Full-Text Search Expression Functions
2//!
3//! This module provides expression functions for full-text search, including:
4//! - score(): Get the relevance score of a document
5//! - highlight(): Get highlighted text fragments
6//! - matched_fields(): Get the list of matched fields
7//! - snippet(): Get a text snippet
8
9use crate::executor::expression::functions::signature::{FunctionSignature, ValueType};
10use crate::executor::expression::{ExpressionError, ExpressionErrorType};
11use linkrs_core::Value;
12use linkrs_fulltext::FulltextSearchEntry;
13use std::collections::HashMap;
14
15/// Full-text search function enumeration
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17pub enum FulltextFunction {
18    /// score() - Get document relevance score
19    Score,
20    /// highlight(field, [pre_tag], [post_tag], [fragment_size]) - Get highlighted text
21    Highlight,
22    /// matched_fields() - Get list of matched fields
23    MatchedFields,
24    /// snippet(field, [max_len]) - Get text snippet
25    Snippet,
26    /// rank() - Get the rank of the document in search results
27    Rank,
28    /// search_match(field, query) - Check if field matches query pattern
29    SearchMatch,
30    /// search_score(field) - Get field-specific relevance score
31    FieldScore,
32}
33
34impl FulltextFunction {
35    /// Get function name
36    pub fn name(&self) -> &'static str {
37        match self {
38            FulltextFunction::Score => "score",
39            FulltextFunction::Highlight => "highlight",
40            FulltextFunction::MatchedFields => "matched_fields",
41            FulltextFunction::Snippet => "snippet",
42            FulltextFunction::Rank => "rank",
43            FulltextFunction::SearchMatch => "search_match",
44            FulltextFunction::FieldScore => "field_score",
45        }
46    }
47
48    /// Get function signature
49    pub fn signature(&self) -> FunctionSignature {
50        match self {
51            FulltextFunction::Score => {
52                FunctionSignature::new("score", vec![], Some(ValueType::Float), false)
53            }
54            FulltextFunction::Highlight => FunctionSignature::new(
55                "highlight",
56                vec![
57                    ValueType::String,
58                    ValueType::String,
59                    ValueType::String,
60                    ValueType::Int,
61                ],
62                Some(ValueType::String),
63                true,
64            ),
65            FulltextFunction::MatchedFields => {
66                FunctionSignature::new("matched_fields", vec![], Some(ValueType::List), false)
67            }
68            FulltextFunction::Snippet => FunctionSignature::new(
69                "snippet",
70                vec![ValueType::String, ValueType::Int],
71                Some(ValueType::String),
72                true,
73            ),
74            FulltextFunction::Rank => {
75                FunctionSignature::new("rank", vec![], Some(ValueType::Int), false)
76            }
77            FulltextFunction::SearchMatch => FunctionSignature::new(
78                "search_match",
79                vec![ValueType::String, ValueType::String],
80                Some(ValueType::Bool),
81                false,
82            ),
83            FulltextFunction::FieldScore => FunctionSignature::new(
84                "field_score",
85                vec![ValueType::String],
86                Some(ValueType::Float),
87                false,
88            ),
89        }
90    }
91
92    /// Get the number of parameters
93    pub fn arity(&self) -> usize {
94        match self {
95            FulltextFunction::Score => 0,
96            FulltextFunction::Highlight => 4,
97            FulltextFunction::MatchedFields => 0,
98            FulltextFunction::Snippet => 2,
99            FulltextFunction::Rank => 0,
100            FulltextFunction::SearchMatch => 2,
101            FulltextFunction::FieldScore => 1,
102        }
103    }
104
105    /// Check if variable parameters are accepted
106    pub fn is_variadic(&self) -> bool {
107        match self {
108            FulltextFunction::Score => false,
109            FulltextFunction::Highlight => true,
110            FulltextFunction::MatchedFields => false,
111            FulltextFunction::Snippet => true,
112            FulltextFunction::Rank => false,
113            FulltextFunction::SearchMatch => false,
114            FulltextFunction::FieldScore => false,
115        }
116    }
117
118    /// Get function description
119    pub fn description(&self) -> &'static str {
120        match self {
121            FulltextFunction::Score => "Get the relevance score of a document in full-text search",
122            FulltextFunction::Highlight => "Get highlighted text fragments for a field",
123            FulltextFunction::MatchedFields => "Get the list of matched fields",
124            FulltextFunction::Snippet => "Get a text snippet from a field",
125            FulltextFunction::Rank => "Get the rank position of a document in search results",
126            FulltextFunction::SearchMatch => "Check if a field matches a search query pattern",
127            FulltextFunction::FieldScore => "Get the field-specific relevance score",
128        }
129    }
130
131    /// Execute the function
132    pub fn execute(
133        &self,
134        args: &[Value],
135        context: &FulltextExecutionContext,
136    ) -> Result<Value, ExpressionError> {
137        match self {
138            FulltextFunction::Score => self.execute_score(args, context),
139            FulltextFunction::Highlight => self.execute_highlight(args, context),
140            FulltextFunction::MatchedFields => self.execute_matched_fields(args, context),
141            FulltextFunction::Snippet => self.execute_snippet(args, context),
142            FulltextFunction::Rank => self.execute_rank(args, context),
143            FulltextFunction::SearchMatch => self.execute_search_match(args, context),
144            FulltextFunction::FieldScore => self.execute_field_score(args, context),
145        }
146    }
147
148    /// Execute score() function
149    fn execute_score(
150        &self,
151        args: &[Value],
152        context: &FulltextExecutionContext,
153    ) -> Result<Value, ExpressionError> {
154        if !args.is_empty() {
155            return Err(ExpressionError::new(
156                ExpressionErrorType::InvalidArgumentCount,
157                format!("score() expects 0 arguments, got {}", args.len()),
158            ));
159        }
160
161        Ok(Value::Double(context.score))
162    }
163
164    /// Execute highlight() function
165    fn execute_highlight(
166        &self,
167        args: &[Value],
168        context: &FulltextExecutionContext,
169    ) -> Result<Value, ExpressionError> {
170        if args.is_empty() {
171            return Err(ExpressionError::new(
172                ExpressionErrorType::InvalidArgumentCount,
173                "highlight() expects at least 1 argument (field name)",
174            ));
175        }
176
177        let field_name = match &args[0] {
178            Value::String(s) => s.to_string(),
179            _ => {
180                return Err(ExpressionError::new(
181                    ExpressionErrorType::TypeError,
182                    "highlight() first argument must be a string (field name)",
183                ));
184            }
185        };
186
187        let pre_tag = args
188            .get(1)
189            .and_then(|v| {
190                if let Value::String(s) = v {
191                    Some(s.to_string())
192                } else {
193                    None
194                }
195            })
196            .unwrap_or_else(|| "<em>".to_string());
197
198        let post_tag = args
199            .get(2)
200            .and_then(|v| {
201                if let Value::String(s) = v {
202                    Some(s.to_string())
203                } else {
204                    None
205                }
206            })
207            .unwrap_or_else(|| "</em>".to_string());
208
209        let fragment_size = args
210            .get(3)
211            .and_then(|v| {
212                if let Value::Int(n) = v {
213                    Some(*n as usize)
214                } else {
215                    None
216                }
217            })
218            .unwrap_or(100);
219
220        // Get highlights from context
221        if let Some(highlights) = &context.highlights {
222            if let Some(field_highlights) = highlights.get(field_name.as_str()) {
223                if !field_highlights.is_empty() {
224                    let highlighted_text = field_highlights.join(" ... ");
225                    return Ok(Value::string(highlighted_text));
226                }
227            }
228        }
229
230        // Return original text if no highlights available
231        if let Some(source) = &context.source {
232            if let Some(Value::String(text)) = source.get(field_name.as_str()) {
233                // Truncate if needed
234                let char_count = text.chars().count();
235                if char_count > fragment_size {
236                    let truncated: String = text.chars().take(fragment_size).collect();
237                    return Ok(Value::string(format!(
238                        "{}{}{}",
239                        pre_tag, truncated, post_tag
240                    )));
241                }
242                return Ok(Value::string(text.clone()));
243            }
244        }
245
246        Ok(Value::Null(linkrs_core::null::NullType::Null))
247    }
248
249    /// Execute matched_fields() function
250    fn execute_matched_fields(
251        &self,
252        args: &[Value],
253        context: &FulltextExecutionContext,
254    ) -> Result<Value, ExpressionError> {
255        if !args.is_empty() {
256            return Err(ExpressionError::new(
257                ExpressionErrorType::InvalidArgumentCount,
258                "matched_fields() expects 0 arguments",
259            ));
260        }
261
262        let fields: Vec<Value> = context
263            .matched_fields
264            .iter()
265            .map(|f| Value::string(f.clone()))
266            .collect();
267
268        Ok(Value::list(linkrs_core::value::list::List {
269            values: fields,
270        }))
271    }
272
273    /// Execute snippet() function
274    fn execute_snippet(
275        &self,
276        args: &[Value],
277        context: &FulltextExecutionContext,
278    ) -> Result<Value, ExpressionError> {
279        if args.is_empty() {
280            return Err(ExpressionError::new(
281                ExpressionErrorType::InvalidArgumentCount,
282                "snippet() expects at least 1 argument (field name)",
283            ));
284        }
285
286        let field_name = match &args[0] {
287            Value::String(s) => s.to_string(),
288            _ => {
289                return Err(ExpressionError::new(
290                    ExpressionErrorType::TypeError,
291                    "snippet() first argument must be a string (field name)",
292                ));
293            }
294        };
295
296        let max_len = args
297            .get(1)
298            .and_then(|v| {
299                if let Value::Int(n) = v {
300                    Some(*n as usize)
301                } else {
302                    None
303                }
304            })
305            .unwrap_or(200);
306
307        // Get text from source
308        if let Some(source) = &context.source {
309            if let Some(Value::String(text)) = source.get(field_name.as_str()) {
310                if text.chars().count() <= max_len {
311                    return Ok(Value::string(text.clone()));
312                }
313
314                // Try to find a good break point
315                let prefix: String = text.chars().take(max_len).collect();
316                let break_point = prefix.rfind(' ').unwrap_or(prefix.len());
317
318                return Ok(Value::string(format!("{}...", &prefix[..break_point])));
319            }
320        }
321
322        Ok(Value::Null(linkrs_core::null::NullType::Null))
323    }
324
325    /// Execute rank() function
326    fn execute_rank(
327        &self,
328        args: &[Value],
329        context: &FulltextExecutionContext,
330    ) -> Result<Value, ExpressionError> {
331        if !args.is_empty() {
332            return Err(ExpressionError::new(
333                ExpressionErrorType::InvalidArgumentCount,
334                format!("rank() expects 0 arguments, got {}", args.len()),
335            ));
336        }
337
338        let rank = if context.score > 0.0 {
339            ((1.0 / context.score) * 100.0) as i64
340        } else {
341            0
342        };
343
344        Ok(Value::BigInt(rank))
345    }
346
347    /// Execute search_match() function
348    fn execute_search_match(
349        &self,
350        args: &[Value],
351        context: &FulltextExecutionContext,
352    ) -> Result<Value, ExpressionError> {
353        if args.len() < 2 {
354            return Err(ExpressionError::new(
355                ExpressionErrorType::InvalidArgumentCount,
356                "search_match() expects 2 arguments (field, query)",
357            ));
358        }
359
360        let field_name = match &args[0] {
361            Value::String(s) => s.to_string(),
362            _ => {
363                return Err(ExpressionError::new(
364                    ExpressionErrorType::TypeError,
365                    "search_match() first argument must be a string (field name)",
366                ));
367            }
368        };
369
370        let query = match &args[1] {
371            Value::String(s) => s.to_lowercase(),
372            _ => {
373                return Err(ExpressionError::new(
374                    ExpressionErrorType::TypeError,
375                    "search_match() second argument must be a string (query)",
376                ));
377            }
378        };
379
380        if let Some(source) = &context.source {
381            if let Some(Value::String(text)) = source.get(field_name.as_str()) {
382                return Ok(Value::Bool(text.to_lowercase().contains(query.as_str())));
383            }
384        }
385
386        Ok(Value::Bool(false))
387    }
388
389    /// Execute field_score() function
390    fn execute_field_score(
391        &self,
392        args: &[Value],
393        context: &FulltextExecutionContext,
394    ) -> Result<Value, ExpressionError> {
395        if args.is_empty() {
396            return Err(ExpressionError::new(
397                ExpressionErrorType::InvalidArgumentCount,
398                "field_score() expects 1 argument (field name)",
399            ));
400        }
401
402        let field_name = match &args[0] {
403            Value::String(s) => s.to_string(),
404            _ => {
405                return Err(ExpressionError::new(
406                    ExpressionErrorType::TypeError,
407                    "field_score() first argument must be a string (field name)",
408                ));
409            }
410        };
411
412        if context.matched_fields.contains(&field_name) {
413            let field_bonus = if context
414                .highlights
415                .as_ref()
416                .is_some_and(|h| h.contains_key(field_name.as_str()))
417            {
418                1.5
419            } else {
420                1.0
421            };
422            return Ok(Value::Double(context.score * field_bonus));
423        }
424
425        Ok(Value::Double(0.0))
426    }
427}
428
429/// Full-text search execution context
430#[derive(Debug, Clone)]
431pub struct FulltextExecutionContext {
432    /// Current document score
433    pub score: f64,
434    /// Highlighted text by field
435    pub highlights: Option<HashMap<String, Vec<String>>>,
436    /// Matched field names
437    pub matched_fields: Vec<String>,
438    /// Source document data
439    pub source: Option<HashMap<String, Value>>,
440}
441
442impl Default for FulltextExecutionContext {
443    fn default() -> Self {
444        Self {
445            score: 0.0,
446            highlights: None,
447            matched_fields: Vec::new(),
448            source: None,
449        }
450    }
451}
452
453impl FulltextExecutionContext {
454    /// Create a new execution context
455    pub fn new() -> Self {
456        Self::default()
457    }
458
459    /// Create context from search result entry
460    pub fn from_search_entry(entry: &FulltextSearchEntry) -> Self {
461        Self {
462            score: entry.score,
463            highlights: entry.highlights.clone(),
464            matched_fields: entry.matched_fields.clone(),
465            source: entry.source.clone(),
466        }
467    }
468
469    /// Set the score
470    pub fn with_score(mut self, score: f64) -> Self {
471        self.score = score;
472        self
473    }
474
475    /// Set the highlights
476    pub fn with_highlights(mut self, highlights: HashMap<String, Vec<String>>) -> Self {
477        self.highlights = Some(highlights);
478        self
479    }
480
481    /// Set the matched fields
482    pub fn with_matched_fields(mut self, fields: Vec<String>) -> Self {
483        self.matched_fields = fields;
484        self
485    }
486
487    /// Set the source
488    pub fn with_source(mut self, source: HashMap<String, Value>) -> Self {
489        self.source = Some(source);
490        self
491    }
492}
493
494/// Register full-text search functions
495pub fn register_fulltext_functions(
496    registry: &mut crate::executor::expression::functions::FunctionRegistry,
497) {
498    registry.register_builtin(
499        crate::executor::expression::functions::BuiltinFunction::Fulltext(FulltextFunction::Score),
500    );
501
502    registry.register_builtin(
503        crate::executor::expression::functions::BuiltinFunction::Fulltext(
504            FulltextFunction::Highlight,
505        ),
506    );
507
508    registry.register_builtin(
509        crate::executor::expression::functions::BuiltinFunction::Fulltext(
510            FulltextFunction::MatchedFields,
511        ),
512    );
513
514    registry.register_builtin(
515        crate::executor::expression::functions::BuiltinFunction::Fulltext(
516            FulltextFunction::Snippet,
517        ),
518    );
519
520    registry.register_builtin(
521        crate::executor::expression::functions::BuiltinFunction::Fulltext(FulltextFunction::Rank),
522    );
523
524    registry.register_builtin(
525        crate::executor::expression::functions::BuiltinFunction::Fulltext(
526            FulltextFunction::SearchMatch,
527        ),
528    );
529
530    registry.register_builtin(
531        crate::executor::expression::functions::BuiltinFunction::Fulltext(
532            FulltextFunction::FieldScore,
533        ),
534    );
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    fn create_test_context() -> FulltextExecutionContext {
542        let mut source = HashMap::new();
543        source.insert("title".into(), Value::string("Database Optimization"));
544        source.insert(
545            "content".to_string(),
546            Value::string("This is a test article about database optimization techniques."),
547        );
548
549        let mut highlights = HashMap::new();
550        highlights.insert(
551            "content".to_string(),
552            vec!["<em>database</em> optimization".to_string()],
553        );
554
555        FulltextExecutionContext {
556            score: 0.85,
557            highlights: Some(highlights),
558            matched_fields: vec!["title".to_string(), "content".to_string()],
559            source: Some(source),
560        }
561    }
562
563    #[test]
564    fn test_score_function() {
565        let func = FulltextFunction::Score;
566        let context = create_test_context();
567
568        let result = func.execute(&[], &context).unwrap();
569        assert!(matches!(result, Value::Double(_)));
570        if let Value::Double(score) = result {
571            assert!((score - 0.85).abs() < 0.001);
572        }
573    }
574
575    #[test]
576    fn test_highlight_function() {
577        let func = FulltextFunction::Highlight;
578        let context = create_test_context();
579
580        let result = func.execute(&[Value::string("content")], &context).unwrap();
581
582        assert!(matches!(result, Value::String(_)));
583        if let Value::String(text) = result {
584            assert!(text.contains("<em>"));
585            assert!(text.contains("</em>"));
586        }
587    }
588
589    #[test]
590    fn test_matched_fields_function() {
591        let func = FulltextFunction::MatchedFields;
592        let context = create_test_context();
593
594        let result = func.execute(&[], &context).unwrap();
595        assert!(matches!(result, Value::List(_)));
596
597        if let Value::List(fields) = result {
598            assert_eq!(fields.len(), 2);
599            assert!(fields.contains(&Value::string("title")));
600            assert!(fields.contains(&Value::string("content")));
601        }
602    }
603
604    #[test]
605    fn test_snippet_function() {
606        let func = FulltextFunction::Snippet;
607        let context = create_test_context();
608
609        let result = func
610            .execute(&[Value::string("content"), Value::Int(50)], &context)
611            .unwrap();
612
613        assert!(matches!(result, Value::String(_)));
614        if let Value::String(text) = result {
615            assert!(text.len() <= 53); // 50 + "..."
616        }
617    }
618
619    #[test]
620    fn test_snippet_non_ascii_truncation() {
621        let func = FulltextFunction::Snippet;
622        let mut source = HashMap::new();
623        source.insert(
624            "content".to_string(),
625            Value::string("你好世界这是一个测试内容"),
626        );
627        let context = FulltextExecutionContext {
628            source: Some(source),
629            ..FulltextExecutionContext::default()
630        };
631        // Byte slicing at max_len would panic mid-character; character
632        // truncation returns the first 5 characters plus an ellipsis.
633        let result = func
634            .execute(&[Value::string("content"), Value::Int(5)], &context)
635            .unwrap();
636        assert_eq!(result, Value::string("你好世界这..."));
637    }
638
639    #[test]
640    fn test_highlight_non_ascii_truncation() {
641        let func = FulltextFunction::Highlight;
642        let mut source = HashMap::new();
643        source.insert("title".into(), Value::string("数据库优化指南"));
644        let context = FulltextExecutionContext {
645            source: Some(source),
646            ..FulltextExecutionContext::default()
647        };
648        let result = func
649            .execute(
650                &[
651                    Value::string("title"),
652                    Value::string("<b>"),
653                    Value::string("</b>"),
654                    Value::Int(2),
655                ],
656                &context,
657            )
658            .unwrap();
659        assert_eq!(result, Value::string("<b>数据</b>"));
660    }
661}
linkrs_query/executor/expression/functions/fulltext.rs

linkrs_query/executor/expression/functions/
fulltext.rs
