RCBC-554: Add support for search score fusion

ext/rcb_search.cxx

@@ -29,6 +29,7 @@
 #include <core/operations/management/search_index_get_documents_count.hxx>
 #include <core/operations/management/search_index_get_stats.hxx>
 #include <core/operations/management/search_index_upsert.hxx>
+#include <core/search_scoring.hxx>
 
 #include <gsl/narrow>
 #include <spdlog/fmt/bundled/core.h>
@@ -926,6 +927,28 @@ cb_Backend_document_search(VALUE self,
     cb_extract_option_bool(req.include_locations, options, "include_locations");
     cb_extract_option_bool(req.show_request, options, "show_request");
 
+    if (VALUE scoring = rb_hash_aref(options, rb_id2sym(rb_intern("scoring"))); !NIL_P(scoring)) {
+      cb_check_type(scoring, T_HASH);
+      VALUE mode = rb_hash_aref(scoring, rb_id2sym(rb_intern("mode")));
+      cb_check_type(mode, T_SYMBOL);
+      VALUE params = rb_hash_aref(scoring, rb_id2sym(rb_intern("params")));
+      if (ID mode_type = rb_sym2id(mode); mode_type == rb_intern("none")) {
+        req.scoring = core::search_scoring_none{};
+      } else if (mode_type == rb_intern("reciprocal_rank_fusion")) {
+        core::search_scoring_reciprocal_rank_fusion rrf{};
+        cb_extract_option_number(rrf.rank_constant, params, "rank_constant");
+        cb_extract_option_number(rrf.window_size, params, "window_size");
+        req.scoring = rrf;
+      } else if (mode_type == rb_intern("relative_score_fusion")) {
+        core::search_scoring_relative_score_fusion rsf{};
+        cb_extract_option_number(rsf.window_size, params, "window_size");
+        req.scoring = rsf;
+      } else {
+        throw ruby_exception(rb_eArgError,
+                             rb_sprintf("unknown search scoring mode: %+" PRIsVALUE, mode));
+      }
+    }
+
     if (VALUE vector_options = rb_hash_aref(search_request, rb_id2sym(rb_intern("vector_search")));
         !NIL_P(vector_options)) {
       cb_check_type(vector_options, T_HASH);

fit-performer/lib/fit/performer/commands/search/options_builder.rb

@@ -178,6 +178,39 @@ def set_serializer
 
             self
           end
+
+          def set_disable_scoring
+            return self unless @raw_options.has_disable_scoring?
+
+            @options.disable_scoring = @raw_options.disable_scoring
+
+            self
+          end
+
+          def set_scoring
+            return self unless @raw_options.has_scoring?
+
+            proto_scoring = @raw_options.scoring
+
+            case proto_scoring.mode
+            when :reciprocal_rank_fusion
+              @options.scoring = Couchbase::SearchScoring.reciprocal_rank_fusion do |s|
+                proto_rrf = proto_scoring.reciprocal_rank_fusion
+                s.rank_constant = proto_rrf.rank_constant if proto_rrf.has_rank_constant?
+                s.window_size = proto_rrf.window_size if proto_rrf.has_window_size?
+              end
+            when :relative_score_fusion
+              @options.scoring = Couchbase::SearchScoring.relative_score_fusion do |s|
+                proto_rsf = proto_scoring.relative_score_fusion
+                s.window_size = proto_rsf.window_size if proto_rsf.has_window_size?
+              end
+            when :none
+              @options.scoring = Couchbase::SearchScoring.none
+            else
+              raise PerformerError, "Unknown search scoring mode `#{proto_scoring.mode}`"
+            end
+            self
+          end
         end
       end
     end

fit-performer/lib/fit/performer/commands/search/search_command.rb

@@ -56,6 +56,8 @@ def set_options
                    .set_raw
                    .set_include_locations
                    .set_serializer
+                   .set_disable_scoring
+                   .set_scoring
                    .set_parent_span(@get_span_fn)
             @cmd_args.append(builder.options)
           end

fit-performer/lib/fit/performer/commands/search/search_query_command.rb

@@ -58,6 +58,8 @@ def set_options
                    .set_facets
                    .set_raw
                    .set_include_locations
+                   .set_disable_scoring
+                   .set_scoring
                    .set_parent_span(@get_span_fn)
             @cmd_args.append(builder.options)
           end

fit-performer/lib/fit/performer/service.rb

@@ -77,6 +77,7 @@ class Service < FIT::Protocol::PerformerService::Service
         :SDK_OBSERVABILITY_RFC_REV_24,
         :SDK_STABLE_OTEL_SEMANTIC_CONVENTIONS,
         :SDK_STABLE_OTEL_SEMANTIC_CONVENTIONS_EMITTED_BY_DEFAULT,
+        :SDK_SEARCH_SCORE_FUSION,
       ].freeze
 
       # We don't currently support transactions. However, the driver calls transactions_factory_create during the

fit-performer/proto/.clang-format

@@ -0,0 +1,7 @@
+# -*- mode: yaml; -*-
+
+# The .proto files are copied verbatim from the FIT protocol repository, so they
+# must not be reformatted by clang-format (which formats .proto files by default)
+---
+DisableFormat: true
+SortIncludes: false

fit-performer/proto/sdk.caps.proto

@@ -152,4 +152,14 @@ 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 hybrid search score fusion (SDK-RFC 52).
+  SDK_SEARCH_SCORE_FUSION = 41;
+
+  // SDK handles query code 2120.
+  SDK_QUERY_2120 = 42;
+
+  // The SDK has support for the Collection.GetReplica API. This includes the new DocumentNotFoundOnReplicaException,
+  // ReplicaIndexOutOfBoundsException and ReplicaIndexCurrentlyUnavailableException exceptions.
+  SDK_GET_REPLICA = 43;
 }

