Merge pull request #474 from estolfo/RUBY-790-count-hint

RUBY-790 Allow count to work with query hints

Author: estolfo

lib/mongo/collection.rb

@@ -1039,19 +1039,24 @@ def stats
     # @option opts [Hash] :query ({}) A query selector for filtering the documents counted.
     # @option opts [Integer] :skip (nil) The number of documents to skip.
     # @option opts [Integer] :limit (nil) The number of documents to limit.
+    # @option opts [String, Array, OrderedHash] :hint hint for query optimizer, usually not necessary if
+    #   using MongoDB > 1.1. This option is only supported with #count in server version > 2.6.
+    # @option opts [String] :named_hint for specifying a named index as a hint, will be overridden by :hint
+    #   if :hint is also provided. This option is only supported with #count in server version > 2.6.
     # @option opts [:primary, :secondary] :read Read preference for this command. See Collection#find for
     #  more details.
     # @option opts [String]  :comment (nil) a comment to include in profiling logs
     #
     # @return [Integer]
     def count(opts={})
       find(opts[:query],
-           :skip  => opts[:skip],
-           :limit => opts[:limit],
-           :read  => opts[:read],
-           :comment => opts[:comment]).count(true)
+           :skip       => opts[:skip],
+           :limit      => opts[:limit],
+           :named_hint => opts[:named_hint] || @hint,
+           :hint       => opts[:hint] || @hint,
+           :read       => opts[:read],
+           :comment    => opts[:comment]).count(true)
     end
-
     alias :size :count
 
     protected

lib/mongo/cursor.rb

@@ -203,7 +203,12 @@ def count(skip_and_limit = false)
         command.merge!(BSON::OrderedHash["skip", @skip]) if @skip != 0
       end
 
+      if @hint
+        hint = @hint.is_a?(String) ? @hint : generate_index_name(@hint)
+      end
+
       command.merge!(BSON::OrderedHash["fields", @fields])
+      command.merge!(BSON::OrderedHash["hint", hint]) if hint
 
       response = @db.command(command, :read => @read, :comment => @comment)
       return response['n'].to_i if Mongo::Support.ok?(response)
@@ -715,5 +720,13 @@ def check_command_cursor
     def compile_regex?
       @compile_regex
     end
+
+    def generate_index_name(spec)
+      indexes = []
+      spec.each_pair do |field, type|
+        indexes.push("#{field}_#{type}")
+      end
+      indexes.join("_")
+    end
   end
 end

test/functional/collection_test.rb

@@ -1008,6 +1008,61 @@ def test_count
     assert_equal 0, @test.count(:skip => 2)
   end
 
+  def test_count_with_hint
+    @test.drop
+    @test.save(:i => 1)
+    @test.save(:i => 2)
+    assert_equal 2, @test.count
+
+    @test.ensure_index(BSON::OrderedHash[:i, Mongo::ASCENDING])
+
+    # Check that a named_hint can be specified
+    assert_equal 1, @test.count(:query => { :i => 1 }, :named_hint => '_id_')
+    assert_equal 2, @test.count(:query => { }, :named_hint => '_id_')
+
+    # Verify that the hint is being sent to the server by providing a bad hint
+    if @version > '2.6'
+      assert_raise Mongo::OperationFailure do
+        @test.count(:query => { :i => 1 }, :hint => 'bad_hint')
+      end
+    else
+      assert_equal 1, @test.count(:query => { :i => 1 }, :hint => 'bad_hint')
+    end
+
+    # Verify that the named_hint is being sent to the server by providing a bad hint
+    if @version > '2.6'
+      assert_raise Mongo::OperationFailure do
+        @test.count(:query => { :i => 1 }, :named_hint => 'bad_hint')
+      end
+    else
+      assert_equal 1, @test.count(:query => { :i => 1 }, :named_hint => 'bad_hint')
+    end
+
+    @test.ensure_index(BSON::OrderedHash[:x, Mongo::ASCENDING], :sparse => true)
+
+    # The sparse index won't have any entries.
+    # Check that count returns 0 when using the hint.
+    expected = @version > '2.6' ? 0 : 1
+    assert_equal expected, @test.count(:query => { :i => 1 }, :hint => { 'x' => 1 })
+    assert_equal expected, @test.count(:query => { :i => 1 }, :hint => 'x')
+    assert_equal expected, @test.count(:query => { :i => 1 }, :named_hint => 'x_1')
+
+    # Verify that the hint / named hint set on the collection is used.
+    @test.hint = { 'x' => 1 }
+    assert_equal expected, @test.count(:query => { :i => 1 })
+
+    @test.hint = 'x'
+    assert_equal expected, @test.count(:query => { :i => 1 })
+
+    # The driver should allow x_1, but the code sets named_hint to @hint without
+    # normalizing.
+    @test.named_hint = 'x'
+    assert_equal expected, @test.count(:query => { :i => 1 })
+
+    assert_equal 2, @test.count(:query => { }, :hint => 'x')
+    assert_equal 2, @test.count(:query => { }, :named_hint => 'x_1')
+  end
+
   # Note: #size is just an alias for #count.
   def test_size
     @test.drop

test/functional/cursor_test.rb

@@ -535,6 +535,59 @@ def test_count_with_fields
     end
   end
 
+  def test_count_with_hint
+    @coll.drop
+    @coll.save(:i => 1)
+    @coll.save(:i => 2)
+    assert_equal 2, @coll.find.count
+
+    @coll.ensure_index(BSON::OrderedHash[:i, Mongo::ASCENDING])
+
+    # Check that a named_hint can be specified
+    assert_equal 1, @coll.find({ :i => 1 }, :named_hint => '_id_').count
+    assert_equal 2, @coll.find({ }, :named_hint => '_id_').count
+
+    # Verify that the hint is being sent to the server by providing a bad hint
+    if @version > '2.6'
+      assert_raise Mongo::OperationFailure do
+        @coll.find({ :i => 1 }, :hint => 'bad_hint').count
+      end
+    else
+      assert_equal 1, @coll.find({ :i => 1 }, :hint => 'bad_hint').count
+    end
+
+    # Verify that the named_hint is being sent to the server by providing a bad hint
+    if @version > '2.6'
+      assert_raise Mongo::OperationFailure do
+        @coll.find({ :i => 1 }, :named_hint => 'bad_hint').count
+      end
+    else
+      assert_equal 1, @coll.find({ :i => 1 }, :named_hint => 'bad_hint').count
+    end
+
+    @coll.ensure_index(BSON::OrderedHash[:x, Mongo::ASCENDING], :sparse => true)
+
+    # The sparse index won't have any entries.
+    # Check that count returns 0 when using the hint.
+    expected = @version > '2.6' ? 0 : 1
+    assert_equal expected, @coll.find({ :i => 1 }, :hint => { 'x' => 1 }).count
+    assert_equal expected, @coll.find({ :i => 1 }, :hint => 'x').count
+    assert_equal expected, @coll.find({ :i => 1 }, :named_hint => 'x_1').count
+
+    # Verify that the hint / named hint set on the collection is used.
+    @coll.hint = { 'x' => 1 }
+    assert_equal expected, @coll.find(:i => 1).count
+
+    @coll.hint = 'x'
+    assert_equal expected, @coll.find(:i => 1).count
+
+    @coll.named_hint = 'x_1'
+    assert_equal expected, @coll.find(:i => 1).count
+
+    assert_equal 2, @coll.find({ }, :hint => 'x').count
+    assert_equal 2, @coll.find({ }, :named_hint => 'x_1').count
+  end
+
   def test_has_next
     @coll.remove
     200.times do |n|