FIT: Base64 vector search in Java & Kotlin

Change-Id: I6e9492ee79ae675604cb82f9b0e04205ee6f85e1
Reviewed-on: https://review.couchbase.org/c/couchbase-jvm-clients/+/210824
Reviewed-by: David Nault <[email protected]>
Tested-by: Build Bot <[email protected]>

Author: dnault

java-fit-performer/src/main/java/com/couchbase/search/SearchHelper.java

@@ -79,7 +79,6 @@
 import com.couchbase.client.protocol.sdk.search.indexmanager.UpsertIndex;
 import com.couchbase.client.protocol.sdk.search.BlockingSearchResult;
 import com.couchbase.client.protocol.sdk.search.Location;
-import com.couchbase.client.protocol.sdk.search.Search;
 import com.couchbase.client.protocol.sdk.search.SearchFacetResult;
 import com.couchbase.client.protocol.sdk.search.SearchFacets;
 import com.couchbase.client.protocol.sdk.search.SearchFragments;
@@ -89,6 +88,7 @@
 import com.couchbase.client.protocol.shared.ContentAs;
 import com.couchbase.stream.ReactiveSearchResultStreamer;
 import com.couchbase.utils.ContentAsUtil;
+import com.google.common.primitives.Floats;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Timestamp;
 import reactor.core.publisher.Mono;
@@ -1439,6 +1439,41 @@ static AnalyzeDocumentOptions createOptions(AnalyzeDocument request, ConcurrentH
     return out;
   }
 
+  // pre-processor doesn't support nested tags, so do this intricate dance
+  // [start:3.6.0]
+  private static VectorQuery toSdk(com.couchbase.client.protocol.sdk.search.VectorQuery fit) {
+    // [end:3.6.0]
+
+    // [start:3.6.3]
+    if (fit.hasBase64VectorQuery()) {
+      VectorQuery sdk = VectorQuery.create(fit.getVectorFieldName(), fit.getBase64VectorQuery());
+      applyVectorQueryOptions(sdk, fit);
+      return sdk;
+    }
+    // [end:3.6.3]
+
+    // [start:3.6.0]
+    var floats = Floats.toArray(fit.getVectorQueryList());
+    VectorQuery sdk = VectorQuery.create(fit.getVectorFieldName(), floats);
+    applyVectorQueryOptions(sdk, fit);
+    return sdk;
+  }
+  // [end:3.6.0]
+
+  // [start:3.6.0]
+  private static void applyVectorQueryOptions(VectorQuery sdk, com.couchbase.client.protocol.sdk.search.VectorQuery fit) {
+    if (fit.hasOptions()) {
+      var opts = fit.getOptions();
+      if (opts.hasNumCandidates()) {
+        sdk.numCandidates(opts.getNumCandidates());
+      }
+      if (opts.hasBoost()) {
+        sdk.boost(opts.getBoost());
+      }
+    }
+  }
+  // [end:3.6.0]
+
   // [start:3.6.0]
   private static SearchRequest convertSearchRequest(com.couchbase.client.protocol.sdk.search.SearchRequest sr) {
     if (sr.hasSearchQuery()) {
@@ -1456,23 +1491,7 @@ else if (sr.hasVectorSearch()) {
 
   private static VectorSearch convertVectorSearch(com.couchbase.client.protocol.sdk.search.VectorSearch vs) {
     var vectors = vs.getVectorQueryList().stream()
-            .map(v -> {
-              var buffer = new float[v.getVectorQueryCount()];
-              for (int i = 0; i < v.getVectorQueryList().size(); i++) {
-                buffer[i] = v.getVectorQuery(i);
-              }
-              var out = VectorQuery.create(v.getVectorFieldName(), buffer);
-              if (v.hasOptions()) {
-                var opts = v.getOptions();
-                if (opts.hasNumCandidates()) {
-                  out = out.numCandidates(opts.getNumCandidates());
-                }
-                if (opts.hasBoost()) {
-                  out = out.boost(opts.getBoost());
-                }
-              }
-              return out;
-            })
+            .map(SearchHelper::toSdk)
             .toList();
     if (vs.hasOptions()) {
       return VectorSearch.create(vectors, convertVectorSearchOptions(vs.getOptions()));

java-fit-performer/src/main/java/com/couchbase/utils/Capabilities.java

@@ -80,6 +80,10 @@ public static List<Caps> sdkImplementationCaps() {
         out.add(Caps.SDK_VECTOR_SEARCH);
         // [end:3.6.0]
 
+        // [start:3.6.3]
+        out.add(Caps.SDK_VECTOR_SEARCH_BASE64);
+        // [end:3.6.3]
+
         return out;
     }
 }

kotlin-fit-performer/src/main/kotlin/com/couchbase/client/performer/kotlin/KotlinPerformer.kt

@@ -55,6 +55,7 @@ class KotlinPerformer : CorePerformer() {
             .addSdkImplementationCaps(Caps.SDK_SEARCH)
             .addSdkImplementationCaps(Caps.SDK_SCOPE_SEARCH)
             .addSdkImplementationCaps(Caps.SDK_VECTOR_SEARCH)
+            .addSdkImplementationCaps(Caps.SDK_VECTOR_SEARCH_BASE64)
             .addSdkImplementationCaps(Caps.SDK_SEARCH_INDEX_MANAGEMENT)
             .addSdkImplementationCaps(Caps.SDK_SCOPE_SEARCH_INDEX_MANAGEMENT)
             .addPerformerCaps(PerformerCaps.CLUSTER_CONFIG_CERT)

kotlin-fit-performer/src/main/kotlin/com/couchbase/client/performer/kotlin/search/SearchHelper.kt

@@ -452,14 +452,27 @@ private fun FitVectorSearch.toSdk(): VectorSearchSpec {
 }
 
 private fun FitVectorQuery.toSdk(): VectorQuery {
-    val query = if (!hasOptions() || !options.hasNumCandidates())
-        SearchSpec.vector(vectorFieldName, vectorQueryList.toFloatArray())
-    else
-        SearchSpec.vector(
-            field = vectorFieldName,
-            vector = vectorQueryList.toFloatArray(),
-            numCandidates = options.numCandidates,
-        )
+    val query: VectorQuery
+
+    if (hasBase64VectorQuery()) {
+        query = if (!hasOptions() || !options.hasNumCandidates())
+            SearchSpec.vector(vectorFieldName, base64VectorQuery)
+        else
+            SearchSpec.vector(
+                field = vectorFieldName,
+                vector = base64VectorQuery,
+                numCandidates = options.numCandidates,
+            )
+    } else {
+        query = if (!hasOptions() || !options.hasNumCandidates())
+            SearchSpec.vector(vectorFieldName, vectorQueryList.toFloatArray())
+        else
+            SearchSpec.vector(
+                field = vectorFieldName,
+                vector = vectorQueryList.toFloatArray(),
+                numCandidates = options.numCandidates,
+            )
+    }
 
     return query.maybeBoost(hasOptions() && options.hasBoost(), options.boost)
 }