fit-performer/proto/sdk.cluster.query.index_manager.proto

@@ -8,8 +8,6 @@ option ruby_package = "FIT::Protocol::SDK::Cluster::Query::IndexManager";
 option java_multiple_files = true;
 
 import "sdk.query.index_manager.proto";
-import "sdk.query.index_manager.options.proto";
-
 
 // This file is for QueryIndexManager: cluster.queryIndexes()
 // For Collection-level query index manager (collection.queryIndexes()), see sdk.collection.query.index_manager.proto.

fit-performer/proto/sdk.kv.commands.proto

@@ -13,6 +13,7 @@ import "shared.content.proto";
 import "shared.basic.proto";
 import "google/protobuf/duration.proto";
 import "streams.top_level.proto";
+import "sdk.kv.replicas.proto";
 
 
 // Performer will return a MutationResult
@@ -68,6 +69,15 @@ message GetAnyReplica {
   shared.ContentAs content_as = 3;
 }
 
+// Performer will return a GetReplicaResult
+message GetReplica {
+  shared.DocLocation location = 1;
+  replicas.GetReplicaStrategy strategy = 2;
+  optional GetReplicaOptions options = 3;
+
+  shared.ContentAs content_as = 4;
+}
+
 // Performer will return an empty result with success set
 message Unlock {
   shared.DocLocation location = 1;
@@ -145,4 +155,4 @@ message GetReplicaResult {
   // Not part of the SDK response - identifies what stream this is from.
   // This is only needed when returning the result of collection.getAllReplicas()
   optional string stream_id = 5;
-}
+}

fit-performer/proto/sdk.kv.options.proto

@@ -98,4 +98,10 @@ message GetAnyReplicaOptions {
     optional shared.ReadPreference read_preference = 3;
 
     optional string parent_span_id = 4;
-}
+}
+
+message GetReplicaOptions {
+    optional int32 timeout_msecs = 1;
+    optional shared.Transcoder transcoder = 2;
+    optional string parent_span_id = 3;
+}

