Added float fields to schema with tests (quickwit-oss#36)

* Added float fields to schema with tests

* Fixed typo

src/document.rs

@@ -292,6 +292,15 @@ impl Document {
  add_value(self, field_name, value);
  }
 
+ /// Add a float value to the document.
+ ///
+ /// Args:
+ /// field_name (str): The field name for which we are adding the value.
+ /// value (f64): The float that will be added to the document.
+ fn add_float(&mut self, field_name: String, value: f64) {
+ add_value(self, field_name, value);
+ }
+
  /// Add a date value to the document.
  ///
  /// Args:

src/schemabuilder.rs

@@ -119,7 +119,7 @@ impl SchemaBuilder {
  ) -> PyResult<Self> {
  let builder = &mut self.builder;
 
- let opts = SchemaBuilder::build_int_option(stored, indexed, fast)?;
+ let opts = SchemaBuilder::build_numeric_option(stored, indexed, fast)?;
 
  if let Some(builder) = builder.write().unwrap().as_mut() {
  builder.add_i64_field(name, opts);
@@ -131,6 +131,28 @@ impl SchemaBuilder {
  Ok(self.clone())
  }
 
+ #[args(stored = false, indexed = false)]
+ fn add_float_field(
+ &mut self,
+ name: &str,
+ stored: bool,
+ indexed: bool,
+ fast: Option<&str>,
+ ) -> PyResult<Self> {
+ let builder = &mut self.builder;
+
+ let opts = SchemaBuilder::build_numeric_option(stored, indexed, fast)?;
+
+ if let Some(builder) = builder.write().unwrap().as_mut() {
+ builder.add_f64_field(name, opts);
+ } else {
+ return Err(exceptions::PyValueError::new_err(
+ "Schema builder object isn't valid anymore.",
+ ));
+ }
+ Ok(self.clone())
+ }
+
  /// Add a new unsigned integer field to the schema.
  ///
  /// Args:
@@ -161,7 +183,7 @@ impl SchemaBuilder {
  ) -> PyResult<Self> {
  let builder = &mut self.builder;
 
- let opts = SchemaBuilder::build_int_option(stored, indexed, fast)?;
+ let opts = SchemaBuilder::build_numeric_option(stored, indexed, fast)?;
 
  if let Some(builder) = builder.write().unwrap().as_mut() {
  builder.add_u64_field(name, opts);
@@ -203,7 +225,7 @@ impl SchemaBuilder {
  ) -> PyResult<Self> {
  let builder = &mut self.builder;
 
- let opts = SchemaBuilder::build_int_option(stored, indexed, fast)?;
+ let opts = SchemaBuilder::build_numeric_option(stored, indexed, fast)?;
 
  if let Some(builder) = builder.write().unwrap().as_mut() {
  builder.add_date_field(name, opts);
@@ -319,7 +341,7 @@ impl SchemaBuilder {
 }
 
 impl SchemaBuilder {
- fn build_int_option(
+ fn build_numeric_option(
  stored: bool,
  indexed: bool,
  fast: Option<&str>,

tests/tantivy_test.py

@@ -12,6 +12,14 @@ def schema():
  .build()
  )
 
+def schema_numeric_fields():
+ return (
+ SchemaBuilder()
+ .add_integer_field("id", stored=True, indexed=True)
+ .add_float_field("rating", stored=True, indexed=True)
+ .add_text_field("body", stored=True)
+ .build()
+ )
 
 def create_index(dir=None):
  # assume all tests will use the same documents for now
@@ -66,6 +74,46 @@ def create_index(dir=None):
  index.reload()
  return index
 
+def create_index_with_numeric_fields(dir=None):
+ index = Index(schema_numeric_fields(), dir)
+ writer = index.writer()
+
+ doc = Document()
+ doc.add_integer("id", 1)
+ doc.add_float("rating", 3.5)
+ doc.add_text(
+ "body",
+ (
+ "He was an old man who fished alone in a skiff in"
+ "the Gulf Stream and he had gone eighty-four days "
+ "now without taking a fish."
+ ),
+ )
+ writer.add_document(doc)
+ doc = Document.from_dict(
+ {
+ "id": 2,
+ "rating": 4.5,
+ "body": (
+ "A few miles south of Soledad, the Salinas River drops "
+ "in close to the hillside bank and runs deep and "
+ "green. The water is warm too, for it has slipped "
+ "twinkling over the yellow sands in the sunlight "
+ "before reaching the narrow pool. On one side of the "
+ "river the golden foothill slopes curve up to the "
+ "strong and rocky Gabilan Mountains, but on the valley "
+ "side the water is lined with trees—willows fresh and "
+ "green with every spring, carrying in their lower leaf "
+ "junctures the debris of the winter’s flooding; and "
+ "sycamores with mottled, white, recumbent limbs and "
+ "branches that arch over the pool"
+ ),
+ }
+ )
+ writer.add_document(doc)
+ writer.commit()
+ index.reload()
+ return index
 
 @pytest.fixture()
 def dir_index(tmpdir):
@@ -77,6 +125,11 @@ def ram_index():
  return create_index()
 
 
+@pytest.fixture(scope="class")
+def ram_index_numeric_fields():
+ return create_index_with_numeric_fields()
+
+
 class TestClass(object):
  def test_simple_search_in_dir(self, dir_index):
  _, index = dir_index
@@ -120,6 +173,25 @@ def test_and_query(self, ram_index):
 
  assert len(result.hits) == 1
 
+ def test_and_query_numeric_fields(self, ram_index_numeric_fields):
+ index = ram_index_numeric_fields
+ searcher = index.searcher()
+
+ # 1 result
+ float_query = index.parse_query("3.5", ["rating"])
+ result = searcher.search(float_query)
+ assert len(result.hits) == 1
+ assert searcher.doc(result.hits[0][1])['rating'][0] == 3.5
+
+ integer_query = index.parse_query("1", ["id"])
+ result = searcher.search(integer_query)
+ assert len(result.hits) == 1
+
+ # 0 result
+ integer_query = index.parse_query("10", ["id"])
+ result = searcher.search(integer_query)
+ assert len(result.hits) == 0
+
  def test_and_query_parser_default_fields(self, ram_index):
  query = ram_index.parse_query("winter", default_field_names=["title"])
  assert repr(query) == """Query(TermQuery(Term(type=Str, field=0, "winter")))"""
@@ -279,8 +351,9 @@ def test_create_readers(self):
 
 
 class TestSearcher(object):
- def test_searcher_repr(self, ram_index):
+ def test_searcher_repr(self, ram_index, ram_index_numeric_fields):
  assert repr(ram_index.searcher()) == "Searcher(num_docs=3, num_segments=1)"
+ assert repr(ram_index_numeric_fields.searcher()) == "Searcher(num_docs=2, num_segments=1)"
 
 
 class TestDocument(object):