Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions performer/proto/operational/sdk.caps.proto
Original file line number Diff line number Diff line change
Expand Up @@ -152,4 +152,8 @@ enum Caps {

// The SDK supports the RESUMING status in the result of functions_status() of the eventing management API.
SDK_EVENTING_RESUMING_FUNCTION_STATUS = 40;

// The SDK supports score fusion for hybrid search (SDK-RFC 52).
// Support for SDK_VECTOR_SEARCH is a prerequisite.
SDK_SEARCH_SCORE_FUSION = 41;
}
27 changes: 25 additions & 2 deletions performer/proto/operational/sdk.search.proto
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,7 @@ option go_package = "github.com/couchbaselabs/transactions-fit-performer/protoco
option java_multiple_files = true;
option ruby_package = "FIT::Protocol::SDK::Search";

import "shared.doc_location.proto";
import "shared.content.proto";
import "shared.collection.proto";
import "shared.basic.proto";
import "google/protobuf/timestamp.proto";
import "streams.top_level.proto";
Expand Down Expand Up @@ -91,6 +89,27 @@ enum VectorQueryCombination {
OR = 1;
}

message SearchScoringReciprocalRankFusion {
optional uint32 rank_constant = 1;
optional uint32 window_size = 2;
}

message SearchScoringRelativeScoreFusion {
optional uint32 window_size = 1;
}

message SearchScoringNone {
}

// Controls the top-level search score field.
message SearchScoring {
oneof mode {
SearchScoringReciprocalRankFusion reciprocal_rank_fusion = 1;
SearchScoringRelativeScoreFusion relative_score_fusion = 2;
// Equivalent to SearchOptions.disable_scoring.
SearchScoringNone none = 3;
}
}

// Executing a `cluster.searchQuery()` FTS query.
//
Expand Down Expand Up @@ -321,6 +340,10 @@ message SearchOptions {
map<string, string> raw = 12;
optional bool include_locations = 13;
optional protocol.shared.JsonSerializer serialize = 14;
// Top-level search score mode. Mutually exclusive with disable_scoring.
optional SearchScoring scoring = 15;
// Maps to the SDK's deprecated disableScoring option.
optional bool disable_scoring = 16 [deprecated = true];
}

enum MatchOperator {
Expand Down
1 change: 1 addition & 0 deletions performer/src/performer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ impl PerformerService for Performer {
sdk::Caps::SdkSearch,
sdk::Caps::SdkVectorSearch,
sdk::Caps::SdkVectorSearchBase64,
sdk::Caps::SdkSearchScoreFusion,
sdk::Caps::SdkScopeSearch,
sdk::Caps::SdkSearchRfcRevision11,
sdk::Caps::SdkScopeSearchIndexManagement,
Expand Down
44 changes: 42 additions & 2 deletions performer/src/translations/search_options.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
use crate::errors::error::{Error, Result};
use crate::proto::protocol::sdk::search::search_facet::Facet;
use crate::proto::protocol::sdk::search::{
search_sort, Highlight, HighlightStyle, SearchFacet, SearchGeoDistanceUnits, SearchOptions,
SearchScanConsistency, SearchSort, VectorQueryCombination, VectorSearchOptions,
search_scoring, search_sort, Highlight, HighlightStyle, SearchFacet, SearchGeoDistanceUnits,
SearchOptions, SearchScanConsistency, SearchScoring, SearchSort, VectorQueryCombination,
VectorSearchOptions,
};
use chrono::{DateTime, FixedOffset, Utc};
use couchbase::options::search_options::ScanConsistency::NotBounded;
use couchbase::search::scoring::{ReciprocalRankScoreFusion, RelativeScoreScoreFusion, Scoring};
use couchbase::search::sort::{Sort, SortFieldMode, SortFieldType, SortGeoDistanceUnit};
use prost_types::Timestamp;
use std::time::SystemTime;
Expand Down Expand Up @@ -72,10 +74,48 @@ impl TryFrom<SearchOptions> for couchbase::options::search_options::SearchOption
if let Some(_serializer) = opts.serialize {
return Err(Error::unimplemented("serializer is unimplemented"));
}
if let Some(scoring) = opts.scoring {
copts = copts.scoring(scoring.try_into()?);
}
#[allow(deprecated)]
if let Some(disable_scoring) = opts.disable_scoring {
copts = copts.disable_scoring(disable_scoring);
}
Ok(copts)
}
}

impl TryFrom<SearchScoring> for Scoring {
type Error = Box<Error>;

fn try_from(proto: SearchScoring) -> Result<Self> {
let mode = proto
.mode
.ok_or_else(|| Error::invalid_argument("Search scoring mode is not set"))?;

match mode {
search_scoring::Mode::None(_) => Ok(Scoring::None),
search_scoring::Mode::ReciprocalRankFusion(rrf) => {
let mut fusion = ReciprocalRankScoreFusion::new();
if let Some(rank_constant) = rrf.rank_constant {
fusion = fusion.rank_constant(rank_constant);
}
if let Some(window_size) = rrf.window_size {
fusion = fusion.window_size(window_size);
}
Ok(Scoring::ReciprocalRankFusion(fusion))
}
search_scoring::Mode::RelativeScoreFusion(rsf) => {
let mut fusion = RelativeScoreScoreFusion::new();
if let Some(window_size) = rsf.window_size {
fusion = fusion.window_size(window_size);
}
Ok(Scoring::RelativeScoreFusion(fusion))
}
}
}
}