fit-performer/proto/sdk.kv.replicas.proto

@@ -0,0 +1,30 @@
+syntax = "proto3";
+
+package protocol.sdk.kv.replicas;
+option csharp_namespace = "Couchbase.Grpc.Protocol.Sdk.Kv.Replicas";
+option java_package = "com.couchbase.client.protocol.sdk.kv.replicas";
+option go_package = "github.com/couchbaselabs/transactions-fit-performer/protocol/sdk/kv/replicas";
+option ruby_package = "FIT::Protocol::SDK::KV::Replicas";
+option java_multiple_files = true;
+
+enum ReplicaIndex {
+  FIRST = 0;
+  SECOND = 1;
+  THIRD = 2;
+}
+
+message GetReplicaStrategyFromIndexOptions {
+  optional bool wrap = 1;
+}
+
+// Equivalent to a GetReplicaStrategy.fromIndex call in the SDK
+message GetReplicaStrategyFromIndex {
+  ReplicaIndex index = 1;
+  optional GetReplicaStrategyFromIndexOptions options = 2;
+}
+
+message GetReplicaStrategy {
+  oneof strategy {
+    GetReplicaStrategyFromIndex from_index = 1;
+  }
+}

fit-performer/proto/sdk.search.proto

@@ -89,6 +89,26 @@ 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 `score` field in the FTS request.
+message SearchScoring {
+  oneof mode {
+    SearchScoringReciprocalRankFusion reciprocal_rank_fusion = 1;
+    SearchScoringRelativeScoreFusion relative_score_fusion = 2;
+    SearchScoringNone none = 3;
+  }
+}
 
 // Executing a `cluster.searchQuery()` FTS query.
 //
