Skip to content

Add Float Support #65

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Mar 18, 2023
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/document.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
30 changes: 26 additions & 4 deletions src/schemabuilder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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:
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -319,7 +341,7 @@ impl SchemaBuilder {
}

impl SchemaBuilder {
fn build_int_option(
fn build_numeric_option(
stored: bool,
indexed: bool,
fast: Option<&str>,
Expand Down
75 changes: 74 additions & 1 deletion tests/tantivy_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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")))"""
Expand Down Expand Up @@ -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):
Expand Down