pegasust
/
tantivy-py
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Better support bytes, IPs, and JSON (#152)

master
Chris Tam 2023-11-19 06:08:33 -05:00 committed by GitHub
parent 4ac17da8f6
commit eeaad34a98
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 212 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
src/document.rs
View File

@ -6,9 +6,10 @@ use pyo3::{
basic::CompareOp,
prelude::*,
types::{
PyAny, PyBool, PyDateAccess, PyDateTime, PyDict, PyList, PyTimeAccess,
PyTuple,
PyAny, PyBool, PyDateAccess, PyDateTime, PyDict, PyInt, PyList,
PyTimeAccess, PyTuple,
},
Python,
};
use chrono::{offset::TimeZone, NaiveDateTime, Utc};
@ -23,7 +24,8 @@ use serde_json::Value as JsonValue;
use std::{
collections::{BTreeMap, HashMap},
fmt,
net::Ipv6Addr,
net::{IpAddr, Ipv6Addr},
str::FromStr,
};
pub(crate) fn extract_value(any: &PyAny) -> PyResult<Value> {
@ -50,6 +52,11 @@ pub(crate) fn extract_value(any: &PyAny) -> PyResult<Value> {
if let Ok(b) = any.extract::<Vec<u8>>() {
return Ok(Value::Bytes(b));
}
if let Ok(dict) = any.downcast::<PyDict>() {
if let Ok(json) = pythonize::depythonize(dict) {
return Ok(Value::JsonObject(json));
}
}
Err(to_pyerr(format!("Value unsupported {any:?}")))
}
@ -105,7 +112,37 @@ pub(crate) fn extract_value_for_type(
.map_err(to_pyerr_for_type("Facet", field_name, any))?
.inner,
),
_ => return Err(to_pyerr(format!("Value unsupported {:?}", any))),
tv::schema::Type::Bytes => Value::Bytes(
any.extract::<Vec<u8>>()
.map_err(to_pyerr_for_type("Bytes", field_name, any))?,
),
tv::schema::Type::Json => {
if let Ok(json_str) = any.extract::<&str>() {
return serde_json::from_str(json_str)
.map(Value::JsonObject)
.map_err(to_pyerr_for_type("Json", field_name, any));
}
Value::JsonObject(
any.downcast::<PyDict>()
.map(|dict| pythonize::depythonize(&dict))
.map_err(to_pyerr_for_type("Json", field_name, any))?
.map_err(to_pyerr_for_type("Json", field_name, any))?,
)
}
tv::schema::Type::IpAddr => {
let val = any
.extract::<&str>()
.map_err(to_pyerr_for_type("IpAddr", field_name, any))?;
IpAddr::from_str(val)
.map(|addr| match addr {
IpAddr::V4(addr) => addr.to_ipv6_mapped(),
IpAddr::V6(addr) => addr,
})
.map(Value::IpAddr)
.map_err(to_pyerr_for_type("IpAddr", field_name, any))?
}
};
Ok(value)
@ -126,6 +163,20 @@ fn extract_value_single_or_list_for_type(
) -> PyResult<Vec<Value>> {
// Check if a numeric fast field supports multivalues.
if let Ok(values) = any.downcast::<PyList>() {
// Process an array of integers as a single entry if it is a bytes field.
if field_type.value_type() == tv::schema::Type::Bytes
&& values
.get_item(0)
.map(|v| v.is_instance_of::<PyInt>())
.unwrap_or(false)
{
return Ok(vec![extract_value_for_type(
values,
field_type.value_type(),
field_name,
)?]);
}
values
.iter()
.map(|any| {

92
src/schemabuilder.rs
View File

@ -2,11 +2,11 @@
use pyo3::{exceptions, prelude::*};
use tantivy::schema;
use crate::schema::Schema;
use std::sync::{Arc, RwLock};
use tantivy::schema::{DateOptions, INDEXED};
use tantivy::schema::{
self, BytesOptions, DateOptions, IpAddrOptions, INDEXED,
};
/// Tantivy has a very strict schema.
/// You need to specify in advance whether a field is indexed or not,
@ -357,17 +357,43 @@ impl SchemaBuilder {
/// Add a fast bytes field to the schema.
///
/// Bytes field are not searchable and are only used
/// as fast field, to associate any kind of payload
/// to a document.
///
/// Args:
/// name (str): The name of the field.
fn add_bytes_field(&mut self, name: &str) -> PyResult<Self> {
/// stored (bool, optional): If true sets the field as stored, the
/// content of the field can be later restored from a Searcher.
/// Defaults to False.
/// indexed (bool, optional): If true sets the field to be indexed.
/// fast (str, optional): Set the bytes options as a fast field. A fast
/// field is a column-oriented fashion storage for tantivy. It is
/// designed for the fast random access of some document fields
/// given a document id.
#[pyo3(signature = (
name,
stored = false,
indexed = false,
fast = false
))]
fn add_bytes_field(
&mut self,
name: &str,
stored: bool,
indexed: bool,
fast: bool,
) -> PyResult<Self> {
let builder = &mut self.builder;
let mut opts = BytesOptions::default();
if stored {
opts = opts.set_stored();
}
if indexed {
opts = opts.set_indexed();
}
if fast {
opts = opts.set_fast();
}
if let Some(builder) = builder.write().unwrap().as_mut() {
builder.add_bytes_field(name, INDEXED);
builder.add_bytes_field(name, opts);
} else {
return Err(exceptions::PyValueError::new_err(
"Schema builder object isn't valid anymore.",
@ -376,6 +402,54 @@ impl SchemaBuilder {
Ok(self.clone())
}
/// Add an IP address field to the schema.
///
/// Args:
/// name (str): The name of the field.
/// stored (bool, optional): If true sets the field as stored, the
/// content of the field can be later restored from a Searcher.
/// Defaults to False.
/// indexed (bool, optional): If true sets the field to be indexed.
/// fast (str, optional): Set the IP address options as a fast field. A
/// fast field is a column-oriented fashion storage for tantivy. It
/// is designed for the fast random access of some document fields
/// given a document id.
#[pyo3(signature = (
name,
stored = false,
indexed = false,
fast = false
))]
fn add_ip_addr_field(
&mut self,
name: &str,
stored: bool,
indexed: bool,
fast: bool,
) -> PyResult<Self> {
let builder = &mut self.builder;
let mut opts = IpAddrOptions::default();
if stored {
opts = opts.set_stored();
}
if indexed {
opts = opts.set_indexed();
}
if fast {
opts = opts.set_fast();
}
if let Some(builder) = builder.write().unwrap().as_mut() {
builder.add_ip_addr_field(name, opts);
} else {
return Err(exceptions::PyValueError::new_err(
"Schema builder object isn't valid anymore.",
));
}
Ok(self.clone())
}
/// Finalize the creation of a Schema.
///
/// Returns a Schema object. After this is called the SchemaBuilder cannot

83
tests/tantivy_test.py
View File

@ -2,6 +2,7 @@ from io import BytesIO
import copy
import datetime
import json
import tantivy
import pickle
import pytest
@ -365,7 +366,9 @@ class TestClass(object):
searched_doc = index.searcher().doc(doc_address)
assert searched_doc["title"] == ["Test title"]
result = searcher.search(query, 10, order_by_field="order", order=tantivy.Order.Asc)
result = searcher.search(
query, 10, order_by_field="order", order=tantivy.Order.Asc
)
assert len(result.hits) == 3
@ -443,7 +446,7 @@ class TestClass(object):
assert searcher.num_segments < 8
def test_doc_from_dict_schema_validation(self):
def test_doc_from_dict_numeric_validation(self):
schema = (
SchemaBuilder()
.add_unsigned_field("unsigned")
@ -504,6 +507,70 @@ class TestClass(object):
schema,
)
def test_doc_from_dict_bytes_validation(self):
schema = SchemaBuilder().add_bytes_field("bytes").build()
good = Document.from_dict({"bytes": b"hello"}, schema)
good = Document.from_dict({"bytes": [[1, 2, 3], [4, 5, 6]]}, schema)
good = Document.from_dict({"bytes": [1, 2, 3]}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict({"bytes": [1, 2, 256]}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict({"bytes": "hello"}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict({"bytes": [1024, "there"]}, schema)
def test_doc_from_dict_ip_addr_validation(self):
schema = SchemaBuilder().add_ip_addr_field("ip").build()
good = Document.from_dict({"ip": "127.0.0.1"}, schema)
good = Document.from_dict({"ip": "::1"}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict({"ip": 12309812348}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict({"ip": "256.100.0.1"}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict(
{"ip": "1234:5678:9ABC:DEF0:1234:5678:9ABC:DEF0:1234"}, schema
)
with pytest.raises(ValueError):
bad = Document.from_dict(
{"ip": "1234:5678:9ABC:DEF0:1234:5678:9ABC:GHIJ"}, schema
)
def test_doc_from_dict_json_validation(self):
# Test implicit JSON
good = Document.from_dict({"dict": {"hello": "world"}})
schema = SchemaBuilder().add_json_field("json").build()
good = Document.from_dict({"json": {}}, schema)
good = Document.from_dict({"json": {"hello": "world"}}, schema)
good = Document.from_dict(
{"nested": {"hello": ["world", "!"]}, "numbers": [1, 2, 3]}, schema
)
list_of_jsons = [
{"hello": "world"},
{"nested": {"hello": ["world", "!"]}, "numbers": [1, 2, 3]},
]
good = Document.from_dict({"json": list_of_jsons}, schema)
good = Document.from_dict({"json": json.dumps(list_of_jsons[1])}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict({"json": 123}, schema)
with pytest.raises(ValueError):
bad = Document.from_dict({"json": "hello"}, schema)
def test_search_result_eq(self, ram_index, spanish_index):
eng_index = ram_index
eng_query = eng_index.parse_query("sea whale", ["title", "body"])
@ -650,10 +717,6 @@ class TestDocument(object):
doc = tantivy.Document(facet=facet)
assert doc["facet"][0].to_path() == ["asia/oceania", "fiji"]
def test_document_error(self):
with pytest.raises(ValueError):
tantivy.Document(name={})
def test_document_eq(self):
doc1 = tantivy.Document(name="Bill", reference=[1, 2])
doc2 = tantivy.Document.from_dict({"name": "Bill", "reference": [1, 2]})
@ -848,9 +911,11 @@ class TestSnippets(object):
result = searcher.search(query)
assert len(result.hits) == 1
snippet_generator = SnippetGenerator.create(searcher, query, doc_schema, "title")
snippet_generator = SnippetGenerator.create(
searcher, query, doc_schema, "title"
)
for (score, doc_address) in result.hits:
for score, doc_address in result.hits:
doc = searcher.doc(doc_address)
snippet = snippet_generator.snippet_from_doc(doc)
highlights = snippet.highlighted()
@ -859,4 +924,4 @@ class TestSnippets(object):
assert first.start == 20
assert first.end == 23
html_snippet = snippet.to_html()
assert html_snippet == 'The Old Man and the <b>Sea</b>'
assert html_snippet == "The Old Man and the <b>Sea</b>"