@@ -319,6 +339,8 @@ message SearchOptions {
   map<string, string> raw = 12;
   optional bool include_locations = 13;
   optional protocol.shared.JsonSerializer serialize = 14;
+  optional SearchScoring scoring = 15;
+  optional bool disable_scoring = 16;
 }
 
 enum MatchOperator {

fit-performer/proto/sdk.workload.proto

@@ -118,6 +118,8 @@ message CollectionLevelCommand {
 
     kv.GetAllReplicas get_all_replicas = 31;
 
+    kv.GetReplica get_replica = 33;
+
     kv.lookup_in.LookupInAnyReplica lookup_in_any_replica = 19;
 
     BinaryCollectionLevelCommand binary = 32;

fit-performer/proto/shared.exceptions.proto

@@ -84,6 +84,9 @@ enum CouchbaseExceptionType {
   SDK_XATTR_UNKNOWN_VIRTUAL_ATTRIBUTE_EXCEPTION = 127;
   SDK_XATTR_CANNOT_MODIFY_VIRTUAL_ATTRIBUTE_EXCEPTION = 128;
   SDK_XATTR_NO_ACCESS_EXCEPTION = 130;
+  SDK_DOCUMENT_NOT_FOUND_ON_REPLICA_EXCEPTION = 132;
+  SDK_REPLICA_INDEX_OUT_OF_BOUNDS_EXCEPTION = 133;
+  SDK_REPLICA_INDEX_CURRENTLY_UNAVAILABLE_EXCEPTION = 134;
 
   // Query
   SDK_PLANNING_FAILURE_EXCEPTION = 201;

lib/couchbase/options.rb

@@ -2329,7 +2329,12 @@ class Search < Base
       attr_accessor :highlight_style # @return [Symbol]
       attr_accessor :highlight_fields # @return [Array<String>]
       attr_accessor :fields # @return [Array<String>]
-      attr_accessor :disable_scoring # @return [Boolean]
+
+      # @deprecated Use +scoring+ with {SearchScoring.none} instead
+      # @return [Boolean]
+      attr_accessor :disable_scoring
+
+      attr_accessor :scoring # @return [SearchScoring, nil]
       attr_accessor :include_locations # @return [Boolean]
       attr_accessor :collections # @return [Array<String>, nil]
       attr_accessor :sort # @return [Array<String, Cluster::SearchSort>]
@@ -2346,7 +2351,12 @@ class Search < Base
       # @param [Array<String>] fields list of field values which should be retrieved for result documents, provided they
       #   were stored while indexing
       # @param [MutationState] mutation_state the mutation tokens this query should be consistent with
-      # @param [Boolean] disable_scoring If set to true, the server will not perform any scoring on the hits
+      # @param [Boolean] disable_scoring DEPRECATED: Use +scoring+ with {SearchScoring.none} instead. If set to true,
+      #   the server will not perform any scoring on the hits.
+      # @param [SearchScoring, nil] scoring specifies the scoring mode used for the request. For a
+      #   hybrid search (a traditional FTS query combined with one or more vector queries) a fusion strategy
+      #   controls how the FTS and vector result sets are merged into a single ranked list. Must not be used
+      #   together with +disable_scoring+.
       # @param [Boolean] include_locations UNCOMMITTED: If set to true, will include the vector of search_location in rows
       # @param [Array<String>, nil] collections list of collections by which to filter the results
       # @param [Array<String, Cluster::SearchSort>] sort Ordering rules to apply to the results. The list might contain
@@ -2372,6 +2382,7 @@ def initialize(limit: nil,
                      fields: nil,
                      mutation_state: nil,
                      disable_scoring: false,
+                     scoring: nil,
                      include_locations: false,
                      collections: nil,
                      sort: nil,
@@ -2389,6 +2400,7 @@ def initialize(limit: nil,
         @highlight_fields = highlight_fields
         @fields = fields
         @disable_scoring = disable_scoring
+        @scoring = scoring
         @include_locations = include_locations
         @collections = collections
         @sort = sort
@@ -2432,14 +2444,22 @@ def scan_consistency=(level)
       # @return [Symbol]
       attr_reader :scan_consistency
 
+      # @api private
+      def validate_scoring
+        raise Error::InvalidArgument, "disable_scoring and scoring must not be used together" if @disable_scoring && !@scoring.nil?
+      end
+
       # @api private
       def to_backend(show_request: nil)
+        validate_scoring
+
         {
           timeout: Utils::Time.extract_duration(@timeout),
           limit: @limit,
           skip: @skip,
           explain: @explain,
           disable_scoring: @disable_scoring,
+          scoring: @scoring&.to_backend,
           include_locations: @include_locations,
           collections: @collections,
           highlight_style: @highlight_style,

lib/couchbase/protostellar/request_generator/search.rb

@@ -16,6 +16,7 @@
 
 require "google/protobuf/well_known_types"
 
+require "couchbase/errors"
 require "couchbase/search_options"
 
 require "couchbase/protostellar/generated/search/v1/search_pb"
@@ -48,7 +49,7 @@ def search_query_request(index_name, query, options)
             scan_consistency: SCAN_CONSISTENCY_MAP[options.scan_consistency],
             include_explanation: options.explain,
             highlight_style: HIGHLIGHT_STYLE_MAP[options.highlight_style],
-            disable_scoring: options.disable_scoring,
+            disable_scoring: get_disable_scoring(options),
             include_locations: options.include_locations,
           }
 
@@ -344,6 +345,21 @@ def get_sort(options)
           end
         end
 
+        # The only scoring mode supported by the protocol is {Couchbase::SearchScoring.none}.
+        # This will be updated to support all types of scoring once the protostellar protocol supports them.
+        def get_disable_scoring(options)
+          return options.disable_scoring if options.scoring.nil?
+
+          options.validate_scoring
+
+          unless options.scoring.is_a?(Couchbase::SearchScoring::SearchScoringNone)
+            raise Couchbase::Error::FeatureNotAvailable,
+                  "The #{Protostellar::NAME} protocol does not support #{options.scoring.class} scoring"
+          end
+
+          true
+        end
+
         def get_facets(options)
           return {} if options.facets.nil?