Added 2 attributes in partialFilterExpression to optimize IndexScan - #370
Added 2 attributes in partialFilterExpression to optimize IndexScan#370RaghuveeRR07 wants to merge 2 commits into
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughSummary by CodeRabbit
WalkthroughChanged MongoDB indexes in state-manager/app/models/db/state.py: removed the Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related issues
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @RaghuveeRR07, 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 optimizes database indexing by modifying an existing index and removing another. The "uniq_fanout_retry" index is now more efficient by applying only to a specific subset of documents, which should improve query performance for relevant operations. The removal of "run_id_status_index" further streamlines the database schema.
Highlights
- Index Modification: The "uniq_fanout_retry" index in "state.py" has been updated to include a "partialFilterExpression".
- Index Filtering Conditions: This partial filter ensures the index only applies to documents where "fanout_id" exists and "status" is one of "CREATED", "QUEUED", or "EXECUTED", optimizing its use.
- Index Removal: The "run_id_status_index" has been removed, streamlining the index definitions.
Using Gemini Code Assist
The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
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
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request optimizes a MongoDB index by adding a partialFilterExpression, which is a good approach to reduce index size. However, the change also removes the run_id_status_index, which raises a performance concern for queries that might rely on it. My review includes feedback on this potential issue, along with suggestions to improve code style and maintainability by removing unnecessary whitespace and replacing hardcoded strings with enum members.
| "fanout_id": { "$exists": True }, | ||
| "status": { "$in": ["CREATED", "QUEUED", "EXECUTED"] } | ||
| } | ||
| ) |
There was a problem hiding this comment.
Removing the run_id_status_index could cause significant performance degradation for queries that filter by run_id and status. These queries would no longer be indexed and might result in slow collection scans. If this index is still in use, it should be retained or replaced with a suitable alternative.
| import time | ||
| import uuid | ||
|
|
||
There was a problem hiding this comment.
This line contains unnecessary whitespace. It should be a blank line without any indentation to conform with PEP 8 style guidelines.1
Style Guide References
Footnotes
-
PEP 8 advises against extraneous whitespace, including on blank lines. It also recommends blank lines to separate logical sections, like imports from class definitions. ↩
| name="uniq_fanout_retry", | ||
| partialFilterExpression={ | ||
| "fanout_id": { "$exists": True }, | ||
| "status": { "$in": ["CREATED", "QUEUED", "EXECUTED"] } |
There was a problem hiding this comment.
To improve maintainability and prevent potential bugs, it's better to use the StateStatusEnum members directly instead of hardcoding string values for the status. This ensures that if the enum values change, this index definition will be updated automatically.
| "status": { "$in": ["CREATED", "QUEUED", "EXECUTED"] } | |
| "status": { "$in": [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED] } |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (2)
state-manager/app/models/db/state.py (2)
12-12: Trim stray indentation on blank line.Blank line has spaces; drop them to satisfy linters.
95-99: Re-check impact of removing run_id_status_index on query patterns.If you have hot paths filtering by run_id and status, removing that index can regress latency. Consider retaining a supporting index (possibly partial on active statuses).
To assess usage across the repo:
#!/bin/bash # Find code filtering by both run_id and status (string and ODM forms) rg -nP -C2 '(run_id).{0,120}(status)|(status).{0,120}(run_id)' rg -nP -C2 'State\.(run_id|status)\b'If these are common and latency-sensitive, add:
IndexModel([("run_id", 1), ("status", 1)], name="run_id_status_index_active", partialFilterExpression={"status": {"$in": ["CREATED","QUEUED","EXECUTED"]}})
| name="uniq_fanout_retry", | ||
| partialFilterExpression={ | ||
| "fanout_id": { "$exists": True }, | ||
| "status": { "$in": [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED] } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Changing a unique index to a partial unique index requires a migration (IndexOptionsConflict otherwise).
MongoDB won’t mutate index options in-place. Reusing the old name “uniq_fanout_retry” will fail until the old index is dropped. Prefer a two-step migration or a temporary new name.
Option A (rename, then remove old):
- name="uniq_fanout_retry",
+ name="uniq_fanout_retry_active",Then, in a migration: create the new index, validate, drop the old one.
Sample mongosh steps to run in prod/Stage:
// 1) Create new partial unique index
db.state.createIndex(
{ node_name:1, namespace_name:1, graph_name:1, identifier:1, run_id:1, retry_count:1, fanout_id:1 },
{
name: "uniq_fanout_retry_active",
unique: true,
partialFilterExpression: {
fanout_id: { $exists: true },
status: { $in: ["CREATED","QUEUED","EXECUTED"] }
}
}
)
// 2) Verify no dupes and query plans look good
// 3) Drop old index
db.state.dropIndex("uniq_fanout_retry")If you prefer to keep the original name, drop the old index first during maintenance, then recreate with the new options.
🧹 Nitpick (assertive)
Guard for duplicate-key on status transitions into the partial set.
Duplicates outside the partial set may surface as E11000 when a doc’s status changes into {CREATED, QUEUED, EXECUTED}. Ensure writers handle this with idempotent upserts/retries.
If needed, I can draft a small retry utility that treats E11000 as a signal to re-read the canonical doc.
| partialFilterExpression={ | ||
| "fanout_id": { "$exists": True }, | ||
| "status": { "$in": [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED] } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Confirm need for fanout_id: {$exists: true}.
Model sets fanout_id with a default UUID (Line 29), so “exists” may be redundant unless legacy docs lack it. If all docs have fanout_id, removing this predicate will enlarge index coverage and simplify the filter.
🤖 Prompt for AI Agents
In state-manager/app/models/db/state.py around lines 96 to 99, the
partialFilterExpression currently requires "fanout_id": { "$exists": True }
which may be redundant because fanout_id is set with a default UUID at model
definition (line 29); verify whether any legacy documents exist without
fanout_id (query the collection for documents where fanout_id is null/absent),
and if none are found remove the "$exists": True clause from the
partialFilterExpression so the index covers more documents; after changing,
rebuild the index and run relevant unit/integration tests or migration scripts
to ensure no regression.
💡 Verification agent
🧩 Analysis chain
Ensure Enum values are BSON-serializable strings (or use .value explicitly).
If StateStatusEnum isn’t a str Enum, PyMongo will reject Enum members in index options. Using .value is safe and explicit.
Apply:
- partialFilterExpression={
- "fanout_id": { "$exists": True },
- "status": { "$in": [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED] }
- }
+ partialFilterExpression={
+ "fanout_id": {"$exists": True},
+ "status": {"$in": [StateStatusEnum.CREATED.value, StateStatusEnum.QUEUED.value, StateStatusEnum.EXECUTED.value]},
+ }Verification script to confirm the enum type:
🏁 Script executed:
#!/bin/bash
# Expect: StateStatusEnum subclasses `str, Enum` or similar.
rg -nP -C2 'class\s+StateStatusEnum\s*\('Length of output: 334
StateStatusEnum subclasses str, so its members are BSON-serializable; using .value is purely optional for clarity.
🤖 Prompt for AI Agents
In state-manager/app/models/db/state.py around lines 96 to 99, the
partialFilterExpression currently uses StateStatusEnum members with .value;
since StateStatusEnum subclasses str and its members are BSON-serializable,
remove the unnecessary .value calls and use the enum members directly (e.g.,
StateStatusEnum.CREATED) for clarity and consistency, ensuring the resulting
partialFilterExpression still contains the intended string values.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
| name="run_id_status_index" | ||
| name="uniq_fanout_retry", | ||
| partialFilterExpression={ | ||
| "fanout_id": { "$exists": True }, |
There was a problem hiding this comment.
@RaghuveeRR07 why this? if you see the model fanout_id always exists
| name="uniq_fanout_retry", | ||
| partialFilterExpression={ | ||
| "fanout_id": { "$exists": True }, | ||
| "status": { "$in": [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED] } |
There was a problem hiding this comment.
why only these CREATED, QUEUED, EXECUTED status?
|
Hey @RaghuveeRR07 wanted to check if you got a chance to improve this implementation as per our last discussion? |
|
Closing this PR for now. @RaghuveeRR07 please re-open this as you complete fix for this. |
2 participants
Index now applies only to documents where: