Skip to content

Commit 913cd58

Browse files
authored
[ENH] Builders for SearchPayload (#5698)
## Description of changes _Summarize the changes made by this PR._ - Improvements & Bug fixes - Add a few helper builders for `SearchPayload` - New functionality - N/A ## Test plan _How are these changes tested?_ - [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Migration plan _Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_ ## Observability plan _What is the plan to instrument and monitor this change?_ ## Documentation Changes _Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_
1 parent ab0ca09 commit 913cd58

File tree

1 file changed

+31
-2
lines changed

1 file changed

+31
-2
lines changed

rust/types/src/execution/plan.rs

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,12 @@ use super::{
44
Filter, KnnBatch, KnnProjection, Limit, Projection, Rank, Scan, ScanToProtoError, Select,
55
},
66
};
7-
use crate::{chroma_proto, validators::validate_rank};
7+
use crate::{
8+
chroma_proto,
9+
operator::{Key, RankExpr},
10+
validators::validate_rank,
11+
Where,
12+
};
813
use serde::{Deserialize, Serialize};
914
use thiserror::Error;
1015
#[cfg(feature = "utoipa")]
@@ -146,7 +151,7 @@ impl TryFrom<Knn> for chroma_proto::KnnPlan {
146151
}
147152
}
148153

149-
#[derive(Clone, Debug, Deserialize, Serialize, Validate)]
154+
#[derive(Clone, Debug, Default, Deserialize, Serialize, Validate)]
150155
pub struct SearchPayload {
151156
#[serde(default)]
152157
pub filter: Filter,
@@ -159,6 +164,30 @@ pub struct SearchPayload {
159164
pub select: Select,
160165
}
161166

167+
impl SearchPayload {
168+
pub fn limit(mut self, limit: Option<u32>, offset: u32) -> Self {
169+
self.limit.limit = limit;
170+
self.limit.offset = offset;
171+
self
172+
}
173+
pub fn rank(mut self, expr: RankExpr) -> Self {
174+
self.rank.expr = Some(expr);
175+
self
176+
}
177+
pub fn select<I, T>(mut self, keys: I) -> Self
178+
where
179+
I: IntoIterator<Item = T>,
180+
T: Into<Key>,
181+
{
182+
self.select.keys = keys.into_iter().map(Into::into).collect();
183+
self
184+
}
185+
pub fn r#where(mut self, r#where: Where) -> Self {
186+
self.filter.where_clause = Some(r#where);
187+
self
188+
}
189+
}
190+
162191
#[cfg(feature = "utoipa")]
163192
impl PartialSchema for SearchPayload {
164193
fn schema() -> RefOr<Schema> {

0 commit comments

Comments
 (0)