Skip to content

Commit

Permalink
Fine-tune handling of GIL for search (#124)
Browse files Browse the repository at this point in the history
  • Loading branch information
adamreichold authored Sep 11, 2023
1 parent 91a422b commit 76512f8
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 57 deletions.
2 changes: 1 addition & 1 deletion src/document.rs
Original file line number Diff line number Diff line change
Expand Up @@ -740,7 +740,7 @@ impl Document {
{
self.field_values
.entry(field_name)
.or_insert_with(Vec::new)
.or_default()
.push(Value::from(value));
}

Expand Down
16 changes: 5 additions & 11 deletions src/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -308,19 +308,13 @@ impl Index {
Ok(())
}

/// Acquires a Searcher from the searcher pool.
/// Returns a searcher
///
/// If no searcher is available during the call, note that
/// this call will block until one is made available.
///
/// Searcher are automatically released back into the pool when
/// they are dropped. If you observe this function to block forever
/// you probably should configure the Index to have a larger
/// searcher pool, or you are holding references to previous searcher
/// for ever.
fn searcher(&self, py: Python) -> Searcher {
/// This method should be called every single time a search query is performed.
/// The same searcher must be used for a given query, as it ensures the use of a consistent segment set.
fn searcher(&self) -> Searcher {
Searcher {
inner: py.allow_threads(|| self.reader.searcher()),
inner: self.reader.searcher(),
}
}

Expand Down
99 changes: 54 additions & 45 deletions src/searcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -133,66 +133,75 @@ impl Searcher {
#[pyo3(signature = (query, limit = 10, count = true, order_by_field = None, offset = 0))]
fn search(
&self,
_py: Python,
py: Python,
query: &Query,
limit: usize,
count: bool,
order_by_field: Option<&str>,
offset: usize,
) -> PyResult<SearchResult> {
let mut multicollector = MultiCollector::new();
py.allow_threads(move || {
let mut multicollector = MultiCollector::new();

let count_handle = if count {
Some(multicollector.add_collector(Count))
} else {
None
};
let count_handle = if count {
Some(multicollector.add_collector(Count))
} else {
None
};

let (mut multifruit, hits) = {
if let Some(order_by) = order_by_field {
let collector = TopDocs::with_limit(limit)
.and_offset(offset)
.order_by_u64_field(order_by);
let top_docs_handle = multicollector.add_collector(collector);
let ret = self.inner.search(query.get(), &multicollector);
let (mut multifruit, hits) = {
if let Some(order_by) = order_by_field {
let collector = TopDocs::with_limit(limit)
.and_offset(offset)
.order_by_u64_field(order_by);
let top_docs_handle =
multicollector.add_collector(collector);
let ret = self.inner.search(query.get(), &multicollector);

match ret {
Ok(mut r) => {
let top_docs = top_docs_handle.extract(&mut r);
let result: Vec<(Fruit, DocAddress)> = top_docs
.iter()
.map(|(f, d)| {
(Fruit::Order(*f), DocAddress::from(d))
})
.collect();
(r, result)
match ret {
Ok(mut r) => {
let top_docs = top_docs_handle.extract(&mut r);
let result: Vec<(Fruit, DocAddress)> = top_docs
.iter()
.map(|(f, d)| {
(Fruit::Order(*f), DocAddress::from(d))
})
.collect();
(r, result)
}
Err(e) => {
return Err(PyValueError::new_err(e.to_string()))
}
}
Err(e) => return Err(PyValueError::new_err(e.to_string())),
}
} else {
let collector = TopDocs::with_limit(limit).and_offset(offset);
let top_docs_handle = multicollector.add_collector(collector);
let ret = self.inner.search(query.get(), &multicollector);
} else {
let collector =
TopDocs::with_limit(limit).and_offset(offset);
let top_docs_handle =
multicollector.add_collector(collector);
let ret = self.inner.search(query.get(), &multicollector);

match ret {
Ok(mut r) => {
let top_docs = top_docs_handle.extract(&mut r);
let result: Vec<(Fruit, DocAddress)> = top_docs
.iter()
.map(|(f, d)| {
(Fruit::Score(*f), DocAddress::from(d))
})
.collect();
(r, result)
match ret {
Ok(mut r) => {
let top_docs = top_docs_handle.extract(&mut r);
let result: Vec<(Fruit, DocAddress)> = top_docs
.iter()
.map(|(f, d)| {
(Fruit::Score(*f), DocAddress::from(d))
})
.collect();
(r, result)
}
Err(e) => {
return Err(PyValueError::new_err(e.to_string()))
}
}
Err(e) => return Err(PyValueError::new_err(e.to_string())),
}
}
};
};

let count = count_handle.map(|h| h.extract(&mut multifruit));
let count = count_handle.map(|h| h.extract(&mut multifruit));

Ok(SearchResult { hits, count })
Ok(SearchResult { hits, count })
})
}

/// Returns the overall number of documents in the index.
Expand Down

0 comments on commit 76512f8

Please sign in to comment.