Snippet generator (fixes #36 and #63) (#122)

Co-authored-by: Justin Greene <justin.greene@intouchsol.com>
Co-authored-by: Paul Masurel <paul@quickwit.io>
Co-authored-by: Pratyush Mittal <pratyushmittal@gmail.com>
Co-authored-by: mukeshsahnis <er.mks89@gmail.com>

.gitignore

@@ -9,6 +9,6 @@ tantivy.so
 tantivy.dylib
 tantivy/tantivy.cpython*.so
 tantivy.egg-info/
-
-# Exclude the mkdocs site directory
-site/
+.venv
+.envrc
+site/

README.md

@@ -23,6 +23,21 @@ only supports python3.
 
 # Development
 
+For compiling Python module:
+
+```bash
+# create virtual env
+python -m venv .venv
+source .venv/bin/activate
+
+# install maturin, the build tool for PyO3
+pip install maturin
+
+# compile and install python module in venv
+maturin develop
+```
+
+
 Setting up a development environment can be done in a virtual environment using
 [`nox`](https://nox.thea.codes) or using local packages using the provided `Makefile`.
 

src/document.rs

@@ -797,7 +797,7 @@ impl Document {
  Ok(())
  }
 
- fn iter_values_for_field<'a>(
+ pub fn iter_values_for_field<'a>(
  &'a self,
  field: &str,
  ) -> impl Iterator<Item = &'a Value> + 'a {

src/lib.rs

@@ -8,13 +8,16 @@ mod query;
 mod schema;
 mod schemabuilder;
 mod searcher;
+mod snippet;
 
 use document::Document;
 use facet::Facet;
 use index::Index;
+use query::Query;
 use schema::Schema;
 use schemabuilder::SchemaBuilder;
 use searcher::{DocAddress, SearchResult, Searcher};
+use snippet::{Snippet, SnippetGenerator};
 
 /// Python bindings for the search engine library Tantivy.
 ///
@@ -76,6 +79,9 @@ fn tantivy(_py: Python, m: &PyModule) -> PyResult<()> {
  m.add_class::<Index>()?;
  m.add_class::<DocAddress>()?;
  m.add_class::<Facet>()?;
+ m.add_class::<Query>()?;
+ m.add_class::<Snippet>()?;
+ m.add_class::<SnippetGenerator>()?;
  Ok(())
 }
 

src/snippet.rs

@@ -0,0 +1,81 @@
+use crate::to_pyerr;
+use pyo3::prelude::*;
+use tantivy as tv;
+
+/// Tantivy schema.
+///
+/// The schema is very strict. To build the schema the `SchemaBuilder` class is
+/// provided.
+#[pyclass]
+pub(crate) struct Snippet {
+ pub(crate) inner: tv::Snippet,
+}
+
+#[pyclass]
+pub(crate) struct Range {
+ #[pyo3(get)]
+ start: usize,
+ #[pyo3(get)]
+ end: usize,
+}
+
+#[pymethods]
+impl Snippet {
+ pub fn to_html(&self) -> PyResult<String> {
+ Ok(self.inner.to_html())
+ }
+
+ pub fn highlighted(&self) -> Vec<Range> {
+ let highlighted = self.inner.highlighted();
+ let results = highlighted
+ .iter()
+ .map(|r| Range {
+ start: r.start,
+ end: r.end,
+ })
+ .collect::<Vec<_>>();
+ results
+ }
+}
+
+#[pyclass]
+pub(crate) struct SnippetGenerator {
+ pub(crate) field_name: String,
+ pub(crate) inner: tv::SnippetGenerator,
+}
+
+#[pymethods]
+impl SnippetGenerator {
+ #[staticmethod]
+ pub fn create(
+ searcher: &crate::Searcher,
+ query: &crate::Query,
+ schema: &crate::Schema,
+ field_name: &str,
+ ) -> PyResult<SnippetGenerator> {
+ let field = schema
+ .inner
+ .get_field(field_name)
+ .or(Err("field not found"))
+ .map_err(to_pyerr)?;
+ let generator =
+ tv::SnippetGenerator::create(&searcher.inner, query.get(), field)
+ .map_err(to_pyerr)?;
+
+ return Ok(SnippetGenerator {
+ field_name: field_name.to_string(),
+ inner: generator,
+ });
+ }
+
+ pub fn snippet_from_doc(&self, doc: &crate::Document) -> crate::Snippet {
+ let text: String = doc
+ .iter_values_for_field(&self.field_name)
+ .flat_map(tv::schema::Value::as_text)
+ .collect::<Vec<&str>>()
+ .join(" ");
+
+ let result = self.inner.snippet(&text);
+ Snippet { inner: result }
+ }
+}

tests/tantivy_test.py

@@ -5,8 +5,8 @@
 import tantivy
 import pickle
 import pytest
-
-from tantivy import Document, Index, SchemaBuilder
+import tantivy
+from tantivy import Document, Index, SchemaBuilder, SnippetGenerator
 
 
 def schema():
@@ -784,3 +784,27 @@ def test_doc_address_pickle():
  pickled = pickle.loads(pickle.dumps(orig))
 
  assert orig == pickled
+
+
+class TestSnippets(object):
+ def test_document_snippet(self, dir_index):
+ index_dir, _ = dir_index
+ doc_schema = schema()
+ index = Index(doc_schema, str(index_dir))
+ query = index.parse_query("sea whale", ["title", "body"])
+ searcher = index.searcher()
+ result = searcher.search(query)
+ assert len(result.hits) == 1
+
+ snippet_generator = SnippetGenerator.create(searcher, query, doc_schema, "title")
+
+ for (score, doc_address) in result.hits:
+ doc = searcher.doc(doc_address)
+ snippet = snippet_generator.snippet_from_doc(doc)
+ highlights = snippet.highlighted()
+ assert len(highlights) == 1
+ first = highlights[0]
+ assert first.start == 20
+ assert first.end == 23
+ html_snippet = snippet.to_html()
+ assert html_snippet == 'The Old Man and the <b>Sea</b>'