Improved knn option for OpenSearch

ankane · Sep 3, 2024 · b307b41 · b307b41
1 parent 6b84664
commit b307b41
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 14 deletions.
diff --git a/lib/searchkick/index_options.rb b/lib/searchkick/index_options.rb
@@ -422,25 +422,35 @@ def generate_mappings
 
       (options[:knn] || []).each do |field, knn_options|
         if Searchkick.opensearch?
-          space_type =
-            case knn_options[:distance]
-            when "cosine", nil
-              "cosinesimil"
-            when "euclidean"
-              "l2"
-            else
-              raise ArgumentError, "Unknown distance: #{distance}"
-            end
+          if knn_options[:distance].nil?
+            # avoid server crash if method not specified
+            raise ArgumentError, "Must specify a distance for OpenSearch"
+          end
 
-          mapping[field.to_s] = {
+          vector_options = {
             type: "knn_vector",
-            dimension: knn_options[:dimensions],
-            method: {
+            dimension: knn_options[:dimensions]
+          }
+
+          if !knn_options[:distance].nil?
+            space_type =
+              case knn_options[:distance]
+              when "cosine"
+                "cosinesimil"
+              when "euclidean"
+                "l2"
+              else
+                raise ArgumentError, "Unknown distance: #{distance}"
+              end
+
+            vector_options[:method] = {
               name: "hnsw",
               space_type: space_type,
               engine: "lucene"
             }
-          }
+          end
+
+          mapping[field.to_s] = vector_options
         else
           vector_options = {
             type: "dense_vector",

diff --git a/test/knn_test.rb b/test/knn_test.rb
@@ -76,6 +76,8 @@ def test_distance
   end
 
   def test_unindexed
+    skip if Searchkick.opensearch?
+
     store [{name: "A", vector: [1, 2, 3]}, {name: "B", vector: [-1, -2, -3]}]
     assert_order "*", ["A", "B"], knn: {field: :vector, vector: [1, 2, 3], distance: "cosine", exact: true}
 

diff --git a/test/models/product.rb b/test/models/product.rb
@@ -21,7 +21,7 @@ class Product
     filterable: [:name, :color, :description],
     similarity: "BM25",
     match: ENV["MATCH"] ? ENV["MATCH"].to_sym : nil,
-    knn: Searchkick.knn_support? ? {embedding: {dimensions: 3, distance: "cosine"}, factors: {dimensions: 3, distance: "euclidean"}, vector: {dimensions: 3}} : nil
+    knn: Searchkick.knn_support? ? {embedding: {dimensions: 3, distance: "cosine"}, factors: {dimensions: 3, distance: "euclidean"}}.merge(Searchkick.opensearch? ? {} : {vector: {dimensions: 3}}) : nil
 
   attr_accessor :conversions, :user_ids, :aisle, :details