diff --git a/lib/searchkick/query.rb b/lib/searchkick/query.rb index 806a7760..bd956bbb 100644 --- a/lib/searchkick/query.rb +++ b/lib/searchkick/query.rb @@ -918,7 +918,11 @@ def set_knn(payload, knn, per_page, offset) payload[:query] = { script_score: { - query: filter, + query: { + bool: { + must: [filter, {exists: {field: field}}] + } + }, script: { source: "knn_score", lang: "knn", @@ -956,7 +960,11 @@ def set_knn(payload, knn, per_page, offset) payload[:query] = { script_score: { - query: filter, + query: { + bool: { + must: [filter, {exists: {field: field}}] + } + }, script: { source: source, params: { diff --git a/test/knn_test.rb b/test/knn_test.rb index b9396cb9..0a73fbbd 100644 --- a/test/knn_test.rb +++ b/test/knn_test.rb @@ -7,7 +7,7 @@ def setup end def test_basic - store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [-1, -2, -3]}] + store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [-1, -2, -3]}, {name: "C"}] assert_order "*", ["A", "B"], knn: {field: :embedding, vector: [1, 2, 3]} scores = Product.search(knn: {field: :embedding, vector: [1, 2, 3]}).hits.map { |v| v["_score"] } @@ -20,6 +20,7 @@ def test_where {name: "A", store_id: 1, embedding: [1, 2, 3]}, {name: "B", store_id: 2, embedding: [1, 2, 3]}, {name: "C", store_id: 1, embedding: [-1, -2, -3]}, + {name: "D", store_id: 1} ] assert_order "*", ["A", "C"], knn: {field: :embedding, vector: [1, 2, 3]}, where: {store_id: 1} end @@ -29,13 +30,14 @@ def test_pagination {name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [1, 2, 0]}, {name: "C", embedding: [-1, -2, 0]}, - {name: "D", embedding: [-1, -2, -3]} + {name: "D", embedding: [-1, -2, -3]}, + {name: "E"} ] assert_order "*", ["B", "C"], knn: {field: :embedding, vector: [1, 2, 3]}, limit: 2, offset: 1 end def test_euclidean - store [{name: "A", factors: [1, 2, 3]}, {name: "B", factors: [1, 5, 7]}] + store [{name: "A", factors: [1, 2, 3]}, {name: "B", factors: [1, 5, 7]}, {name: "C"}] assert_order "*", ["A", "B"], knn: {field: :factors, vector: [1, 2, 3]} scores = Product.search(knn: {field: :factors, vector: [1, 2, 3]}).hits.map { |v| v["_score"] } @@ -45,7 +47,7 @@ def test_euclidean end def test_exact_cosine - store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [-1, -2, -3]}] + store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [-1, -2, -3]}, {name: "C"}] assert_order "*", ["A", "B"], knn: {field: :embedding, vector: [1, 2, 3], exact: true} scores = Product.search(knn: {field: :embedding, vector: [1, 2, 3], exact: true}).hits.map { |v| v["_score"] } @@ -55,7 +57,7 @@ def test_exact_cosine end def test_exact_euclidean - store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [1, 5, 7]}] + store [{name: "A", embedding: [1, 2, 3]}, {name: "B", embedding: [1, 5, 7]}, {name: "C"}] assert_order "*", ["A", "B"], knn: {field: :embedding, vector: [1, 2, 3], exact: true, distance: "euclidean"} scores = Product.search(knn: {field: :embedding, vector: [1, 2, 3], exact: true, distance: "euclidean"}).hits.map { |v| v["_score"] } @@ -78,7 +80,7 @@ def test_distance def test_unindexed skip if Searchkick.opensearch? - store [{name: "A", vector: [1, 2, 3]}, {name: "B", vector: [-1, -2, -3]}] + store [{name: "A", vector: [1, 2, 3]}, {name: "B", vector: [-1, -2, -3]}, {name: "C"}] assert_order "*", ["A", "B"], knn: {field: :vector, vector: [1, 2, 3], distance: "cosine", exact: true} scores = Product.search(knn: {field: :vector, vector: [1, 2, 3], distance: "cosine", exact: true}).hits.map { |v| v["_score"] }