impl TryFrom<Highlight> for couchbase::options::search_options::Highlight {
type Error = Box<Error>;

Expand Down
3 changes: 3 additions & 0 deletions sdk/couchbase-core/src/componentconfigs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,9 @@ impl AgentComponentConfigs {
vector_search_enabled: config
.features
.contains(&ParsedConfigFeature::FtsVectorSearch),
score_fusion_enabled: config
.features
.contains(&ParsedConfigFeature::FtsScoreFusion),
},
mgmt_config: MgmtComponentConfig {
endpoints: mgmt_endpoints,
Expand Down
3 changes: 3 additions & 0 deletions sdk/couchbase-core/src/configparser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,9 @@ impl ConfigParser {
if caps.contains(&"vectorSearch".to_string()) {
features.push(ParsedConfigFeature::FtsVectorSearch);
}
if caps.contains(&"scoreFusion".to_string()) {
features.push(ParsedConfigFeature::FtsScoreFusion);
}
}

let cluster_labels = if config.cluster_name.is_some() || config.cluster_uuid.is_some() {
Expand Down
10 changes: 9 additions & 1 deletion sdk/couchbase-core/src/options/search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use crate::retry::{RetryStrategy, DEFAULT_RETRY_STRATEGY};
use crate::searchx;
use crate::searchx::facets::Facet;
use crate::searchx::queries::Query;
use crate::searchx::query_options::{Control, Highlight, KnnOperator, KnnQuery};
use crate::searchx::query_options::{Control, Highlight, KnnOperator, KnnQuery, Params};
use crate::searchx::sort::Sort;
use std::collections::HashMap;
use std::sync::Arc;
Expand All @@ -46,6 +46,7 @@ pub struct SearchOptions {
pub sort: Option<Vec<Sort>>,
pub knn: Option<Vec<KnnQuery>>,
pub knn_operator: Option<KnnOperator>,
pub params: Option<Params>,

pub raw: Option<HashMap<String, serde_json::Value>>,

Expand Down Expand Up @@ -79,6 +80,7 @@ impl SearchOptions {
sort: None,
knn: None,
knn_operator: None,
params: None,
raw: None,
index_name: index_name.into(),
scope_name: None,
Expand Down Expand Up @@ -174,6 +176,11 @@ impl SearchOptions {
self
}

pub fn params(mut self, params: impl Into<Option<Params>>) -> Self {
self.params = params.into();
self
}

pub fn raw(mut self, raw: impl Into<Option<HashMap<String, serde_json::Value>>>) -> Self {
self.raw = raw.into();
self
Expand Down Expand Up @@ -225,6 +232,7 @@ impl From<SearchOptions> for searchx::query_options::QueryOptions {
sort: opts.sort,
knn: opts.knn,
knn_operator: opts.knn_operator,
params: opts.params,
raw: opts.raw,
index_name: opts.index_name,
scope_name: opts.scope_name,
Expand Down
1 change: 1 addition & 0 deletions sdk/couchbase-core/src/parsedconfig.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ pub(crate) struct ParsedConfigNode {
#[derive(Debug, Clone, Eq, PartialEq)]
pub(crate) enum ParsedConfigFeature {
FtsVectorSearch,
FtsScoreFusion,
Unknown,
}

Expand Down
9 changes: 9 additions & 0 deletions sdk/couchbase-core/src/searchcomponent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,13 +67,15 @@ pub(crate) struct SearchComponent<C: Client> {
#[derive(Debug)]
pub(crate) struct SearchComponentState {
pub vector_search_enabled: bool,
pub score_fusion_enabled: bool,
}

pub(crate) struct SearchComponentConfig {
pub endpoints: HashMap<String, NetworkAndCanonicalEndpoint>,
pub authenticator: Authenticator,

pub vector_search_enabled: bool,
pub score_fusion_enabled: bool,
}

#[derive(Debug)]
Expand Down Expand Up @@ -102,6 +104,7 @@ impl<C: Client + 'static> SearchComponent<C> {
retry_manager,
state: ArcSwap::new(Arc::new(SearchComponentState {
vector_search_enabled: config.vector_search_enabled,
score_fusion_enabled: config.score_fusion_enabled,
})),
}
}
Expand All @@ -120,6 +123,7 @@ impl<C: Client + 'static> SearchComponent<C> {

self.state.swap(Arc::new(SearchComponentState {
vector_search_enabled: config.vector_search_enabled,
score_fusion_enabled: config.score_fusion_enabled,
}));
}

Expand Down Expand Up @@ -157,6 +161,7 @@ impl<C: Client + 'static> SearchComponent<C> {
auth,

vector_search_enabled: self.state.load().vector_search_enabled,
score_fusion_enabled: self.state.load().score_fusion_enabled,
tracing: self.tracing.clone(),
}
.query(&copts)
Expand Down Expand Up @@ -475,6 +480,7 @@ impl<C: Client + 'static> SearchComponent<C> {
canonical_endpoint: target.canonical_endpoint,
auth: target.auth,
vector_search_enabled: false,
score_fusion_enabled: false,
tracing: self.tracing.clone(),
};

Expand Down Expand Up @@ -511,6 +517,7 @@ impl<C: Client + 'static> SearchComponent<C> {
auth: target.auth,

vector_search_enabled: self.state.load().vector_search_enabled,
score_fusion_enabled: self.state.load().score_fusion_enabled,
tracing: self.tracing.clone(),
};

Expand Down Expand Up @@ -598,6 +605,7 @@ impl<C: Client + 'static> SearchComponent<C> {
auth,

vector_search_enabled: self.state.load().vector_search_enabled,
score_fusion_enabled: self.state.load().score_fusion_enabled,
tracing: self.tracing.clone(),
})
.await
Expand Down Expand Up @@ -641,6 +649,7 @@ impl<C: Client + 'static> SearchComponent<C> {
auth,

vector_search_enabled: self.state.load().vector_search_enabled,
score_fusion_enabled: self.state.load().score_fusion_enabled,
tracing: self.tracing.clone(),
})
.await
Expand Down
2 changes: 2 additions & 0 deletions sdk/couchbase-core/src/searchx/ensure_index_helper.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ impl<'a> EnsureIndexHelper<'a> {
canonical_endpoint: target.canonical_endpoint.to_string(),
auth: target.auth.clone(),
vector_search_enabled: true,
score_fusion_enabled: true,
tracing: Default::default(),
}
.get_index(&GetIndexOptions {
Expand Down Expand Up @@ -126,6 +127,7 @@ impl<'a> EnsureIndexHelper<'a> {
canonical_endpoint: target.canonical_endpoint.to_string(),
auth: target.auth.clone(),
vector_search_enabled: true,
score_fusion_enabled: true,
tracing: Default::default(),
}
.refresh_config(&RefreshConfigOptions::new())
Expand Down
33 changes: 33 additions & 0 deletions sdk/couchbase-core/src/searchx/query_options.rs
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,31 @@ impl Highlight {
}
}

#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)]
#[non_exhaustive]
pub struct Params {
#[serde(skip_serializing_if = "Option::is_none")]
pub score_rank_constant: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub score_window_size: Option<u32>,
}

impl Params {
pub fn new() -> Self {
Default::default()
}

pub fn score_rank_constant(mut self, value: impl Into<Option<u32>>) -> Self {
self.score_rank_constant = value.into();
self
}

pub fn score_window_size(mut self, value: impl Into<Option<u32>>) -> Self {
self.score_window_size = value.into();
self
}
}

pub type ConsistencyVectors = HashMap<String, HashMap<String, u64>>;

#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)]
Expand Down Expand Up @@ -256,6 +281,8 @@ pub struct QueryOptions {
pub(crate) knn: Option<Vec<KnnQuery>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) knn_operator: Option<KnnOperator>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) params: Option<Params>,

