Fix search (and vector search) index errors (#6153)

* make sure collection is created before creating indexes

* Fix filter error in vector search

* Normalize index definitions for comparison

* address Copilot feedback

Author: jamis

lib/mongoid/search_indexable.rb

@@ -85,17 +85,18 @@ def vector_search(index: nil, path: nil, limit: 10, num_candidates: nil, filter:
               "#{resolved_path} is nil on this document; cannot perform vector search"
       end
 
-      self_filter = { '_id' => { '$ne' => _id } }
-      combined_filter = filter ? { '$and' => [ self_filter, filter ] } : self_filter
+      self_exclusion = { '$match' => { '_id' => { '$ne' => _id } } }
+      post_pipeline = [ self_exclusion, { '$limit' => limit }, *Array(pipeline) ]
+      effective_candidates = num_candidates || (limit * 10)
 
       self.class.vector_search(
         query_vector,
         index: index,
         path: path,
-        limit: limit,
-        num_candidates: num_candidates,
-        filter: combined_filter,
-        pipeline: pipeline
+        limit: limit + 1,
+        num_candidates: effective_candidates,
+        filter: filter,
+        pipeline: post_pipeline
       )
     end
 
@@ -129,19 +130,20 @@ def auto_embed_search(index: nil, path: nil, limit: 10, num_candidates: nil, fil
               "#{resolved_path} is nil on this document; cannot perform auto-embed search"
       end
 
-      self_filter = { '_id' => { '$ne' => _id } }
-      combined_filter = filter ? { '$and' => [ self_filter, filter ] } : self_filter
+      self_exclusion = { '$match' => { '_id' => { '$ne' => _id } } }
+      post_pipeline = [ self_exclusion, { '$limit' => limit }, *Array(pipeline) ]
+      effective_candidates = num_candidates || (limit * 10)
 
       self.class.auto_embed_search(
         text,
         index: index,
         path: path,
-        limit: limit,
-        num_candidates: num_candidates,
-        filter: combined_filter,
+        limit: limit + 1,
+        num_candidates: effective_candidates,
+        filter: filter,
         exact: exact,
         model: model,
-        pipeline: pipeline
+        pipeline: post_pipeline
       )
     end
 

spec/mongoid/search_indexable_spec.rb

@@ -17,10 +17,10 @@ def collection
 
   # Wait for all of the indexes with the given names to be ready; then return
   # the list of index definitions corresponding to those names.
-  def wait_for(*names, &condition)
+  def wait_for(*names, timeout: 300, &condition)
     names.flatten!
 
-    timeboxed_wait do
+    timeboxed_wait(max: timeout) do
       result = collection.search_indexes
       return filter_results(result, names) if names.all? { |name| ready?(result, name, &condition) }
     end
@@ -36,6 +36,14 @@ def wait_for_absence_of(*names)
     end
   end
 
+  # Atlas normalizes latestDefinition by adding "fields"=>{} inside mappings
+  # even when no fields were declared. Strip it before comparing against specs.
+  def normalize_definition(defn)
+    return defn unless defn.is_a?(Hash) && defn['mappings'].is_a?(Hash)
+
+    defn.merge('mappings' => defn['mappings'].reject { |k, v| k.to_s == 'fields' && v == {} })
+  end
+
   private
 
   def timeboxed_wait(step: 5, max: 300)
@@ -57,7 +65,7 @@ def ready?(list, name, &condition)
   end
 
   def filter_results(result, names)
-    result.select { |index| names.include?(index['name']) }
+    names.filter_map { |name| result.find { |index| index['name'] == name } }
   end
 end
 
@@ -309,6 +317,153 @@ def filter_results(result, names)
     end
   end
 
+  describe '#vector_search pipeline construction' do
+    let(:model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        field :embedding, type: Array
+        vector_search_index fields: [ { type: 'vector', path: 'embedding', numDimensions: 3, similarity: 'cosine' } ]
+      end
+    end
+
+    let(:fake_collection) { instance_double(Mongo::Collection) }
+    let(:fake_cursor) { double(map: []) }
+    let(:doc) { model.new(embedding: [ 0.1, 0.2, 0.3 ]) }
+
+    before do
+      allow(model).to receive(:collection).and_return(fake_collection)
+      allow(fake_collection).to receive(:aggregate).and_return(fake_cursor)
+    end
+
+    it 'passes limit + 1 to $vectorSearch so the post-filter never short-counts' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        vs = pipeline.find { |s| s['$vectorSearch'] }
+        expect(vs['$vectorSearch']['limit']).to eq 6
+        fake_cursor
+      end
+
+      doc.vector_search(limit: 5)
+    end
+
+    it 'does not use filter in $vectorSearch for self-exclusion' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        vs = pipeline.find { |s| s['$vectorSearch'] }
+        expect(vs['$vectorSearch']).not_to have_key('filter')
+        fake_cursor
+      end
+
+      doc.vector_search(limit: 5)
+    end
+
+    it 'adds a $match stage after $vectorSearch to exclude self' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        match = pipeline.find { |s| s['$match'] }
+        expect(match).to eq({ '$match' => { '_id' => { '$ne' => doc.id } } })
+        fake_cursor
+      end
+
+      doc.vector_search(limit: 5)
+    end
+
+    it 'adds a $limit stage after $match to cap results at the requested limit' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        match_idx = pipeline.index { |s| s['$match'] }
+        limit_stage = pipeline[match_idx + 1]
+        expect(limit_stage).to eq({ '$limit' => 5 })
+        fake_cursor
+      end
+
+      doc.vector_search(limit: 5)
+    end
+
+    it 'passes a user-provided filter through to $vectorSearch' do
+      user_filter = { 'status' => 'published' }
+
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        vs = pipeline.find { |s| s['$vectorSearch'] }
+        expect(vs['$vectorSearch']['filter']).to eq user_filter
+        fake_cursor
+      end
+
+      doc.vector_search(limit: 5, filter: user_filter)
+    end
+  end
+
+  describe '#auto_embed_search pipeline construction' do
+    let(:model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        auto_embed_field :description, model: 'voyage-4'
+      end
+    end
+
+    let(:fake_collection) { instance_double(Mongo::Collection) }
+    let(:fake_cursor) { double(map: []) }
+    let(:doc) { model.new(description: 'hello world') }
+
+    before do
+      allow(model).to receive(:collection).and_return(fake_collection)
+      allow(fake_collection).to receive(:aggregate).and_return(fake_cursor)
+    end
+
+    it 'passes limit + 1 to $vectorSearch so the post-filter never short-counts' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        vs = pipeline.find { |s| s['$vectorSearch'] }
+        expect(vs['$vectorSearch']['limit']).to eq 6
+        fake_cursor
+      end
+
+      doc.auto_embed_search(limit: 5)
+    end
+
+    it 'does not use filter in $vectorSearch for self-exclusion' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        vs = pipeline.find { |s| s['$vectorSearch'] }
+        expect(vs['$vectorSearch']).not_to have_key('filter')
+        fake_cursor
+      end
+
+      doc.auto_embed_search(limit: 5)
+    end
+
+    it 'adds a $match stage after $vectorSearch to exclude self' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        match = pipeline.find { |s| s['$match'] }
+        expect(match).to eq({ '$match' => { '_id' => { '$ne' => doc.id } } })
+        fake_cursor
+      end
+
+      doc.auto_embed_search(limit: 5)
+    end
+
+    it 'adds a $limit stage after $match to cap results at the requested limit' do
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        match_idx = pipeline.index { |s| s['$match'] }
+        limit_stage = pipeline[match_idx + 1]
+        expect(limit_stage).to eq({ '$limit' => 5 })
+        fake_cursor
+      end
+
+      doc.auto_embed_search(limit: 5)
+    end
+
+    it 'passes a user-provided filter through to $vectorSearch' do
+      user_filter = { 'status' => 'published' }
+
+      expect(fake_collection).to receive(:aggregate) do |pipeline|
+        vs = pipeline.find { |s| s['$vectorSearch'] }
+        expect(vs['$vectorSearch']['filter']).to eq user_filter
+        fake_cursor
+      end
+
+      doc.auto_embed_search(limit: 5, filter: user_filter)
+    end
+  end
+
   # Atlas integration tests — skipped when ATLAS_URI is not set.
 
   context 'Atlas integration' do
@@ -332,7 +487,7 @@ def filter_results(result, names)
       let(:requested_definitions) { model.search_index_specs.map { |spec| spec[:definition].with_indifferent_access } }
       let(:index_names) { model.create_search_indexes }
       let(:actual_indexes) { helper.wait_for(*index_names) }
-      let(:actual_definitions) { actual_indexes.map { |i| i['latestDefinition'] } }
+      let(:actual_definitions) { actual_indexes.map { |i| helper.normalize_definition(i['latestDefinition']) } }
 
       describe '.create_search_indexes' do
         it 'creates the indexes' do
@@ -343,10 +498,10 @@ def filter_results(result, names)
       describe '.search_indexes' do
         before { actual_indexes } # wait for the indices to be created
 
-        let(:queried_definitions) { model.search_indexes.map { |i| i['latestDefinition'] } }
+        let(:queried_definitions) { model.search_indexes.map { |i| helper.normalize_definition(i['latestDefinition']) } }
 
         it 'queries the available search indexes' do
-          expect(queried_definitions).to eq requested_definitions
+          expect(queried_definitions).to match_array(requested_definitions)
         end
       end
 
@@ -390,11 +545,19 @@ def filter_results(result, names)
       let(:vector_helper) { SearchIndexHelper.new(vector_model) }
 
       # Three orthogonal unit vectors as a minimal, predictable dataset.
-      let!(:doc_a) { vector_model.create!(embedding: [ 1.0, 0.0, 0.0 ]) }
-      let!(:doc_b) { vector_model.create!(embedding: [ 0.0, 1.0, 0.0 ]) }
-      let!(:doc_c) { vector_model.create!(embedding: [ 0.0, 0.0, 1.0 ]) }
+      let(:doc_a) { vector_model.create!(embedding: [ 1.0, 0.0, 0.0 ]) }
+      let(:doc_b) { vector_model.create!(embedding: [ 0.0, 1.0, 0.0 ]) }
+      let(:doc_c) { vector_model.create!(embedding: [ 0.0, 0.0, 1.0 ]) }
 
       before do
+        vector_helper # force collection to be dropped, created
+
+        # once the collection has been recreated, populate it with documents
+        doc_a
+        doc_b
+        doc_c
+
+        # prepare the search indexes
         names = vector_model.create_search_indexes
         vector_helper.wait_for(*names)
       end

spec/spec_helper.rb

@@ -101,7 +101,11 @@ def test_model(name: 'TestModel', &block)
         heartbeat_frequency: 180,
         user: SpecConfig.instance.uri.client_options[:user] || MONGOID_ROOT_USER.name,
         password: SpecConfig.instance.uri.client_options[:password] || MONGOID_ROOT_USER.password,
-        auth_source: Mongo::Database::ADMIN
+        auth_source: Mongo::Database::ADMIN,
+
+        # so that we can run an Atlas cluster locally, and run tests with Atlas
+        # support enabled.
+        direct_connection: SpecConfig.instance.uri.uri_options[:direct_connection]
       )
     }
   },