Skip to content

Commit 8d6e378

Browse files
committed
Updated KNN tests for Elasticsearch 9 and OpenSearch 3
1 parent b2050f2 commit 8d6e378

File tree

2 files changed

+30
-12
lines changed

2 files changed

+30
-12
lines changed

lib/searchkick/query.rb

+5
Original file line numberDiff line numberDiff line change
@@ -950,6 +950,11 @@ def set_knn(payload, knn, per_page, offset)
950950
end
951951
else
952952
if exact
953+
# prevent incorrect distances/results with Elasticsearch 9.0.0-rc1
954+
if !below90? && field_options[:distance] == "cosine" && distance != "cosine"
955+
raise ArgumentError, "distance must match searchkick options"
956+
end
957+
953958
# https://github.com/elastic/elasticsearch/blob/main/docs/reference/vectors/vector-functions.asciidoc
954959
source =
955960
case distance

test/knn_test.rb

+25-12
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@ class KnnTest < Minitest::Test
44
def setup
55
skip unless Searchkick.knn_support?
66
super
7+
8+
# prevent null_pointer_exception with OpenSearch 3.0.0-alpha1
9+
Product.reindex if Searchkick.opensearch? && !Searchkick.server_below?("3.0.0", true)
710
end
811

912
def test_basic
@@ -76,19 +79,19 @@ def test_euclidean
7679
end
7780

7881
def test_euclidean_exact
79-
store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [1, 5, 7]}, {name: "C"}]
80-
assert_order "*", ["A", "B"], knn: {field: :embedding, vector: [1, 2, 3], distance: "euclidean"}
82+
store [{name: "A", embedding2: [1, 2, 3]}, {name: "B", embedding2: [1, 5, 7]}, {name: "C"}]
83+
assert_order "*", ["A", "B"], knn: {field: :embedding2, vector: [1, 2, 3], distance: "euclidean"}
8184

82-
scores = Product.search(knn: {field: :embedding, vector: [1, 2, 3], distance: "euclidean"}).hits.map { |v| v["_score"] }
85+
scores = Product.search(knn: {field: :embedding2, vector: [1, 2, 3], distance: "euclidean"}).hits.map { |v| v["_score"] }
8386
assert_in_delta 1.0 / (1 + 0), scores[0]
8487
assert_in_delta 1.0 / (1 + 5**2), scores[1]
8588
end
8689

8790
def test_taxicab_exact
88-
store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [1, 5, 7]}, {name: "C"}]
89-
assert_order "*", ["A", "B"], knn: {field: :embedding, vector: [1, 2, 3], distance: "taxicab"}
91+
store [{name: "A", embedding2: [1, 2, 3]}, {name: "B", embedding2: [1, 5, 7]}, {name: "C"}]
92+
assert_order "*", ["A", "B"], knn: {field: :embedding2, vector: [1, 2, 3], distance: "taxicab"}
9093

91-
scores = Product.search(knn: {field: :embedding, vector: [1, 2, 3], distance: "taxicab"}).hits.map { |v| v["_score"] }
94+
scores = Product.search(knn: {field: :embedding2, vector: [1, 2, 3], distance: "taxicab"}).hits.map { |v| v["_score"] }
9295
assert_in_delta 1.0 / (1 + 0), scores[0]
9396
assert_in_delta 1.0 / (1 + 7), scores[1]
9497
end
@@ -116,10 +119,10 @@ def test_inner_product
116119
end
117120

118121
def test_inner_product_exact
119-
store [{name: "A", embedding: [-1, -2, -3]}, {name: "B", embedding: [1, 5, 7]}, {name: "C"}]
120-
assert_order "*", ["B", "A"], knn: {field: :embedding, vector: [1, 2, 3], distance: "inner_product"}
122+
store [{name: "A", embedding3: [-1, -2, -3]}, {name: "B", embedding3: [1, 5, 7]}, {name: "C"}]
123+
assert_order "*", ["B", "A"], knn: {field: :embedding3, vector: [1, 2, 3], distance: "inner_product"}
121124

122-
scores = Product.search(knn: {field: :embedding, vector: [1, 2, 3], distance: "inner_product"}).hits.map { |v| v["_score"] }
125+
scores = Product.search(knn: {field: :embedding3, vector: [1, 2, 3], distance: "inner_product"}).hits.map { |v| v["_score"] }
123126
assert_in_delta 1 + 32, scores[0]
124127
assert_in_delta 1.0 / (1 + 14), scores[1]
125128
end
@@ -148,15 +151,25 @@ def test_unindexed
148151
Product.search(knn: {field: :embedding, vector: [1, 2, 3], distance: "euclidean", exact: false})
149152
end
150153
assert_equal "distance must match searchkick options for approximate search", error.message
154+
155+
if !Searchkick.server_below?("9.0.0")
156+
error = assert_raises(ArgumentError) do
157+
Product.search(knn: {field: :embedding, vector: [1, 2, 3], distance: "euclidean"})
158+
end
159+
assert_equal "distance must match searchkick options", error.message
160+
end
151161
end
152162

153163
def test_explain
154164
store [{name: "A", embedding: [1, 2, 3], embedding2: [1, 2, 3], embedding3: [1, 2, 3], embedding4: [1, 2, 3]}]
155165

156166
assert_approx true, :embedding, "cosine"
157-
assert_approx false, :embedding, "euclidean"
158-
assert_approx false, :embedding, "inner_product"
159-
assert_approx false, :embedding, "taxicab"
167+
168+
if Searchkick.opensearch? || Searchkick.server_below?("9.0.0")
169+
assert_approx false, :embedding, "euclidean"
170+
assert_approx false, :embedding, "inner_product"
171+
assert_approx false, :embedding, "taxicab"
172+
end
160173

161174
if Searchkick.opensearch?
162175
assert_approx false, :embedding, "chebyshev"

0 commit comments

Comments
 (0)