Add support for custom search_document_id with record loading

lib/searchkick/results.rb

@@ -248,7 +248,7 @@ def with_hit_and_missing_records
           grouped_hits.each do |index, index_hits|
             results[index] = {}
             index_models[index].each do |model|
-              results[index].merge!(results_query(model, index_hits).to_a.index_by { |r| r.id.to_s })
+              results[index].merge!(results_query(model, index_hits).to_a.index_by { |r| record_search_id(r) })
             end
           end
 
@@ -331,7 +331,11 @@ def results_query(records, hits)
         records = options[:scope_results].call(records)
       end
 
-      Searchkick.load_records(records, ids)
+      if records.respond_to?(:find_by_search_document_ids)
+        records.find_by_search_document_ids(ids)
+      else
+        Searchkick.load_records(records, ids)
+      end
     end
 
     def combine_includes(result, inc)
@@ -344,6 +348,11 @@ def combine_includes(result, inc)
       end
     end
 
+    def record_search_id(record)
+      id = record.respond_to?(:search_document_id) ? record.search_document_id : record.id
+      id.to_s
+    end
+
     def base_field(k)
       k.sub(/\.(analyzed|word_start|word_middle|word_end|text_start|text_middle|text_end|exact)\z/, "")
     end

test/search_document_id_test.rb

@@ -0,0 +1,73 @@
+require_relative "test_helper"
+
+class SearchDocumentIdTest < Minitest::Test
+  def test_custom_search_document_id_with_find_by_search_document_ids
+    skip "ActiveRecord only" if mongoid?
+
+    # Define custom search_document_id on Product
+    Product.class_eval do
+      def search_document_id
+        "custom_#{id}"
+      end
+
+      def self.find_by_search_document_ids(search_ids)
+        ids = search_ids.map { |sid| sid.sub(/^custom_/, "") }
+        where(id: ids)
+      end
+    end
+
+    begin
+      Product.reindex
+
+      store_names ["Product A", "Product B"]
+
+      # Test that search with load: true works
+      results = Product.search("product").to_a
+      assert_equal 2, results.size
+      assert_kind_of Product, results.first
+      assert_includes ["Product A", "Product B"], results.first.name
+      assert_includes ["Product A", "Product B"], results.last.name
+    ensure
+      # Clean up - remove the custom methods
+      Product.class_eval do
+        undef_method :search_document_id if method_defined?(:search_document_id)
+
+        class << self
+          undef_method :find_by_search_document_ids if method_defined?(:find_by_search_document_ids)
+        end
+      end
+
+      Product.reindex
+    end
+  end
+
+  def test_custom_search_document_id_without_find_by_search_document_ids_shows_missing_records
+    skip "ActiveRecord only" if mongoid?
+
+    # Define custom search_document_id on Product but NOT find_by_search_document_ids
+    Product.class_eval do
+      def search_document_id
+        "custom_#{id}"
+      end
+    end
+
+    begin
+      Product.reindex
+
+      store_names ["Product A"]
+
+      # Without find_by_search_document_ids, records should be reported as missing
+      assert_warns "Records in search index do not exist in database" do
+        results = Product.search("product").to_a
+        assert_empty results
+      end
+    ensure
+      # Clean up
+      Product.class_eval do
+        undef_method :search_document_id if method_defined?(:search_document_id)
+      end
+
+      Product.reindex
+    end
+  end
+end