Add KNN

CHANGELOG.md

@@ -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

lib/chewy/search/parameters/knn.rb

@@ -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

lib/chewy/search/request.rb

@@ -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
@@ -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,
@@ -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

spec/chewy/search/parameters/knn_spec.rb

@@ -0,0 +1,5 @@
+require 'chewy/search/parameters/hash_storage_examples'
+
+describe Chewy::Search::Parameters::Knn do
+  it_behaves_like :hash_storage, :knn
+end

spec/chewy/search/request_spec.rb

@@ -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