[ENH]: More precise calculation of hnsw cache item weight

[sanketkedia]

Description of changes

Summarize the changes made by this PR.

• Improvements & Bug fixes
  • We undercalculate the weight of hnsw cache item. This PR makes it a tad more accurate by accounting for the graph nodes and edges and also levels other than L0
• New functionality
  • ...

Test plan

How are these changes tested?

• Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Documentation Changes

None

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[sanketkedia]

• [ENH]: More precise calculation of hnsw cache item weight #4657 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[propel-code-bot]

This PR updates the weight estimation logic for HnswIndexRef cache items to more accurately account for the memory footprint, incorporating not just the L0 embeddings but also graph nodes, edges, and higher HNSW levels. The approach introduces a constant for the M parameter (node connections) and doubles the size calculation to approximate additional HNSW layers.

This summary was automatically generated by @propel-code-bot

[entelligence-ai-reviews]

Walkthrough

This update refines the memory usage estimation logic in the weight method of the HnswIndexRef implementation. The new approach accounts for both the L0 graph structure and the embeddings, using a more comprehensive calculation that includes deleted elements and multiplies the total by 2 to estimate higher graph levels. This results in more accurate cache weight calculations for HNSW indices, improving resource management.

Changes

File(s)	Summary
rust/index/src/hnsw_provider.rs	Updated the weight method for HnswIndexRef to improve memory usage estimation by considering both L0 graph and embeddings, using index.len_with_deleted() for element count, and multiplying by 2 for higher graph levels.

Sequence Diagram

This diagram shows the interactions between components:

sequenceDiagram
    title HnswIndexRef Weight Calculation Flow
    
    participant Client
    participant HnswIndexRef
    participant IndexInner
    
    Client->>HnswIndexRef: weight()
    activate HnswIndexRef
    
    HnswIndexRef->>IndexInner: read()
    activate IndexInner
    IndexInner-->>HnswIndexRef: index (RwLockReadGuard)
    deactivate IndexInner
    
    Note over HnswIndexRef: Check if index is empty
    HnswIndexRef->>IndexInner: len_with_deleted()
    activate IndexInner
    IndexInner-->>HnswIndexRef: count (including deleted items)
    deactivate IndexInner
    
    alt index is empty (len_with_deleted() == 0)
        HnswIndexRef-->>Client: return 1
    else index has elements
        HnswIndexRef->>IndexInner: len_with_deleted()
        activate IndexInner
        IndexInner-->>HnswIndexRef: count (including deleted items)
        deactivate IndexInner
        
        HnswIndexRef->>IndexInner: dimensionality()
        activate IndexInner
        IndexInner-->>HnswIndexRef: dimensions
        deactivate IndexInner
        
        Note over HnswIndexRef: Calculate memory usage
        Note over HnswIndexRef: 1. Graph bytes = M * sizeof(u32) * len_with_deleted()
        Note over HnswIndexRef: 2. Embedding bytes = len_with_deleted() * sizeof(f32) * dimensionality
        Note over HnswIndexRef: 3. Total bytes = 2 * (graph_bytes + embedding_bytes)
        Note over HnswIndexRef: 4. Convert to MB
        
        alt as_mb == 0
            HnswIndexRef-->>Client: return 1
        else as_mb > 0
            HnswIndexRef-->>Client: return as_mb
        end
    end
    
    deactivate HnswIndexRef

Loading

▶️ ⚡ AI Code Reviews for VS Code, Cursor, Windsurf
Install the extension

Note for Windsurf
Please change the default marketplace provider to the following in the windsurf settings:
Marketplace Extension Gallery Service URL: https://marketplace.visualstudio.com/_apis/public/gallery
Marketplace Gallery Item URL: https://marketplace.visualstudio.com/items

Entelligence.ai can learn from your feedback. Simply add 👍 / 👎 emojis to teach it your preferences. More shortcuts below

Emoji Descriptions:

• ⚠️ Potential Issue - May require further investigation.
• 🔒 Security Vulnerability - Fix to ensure system safety.
• 💻 Code Improvement - Suggestions to enhance code quality.
• 🔨 Refactor Suggestion - Recommendations for restructuring code.
• ℹ️ Others - General comments and information.

Interact with the Bot:

• Send a message or request using the format:
  @entelligenceai + *your message*

Example: @entelligenceai Can you suggest improvements for this code?

• Help the Bot learn by providing feedback on its responses.
  @entelligenceai + *feedback*

Example: @entelligenceai Do not comment on `save_auth` function !

Also you can trigger various commands with the bot by doing
@entelligenceai command

The current supported commands are

1. config - shows the current config
2. retrigger_review - retriggers the review

More commands to be added soon.