refactor(search): move scoring logic to native Rust module #8492
Conversation
Summary of ChangesHello @werererer, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the application's search capabilities by migrating a performance-critical scoring function to a native Rust module. This strategic move addresses existing performance bottlenecks, particularly for users with extensive note collections, and lays the groundwork for future high-performance features by integrating Rust into the project's architecture. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request is a great initiative to improve search performance by moving the scoring logic to a native Rust module. The overall structure and integration with Node.js via NAPI-RS are well done. My review focuses on the new Rust implementation, where I've found a couple of critical correctness issues and some opportunities for further performance optimization that align with the goals of this PR. Addressing these will ensure the new native module is both correct and as fast as possible.
| fn normalize(s: &str) -> String { | ||
| s.to_lowercase() | ||
| .replace(|c: char| !c.is_alphanumeric() && c != ' ', "") | ||
| } |
There was a problem hiding this comment.
The normalize function in Rust is not consistent with the normalizeSearchText function on the JavaScript side. The JS function removes diacritics (e.g., café -> cafe), but this Rust version does not. This will lead to incorrect search scoring, as exact title matches will fail for titles with diacritics.
To fix this, you should use a Unicode normalization crate. First, add unicode-normalization = "0.1" to your [dependencies] in Cargo.toml. Then, update the normalize function to correctly handle diacritics. You'll also need to add use unicode_normalization::UnicodeNormalization; at the top of this file.
fn normalize(s: &str) -> String {
use unicode_normalization::UnicodeNormalization;
s.to_lowercase()
.nfd()
.filter(|c| !matches!(c, '\u{0300}'..='\u{036f}'))
.collect()
}| fn edit_distance(a: &str, b: &str, max: usize) -> usize { | ||
| let mut costs: Vec<usize> = (0..=b.len()).collect(); | ||
|
|
||
| for (i, ca) in a.chars().enumerate() { | ||
| let mut last = i; | ||
| costs[0] = i + 1; | ||
|
|
||
| for (j, cb) in b.chars().enumerate() { | ||
| let new = if ca == cb { | ||
| last | ||
| } else { | ||
| 1 + last.min(costs[j]).min(costs[j + 1]) | ||
| }; | ||
|
|
||
| last = costs[j + 1]; | ||
| costs[j + 1] = new; | ||
| } | ||
|
|
||
| if *costs.iter().min().unwrap() > max { | ||
| return max + 1; | ||
| } | ||
| } | ||
|
|
||
| costs[b.len()] | ||
| } |
There was a problem hiding this comment.
The implementation of edit_distance appears to be incorrect. The logic for updating costs in the single-row optimization is flawed and will not calculate the Levenshtein distance correctly. This will affect all fuzzy matching scores.
Additionally, the early-exit check *costs.iter().min().unwrap() > max is inefficient as it iterates over the costs vector in every outer loop.
Here is a corrected and optimized implementation based on the standard two-row approach, which is easier to verify for correctness and includes an efficient early-exit.
fn edit_distance(a: &str, b: &str, max: usize) -> usize {
if (a.chars().count() as i32 - b.chars().count() as i32).abs() as usize > max {
return max + 1;
}
let mut v0: Vec<usize> = (0..=b.chars().count()).collect();
let mut v1 = vec![0; b.chars().count() + 1];
for (i, ca) in a.chars().enumerate() {
v1[0] = i + 1;
let mut min_in_row = i + 1;
for (j, cb) in b.chars().enumerate() {
let cost = if ca == cb { 0 } else { 1 };
v1[j + 1] = (v0[j + 1] + 1).min(v1[j] + 1).min(v0[j] + cost);
min_in_row = min_in_row.min(v1[j + 1]);
}
if min_in_row > max {
return max + 1;
}
std::mem::swap(&mut v0, &mut v1);
}
let result = v0[b.chars().count()];
if result > max { max + 1 } else { result }
}| const normalizedQuery = normalizeSearchText( | ||
| fulltextQuery.toLowerCase() | ||
| ); | ||
|
|
||
| const score = computeScore( | ||
| { | ||
| query: fulltextQuery.toLowerCase(), | ||
| tokens, | ||
| normalizedQuery, |
There was a problem hiding this comment.
The fulltextQuery is converted to lower case twice. It's more efficient and cleaner to do this once and reuse the resulting variable.
const lowercasedQuery = fulltextQuery.toLowerCase();
const normalizedQuery = normalizeSearchText(
lowercasedQuery
);
const score = computeScore(
{
query: lowercasedQuery,
tokens,
normalizedQuery,| fn word_match(text: &str, query: &str) -> bool { | ||
| text.contains(&format!(" {query} ")) | ||
| || text.starts_with(&format!("{query} ")) | ||
| || text.ends_with(&format!(" {query}")) | ||
| } |
There was a problem hiding this comment.
This function allocates three new strings on every call due to format!. In a performance-critical path like search scoring, these allocations can add up.
Consider a more allocation-friendly approach. For example, you could build the patterns once and reuse them, or use more advanced string searching techniques that work with slices.
|
|
||
| for chunk in &chunks { | ||
| for token in tokens { | ||
| let norm_token = normalize(token); |
There was a problem hiding this comment.
This normalize(token) call is inside a nested loop. This is inefficient as each token from the search query will be normalized repeatedly for every word (chunk) in the text being scored.
For better performance, you should normalize all tokens once before the loops begin, store them in a new Vec, and then use the pre-normalized tokens inside the loop.
What's the performance improvement with the Rust implementation? |
|
The current implementation didn't improve performance, likely because it is implemented very inefficiently because of the serilization cost or because it didn't covere the most contributing factor to search performance. I will need to move more of scoring into rust to give a more fair evaluation. This PR won't go out of Draft until I see a notable performance increase. |
|
I will try to find another solution based on worker Threads instead This will take a while to cook up. |
2 participants
I maintain a collection with over 250k notes on relatively weak hardware, which makes search and autocomplete noticeably slow. (Up to 4 seconds per query)
This PR introduces Rust modules by replacing one small performance-critical function with a Rust implementation. The goal is to establish an initial integration point for gradually moving hot paths into Rust in follow-up work, in order to better support very large personal knowledge graphs or future workloads involving ontologies or LLM-generated graphs. Rust provides a foundation for exploring future optimizations such as reduced allocation overhead and parallel execution for long-running tasks.
This work is inspired by the previous effort here and may supersede it:
#7222
Open Question
Milestone: Official mobile application #7447
TODO