Skip to content

Commit

Permalink
Add KNN (#890)
Browse files Browse the repository at this point in the history
* Add KNN

* Fix PR number in Changelog

* Fix test for collapse and knn
  • Loading branch information
jkostolansky authored Dec 12, 2023
1 parent 1a5594f commit a3f5026
Show file tree
Hide file tree
Showing 5 changed files with 46 additions and 10 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

### New Features

* [#890](https://github.com/toptal/chewy/pull/890): Add the [`knn`](https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html) option to the request. ([@jkostolansky][])

### Changes

### Bugs Fixed
Expand Down
16 changes: 16 additions & 0 deletions lib/chewy/search/parameters/knn.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
require 'chewy/search/parameters/storage'

module Chewy
module Search
class Parameters
# Just a standard hash storage. Nothing to see here.
#
# @see Chewy::Search::Parameters::HashStorage
# @see Chewy::Search::Request#knn
# @see https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html
class Knn < Storage
include HashStorage
end
end
end
end
17 changes: 14 additions & 3 deletions lib/chewy/search/request.rb
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ class Request
UNDEFINED = Class.new.freeze
EVERFIELDS = %w[_index _type _id _parent _routing].freeze
DELEGATED_METHODS = %i[
query filter post_filter order reorder docvalue_fields
query filter post_filter knn order reorder docvalue_fields
track_scores track_total_hits request_cache explain version profile
search_type preference limit offset terminate_after
timeout min_score source stored_fields search_after
Expand All @@ -41,7 +41,7 @@ class Request
EXTRA_STORAGES = %i[aggs suggest].freeze
# An array of storage names that are changing the returned hist collection in any way.
WHERE_STORAGES = %i[
query filter post_filter none min_score rescore indices_boost collapse
query filter post_filter knn none min_score rescore indices_boost collapse
].freeze

delegate :hits, :wrappers, :objects, :records, :documents,
Expand Down Expand Up @@ -520,7 +520,18 @@ def reorder(value, *values)
# @see https://www.elastic.co/guide/en/elasticsearch/reference/current/collapse-search-results.html
# @param value [Hash]
# @return [Chewy::Search::Request]
%i[request_cache search_type preference timeout limit offset terminate_after min_score ignore_unavailable collapse].each do |name|
#
# @!method knn(value)
# Replaces the value of the `knn` request part.
#
# @example
# PlacesIndex.knn(field: :vector, query_vector: [4, 2], k: 5, num_candidates: 50)
# # => <PlacesIndex::Query {..., :body=>{:knn=>{"field"=>:vector, "query_vector"=>[4, 2], "k"=>5, "num_candidates"=>50}}}>
# @see Chewy::Search::Parameters::Knn
# @see https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html
# @param value [Hash]
# @return [Chewy::Search::Request]
%i[request_cache search_type preference timeout limit offset terminate_after min_score ignore_unavailable collapse knn].each do |name|
define_method name do |value|
modify(name) { replace!(value) }
end
Expand Down
5 changes: 5 additions & 0 deletions spec/chewy/search/parameters/knn_spec.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
require 'chewy/search/parameters/hash_storage_examples'

describe Chewy::Search::Parameters::Knn do
it_behaves_like :hash_storage, :knn
end
16 changes: 9 additions & 7 deletions spec/chewy/search/request_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -314,14 +314,16 @@
end
end

describe '#collapse' do
specify { expect(subject.collapse(foo: {bar: 42}).render[:body]).to include(collapse: {'foo' => {bar: 42}}) }
specify do
expect(subject.collapse(foo: {bar: 42}).collapse(moo: {baz: 43}).render[:body])
.to include(collapse: {'moo' => {baz: 43}})
%i[collapse knn].each do |name|
describe "##{name}" do
specify { expect(subject.send(name, foo: {bar: 42}).render[:body]).to include(name => {'foo' => {bar: 42}}) }
specify do
expect(subject.send(name, foo: {bar: 42}).send(name, moo: {baz: 43}).render[:body])
.to include(name => {'moo' => {baz: 43}})
end
specify { expect(subject.send(name, foo: {bar: 42}).send(name, nil).render[:body]).to be_blank }
specify { expect { subject.send(name, foo: {bar: 42}) }.not_to change { subject.render } }
end
specify { expect(subject.collapse(foo: {bar: 42}).collapse(nil).render[:body]).to be_blank }
specify { expect { subject.collapse(foo: {bar: 42}) }.not_to change { subject.render } }
end

describe '#docvalue_fields' do
Expand Down

0 comments on commit a3f5026

Please sign in to comment.