#[serde(skip_serializing_if = "Option::is_none", flatten)]
pub(crate) raw: Option<HashMap<String, serde_json::Value>>,
Expand Down Expand Up @@ -291,6 +318,7 @@ impl QueryOptions {
sort: None,
knn: None,
knn_operator: None,
params: None,
raw: None,
index_name: index_name.into(),
scope_name: None,
Expand Down Expand Up @@ -384,6 +412,11 @@ impl QueryOptions {
self
}

pub fn params(mut self, params: impl Into<Option<Params>>) -> Self {
self.params = params.into();
self
}

pub fn raw(mut self, raw: impl Into<Option<HashMap<String, serde_json::Value>>>) -> Self {
self.raw = raw.into();
self
Expand Down
11 changes: 11 additions & 0 deletions sdk/couchbase-core/src/searchx/search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ pub struct Search<C: Client> {
pub auth: Auth,

pub vector_search_enabled: bool,
pub score_fusion_enabled: bool,
pub(crate) tracing: Arc<TracingComponent>,
}

Expand Down Expand Up @@ -112,6 +113,16 @@ impl<C: Client> Search<C> {
));
}

if !self.score_fusion_enabled {
if let Some(score) = opts.score.as_ref() {
if score == "rrf" || score == "rsf" {
return Err(Error::new_unsupported_feature_error(
"score fusion".to_string(),
));
}
}
}

let req_uri = if let Some(bucket) = &opts.bucket_name {
if let Some(scope) = &opts.scope_name {
format!(
Expand Down
Loading
Loading