From b307b41a58d1a16fc4c31588e19ec75ff9eaef24 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Tue, 3 Sep 2024 16:32:22 -0700
Subject: [PATCH] Improved knn option for OpenSearch

---
 lib/searchkick/index_options.rb | 36 +++++++++++++++++++++------------
 test/knn_test.rb                |  2 ++
 test/models/product.rb          |  2 +-
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/lib/searchkick/index_options.rb b/lib/searchkick/index_options.rb
index a4013bbc..7b8301dd 100644
--- 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
index d258862c..b9396cb9 100644
--- 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
index d4ae93b0..8c902519 100644
--- 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