Initial pile-dataset
This commit is contained in:
304
crates/pile-dataset/src/index/index_fts.rs
Normal file
304
crates/pile-dataset/src/index/index_fts.rs
Normal file
@@ -0,0 +1,304 @@
|
||||
use jsonpath_rust::JsonPath;
|
||||
use pile_config::{ConfigToml, DatasetFts};
|
||||
use serde_json::Value;
|
||||
use std::{path::PathBuf, sync::LazyLock};
|
||||
use tantivy::{
|
||||
DocAddress, Index, ReloadPolicy, TantivyDocument, TantivyError,
|
||||
collector::Collector,
|
||||
query::QueryParser,
|
||||
schema::{self, Schema, Value as TantivyValue},
|
||||
};
|
||||
use tracing::{debug, trace, warn};
|
||||
|
||||
use crate::{Item, Key};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FtsLookupResult {
|
||||
pub score: f32,
|
||||
pub source_name: String,
|
||||
pub key: String,
|
||||
}
|
||||
|
||||
pub struct DbFtsIndex {
|
||||
pub path: PathBuf,
|
||||
pub cfg: ConfigToml,
|
||||
pub schema: Schema,
|
||||
}
|
||||
|
||||
impl DbFtsIndex {
|
||||
fn fts_cfg(&self) -> &DatasetFts {
|
||||
static DEFAULT: LazyLock<DatasetFts> = LazyLock::new(|| DatasetFts::default());
|
||||
&self.cfg.fts.as_ref().unwrap_or(&DEFAULT)
|
||||
}
|
||||
}
|
||||
|
||||
impl DbFtsIndex {
|
||||
pub fn new(path: impl Into<PathBuf>, cfg: &ConfigToml) -> Self {
|
||||
let path = path.into();
|
||||
let cfg = cfg.clone();
|
||||
|
||||
let schema = {
|
||||
let mut schema_builder = Schema::builder();
|
||||
|
||||
schema_builder.add_text_field("_meta_source", schema::STORED);
|
||||
schema_builder.add_text_field("_meta_key", schema::STORED);
|
||||
|
||||
let default = DatasetFts::default();
|
||||
let fields = &cfg.fts.as_ref().unwrap_or(&default).fields;
|
||||
for (name, field) in fields {
|
||||
if field.tokenize {
|
||||
schema_builder.add_text_field(&name, schema::TEXT);
|
||||
} else {
|
||||
schema_builder.add_text_field(&name, schema::STRING);
|
||||
}
|
||||
}
|
||||
|
||||
schema_builder.build()
|
||||
};
|
||||
|
||||
Self { path, cfg, schema }
|
||||
}
|
||||
|
||||
//
|
||||
// MARK: write
|
||||
//
|
||||
|
||||
/// Turn an entry into a tantivy document
|
||||
pub fn entry_to_document<K: Key>(
|
||||
&self,
|
||||
item: &dyn Item<Key = K>,
|
||||
) -> Result<Option<TantivyDocument>, TantivyError> {
|
||||
let mut doc = TantivyDocument::default();
|
||||
|
||||
{
|
||||
let f_ptr = self.schema.get_field("_meta_source")?;
|
||||
doc.add_text(f_ptr, item.source_name());
|
||||
}
|
||||
|
||||
{
|
||||
let f_ptr = self.schema.get_field("_meta_key")?;
|
||||
doc.add_text(f_ptr, item.key().to_string());
|
||||
}
|
||||
|
||||
let json = item.json().unwrap();
|
||||
let mut empty = true;
|
||||
for (name, _field) in &self.fts_cfg().fields {
|
||||
let val = match self.get_field(&json, name)? {
|
||||
Some(x) => x,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
empty = false;
|
||||
|
||||
let field = self.schema.get_field(name);
|
||||
if let Ok(field) = field {
|
||||
doc.add_text(field, val);
|
||||
}
|
||||
}
|
||||
|
||||
return Ok((!empty).then_some(doc));
|
||||
}
|
||||
|
||||
//
|
||||
// MARK: read
|
||||
//
|
||||
|
||||
pub fn get_field(
|
||||
&self,
|
||||
json: &Value,
|
||||
field_name: &str,
|
||||
) -> Result<Option<String>, std::io::Error> {
|
||||
let field = match self.cfg.schema.get(field_name) {
|
||||
Some(x) => x,
|
||||
None => {
|
||||
warn!("Unknown field {field_name:?}");
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
// Try paths in order, using the first value we find
|
||||
'outer: for path in field.path.as_slice() {
|
||||
let val = match json.query(&path) {
|
||||
Ok(mut x) => {
|
||||
if x.len() > 1 {
|
||||
warn!(
|
||||
message = "Path returned more than one value, this is not supported. Skipping.",
|
||||
?path,
|
||||
field = field_name
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
match x.pop() {
|
||||
Some(x) => x,
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
|
||||
Err(error) => {
|
||||
warn!(
|
||||
message = "Invalid path, skipping",
|
||||
?path,
|
||||
field = field_name,
|
||||
?error
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let mut val = match val {
|
||||
Value::Null => {
|
||||
trace!(
|
||||
message = "Skipping field, is null",
|
||||
field = field_name,
|
||||
path,
|
||||
value = ?val
|
||||
);
|
||||
continue;
|
||||
}
|
||||
x => x.clone(),
|
||||
};
|
||||
|
||||
for post in &field.post {
|
||||
val = match post.apply(&val) {
|
||||
Some(x) => x,
|
||||
None => return Ok(None),
|
||||
};
|
||||
}
|
||||
|
||||
loop {
|
||||
val = match val {
|
||||
Value::Array(ref mut x) => {
|
||||
if x.len() == 1 {
|
||||
x.pop().unwrap()
|
||||
} else if x.len() > 1 {
|
||||
debug!(
|
||||
message = "Skipping field, is array with more than one element",
|
||||
field = field_name,
|
||||
path,
|
||||
value = ?val
|
||||
);
|
||||
continue 'outer;
|
||||
} else {
|
||||
debug!(
|
||||
message = "Skipping field, is empty array",
|
||||
field = field_name,
|
||||
path,
|
||||
value = ?val
|
||||
);
|
||||
continue 'outer;
|
||||
}
|
||||
}
|
||||
Value::Null => {
|
||||
trace!(
|
||||
message = "Skipping field, is null",
|
||||
field = field_name,
|
||||
path,
|
||||
value = ?val
|
||||
);
|
||||
continue 'outer;
|
||||
}
|
||||
Value::Object(_) => {
|
||||
trace!(
|
||||
message = "Skipping field, is object",
|
||||
field = field_name,
|
||||
path,
|
||||
value = ?val
|
||||
);
|
||||
continue 'outer;
|
||||
}
|
||||
Value::Bool(x) => return Ok(Some(x.to_string())),
|
||||
Value::Number(x) => return Ok(Some(x.to_string())),
|
||||
Value::String(x) => return Ok(Some(x)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
/// Run the given query on this table's ftx index.
|
||||
/// Returns `None` if this table does not have an fts index.
|
||||
///
|
||||
/// Note that concurrent writes to the table may invalidate `Entry::Id`.
|
||||
/// When you use [Table::get_meta] on this method's results, there is
|
||||
/// NO GUARANTEE that the entries returned by this method will exist.
|
||||
pub fn lookup<C>(
|
||||
&self,
|
||||
query: impl Into<String>,
|
||||
collector: impl AsRef<C> + Send + 'static,
|
||||
) -> Result<Vec<FtsLookupResult>, TantivyError>
|
||||
where
|
||||
C: Collector,
|
||||
C::Fruit: IntoIterator<Item = (f32, DocAddress)>,
|
||||
{
|
||||
if !self.path.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
if !self.path.is_dir() {
|
||||
warn!("fts index at {} is not a directory?!", self.path.display());
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let query: String = query.into();
|
||||
let index = Index::open_in_dir(&self.path)?;
|
||||
let reader = index
|
||||
.reader_builder()
|
||||
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
||||
.try_into()?;
|
||||
|
||||
let schema = index.schema();
|
||||
|
||||
let mut fields = Vec::new();
|
||||
for f in self.schema.fields() {
|
||||
fields.push(f.0)
|
||||
}
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let query_parser = QueryParser::for_index(&index, fields);
|
||||
let query = query_parser.parse_query(&query)?;
|
||||
let res = searcher.search(&query, collector.as_ref())?;
|
||||
|
||||
let mut out = Vec::new();
|
||||
for (score, doc) in res {
|
||||
let retrieved_doc: TantivyDocument = searcher.doc(doc)?;
|
||||
|
||||
#[expect(clippy::unwrap_used)]
|
||||
let source_name = {
|
||||
let f_source = schema.get_field("_meta_source").unwrap();
|
||||
|
||||
let mut source = None;
|
||||
for v in retrieved_doc.get_all(f_source) {
|
||||
assert!(source.is_none()); // Must only exist once
|
||||
let v = v.as_str().unwrap().to_owned();
|
||||
source = Some(v)
|
||||
}
|
||||
|
||||
source.unwrap()
|
||||
};
|
||||
|
||||
#[expect(clippy::unwrap_used)]
|
||||
let key = {
|
||||
let f_source = schema.get_field("_meta_key").unwrap();
|
||||
|
||||
let mut source = None;
|
||||
for v in retrieved_doc.get_all(f_source) {
|
||||
assert!(source.is_none()); // Must only exist once
|
||||
let v = v.as_str().unwrap().to_owned();
|
||||
source = Some(v)
|
||||
}
|
||||
|
||||
source.unwrap()
|
||||
};
|
||||
|
||||
out.push(FtsLookupResult {
|
||||
score,
|
||||
source_name,
|
||||
key,
|
||||
});
|
||||
}
|
||||
|
||||
return Ok(out);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user