use serde::{Deserialize, Serialize};
use crate::core::{ApplyDefaultLocator, HasLocator};
use crate::{DistanceMetric, FilterExpr, StorageLocator};
#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
pub enum Projection {
#[default]
All,
Columns(Vec<String>),
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct QueryScope {
pub locator: StorageLocator,
pub filter: Option<FilterExpr>,
pub projection: Projection,
}
impl QueryScope {
#[must_use]
pub fn new(locator: StorageLocator) -> Self {
Self {
locator,
filter: None,
projection: Projection::All,
}
}
#[must_use]
pub fn filter(mut self, filter: FilterExpr) -> Self {
self.filter = Some(filter);
self
}
#[must_use]
pub fn projection(mut self, projection: Projection) -> Self {
self.projection = projection;
self
}
}
impl From<StorageLocator> for QueryScope {
fn from(locator: StorageLocator) -> Self {
Self::new(locator)
}
}
impl HasLocator for QueryScope {
fn locator_mut(&mut self) -> &mut StorageLocator {
&mut self.locator
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct ScopedQuery<Options> {
pub scope: QueryScope,
pub options: Options,
}
impl<Options> ScopedQuery<Options> {
#[must_use]
pub fn new(scope: impl Into<QueryScope>, options: Options) -> Self {
Self {
scope: scope.into(),
options,
}
}
pub(crate) fn collection(&self) -> &str {
&self.scope.locator.collection
}
#[must_use]
pub fn with_scope(mut self, scope: QueryScope) -> Self {
self.scope = scope;
self
}
#[must_use]
pub fn filter(mut self, filter: FilterExpr) -> Self {
self.scope = self.scope.filter(filter);
self
}
#[must_use]
pub fn projection(mut self, projection: Projection) -> Self {
self.scope = self.scope.projection(projection);
self
}
#[must_use]
pub fn map_options(mut self, update: impl FnOnce(&mut Options)) -> Self {
update(&mut self.options);
self
}
}
impl<Options> ApplyDefaultLocator for ScopedQuery<Options> {
fn apply_default_locator(mut self, default: &StorageLocator) -> Self {
self.scope = self.scope.apply_default_locator(default);
self
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct VectorSearchParams {
pub limit: usize,
pub offset: Option<usize>,
pub distance: Option<DistanceMetric>,
pub ef: Option<usize>,
}
impl Default for VectorSearchParams {
fn default() -> Self {
Self {
limit: 10,
offset: None,
distance: None,
ef: None,
}
}
}
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
pub struct VectorQueryOptions {
pub vector_field: Option<String>,
pub params: VectorSearchParams,
}
impl VectorQueryOptions {
pub fn set_field(&mut self, field: impl Into<String>) {
self.vector_field = Some(field.into());
}
pub fn set_limit(&mut self, limit: usize) {
self.params.limit = limit;
}
pub fn set_params(&mut self, params: VectorSearchParams) {
self.params = params;
}
}
impl ScopedQuery<VectorQueryOptions> {
#[must_use]
pub fn vector(locator: StorageLocator) -> Self {
Self::new(locator, VectorQueryOptions::default())
}
#[must_use]
pub fn vector_field(self, field: impl Into<String>) -> Self {
self.map_options(|options| options.set_field(field))
}
#[must_use]
pub fn limit(self, limit: usize) -> Self {
self.map_options(|options| options.set_limit(limit))
}
#[must_use]
pub fn params(self, params: VectorSearchParams) -> Self {
self.map_options(|options| options.set_params(params))
}
}
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
pub struct TextQueryOptions {
pub text_field: Option<String>,
pub limit: Option<usize>,
pub offset: Option<usize>,
}
impl TextQueryOptions {
pub fn set_field(&mut self, field: impl Into<String>) {
self.text_field = Some(field.into());
}
pub fn set_limit(&mut self, limit: usize) {
self.limit = Some(limit);
}
pub fn set_offset(&mut self, offset: usize) {
self.offset = Some(offset);
}
}
impl ScopedQuery<TextQueryOptions> {
#[must_use]
pub fn text(locator: StorageLocator) -> Self {
Self::new(locator, TextQueryOptions::default())
}
#[must_use]
pub fn text_field(self, field: impl Into<String>) -> Self {
self.map_options(|options| options.set_field(field))
}
#[must_use]
pub fn limit(self, limit: usize) -> Self {
self.map_options(|options| options.set_limit(limit))
}
#[must_use]
pub fn offset(self, offset: usize) -> Self {
self.map_options(|options| options.set_offset(offset))
}
}