-
Notifications
You must be signed in to change notification settings - Fork 1.6k
[ENH] Add support to convert Get/QueryResult to pandas dataframe #4304
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Reviewer ChecklistPlease leverage this checklist to ensure your code review is thorough before approving Testing, Bugs, Errors, Logs, Documentation
System Compatibility
Quality
|
This stack of pull requests is managed by Graphite. Learn more about stacking. |
003dfae
to
ba51cf1
Compare
ba51cf1
to
b486c1b
Compare
b486c1b
to
5b56585
Compare
chromadb/api/types.py
Outdated
@@ -419,6 +419,10 @@ class QueryResult(TypedDict): | |||
included: Include | |||
|
|||
|
|||
# Type alias for union of result types | |||
GetOrQueryResult = Union["QueryResult", "GetResult"] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we should do two different types here. Commingling them makes it more complicated than necessary.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
separated, now theres 2 functions: query_result_to_df and get_result_to_df
5b56585
to
a42164b
Compare
chromadb/utils/results.py
Outdated
|
||
data_for_df["id"] = query_result["ids"][i] | ||
|
||
if "documents" in included_fields and query_result["documents"] is not None: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we generalize this behavior into a helper? Its all the same
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
updated
4647db9
to
3b3363d
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- Please add tests
- Pandas is not a explicit dep of chroma. Maybe we should error if its not installed.
97ed1a2
to
b705d6d
Compare
b705d6d
to
a5eb712
Compare
Description of changes
This PR adds support to convert both get and query results to pandas dataframes. As a byproduct, users can now get results in row format.
Test plan
How are these changes tested?
pytest
for python,yarn test
for js,cargo test
for rustDocumentation Changes
Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs repository?