diff --git a/lib/mongo/collection.rb b/lib/mongo/collection.rb index 12e53b76fb..65d1702058 100644 --- a/lib/mongo/collection.rb +++ b/lib/mongo/collection.rb @@ -1039,6 +1039,10 @@ 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 @@ -1046,12 +1050,13 @@ def stats # @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 diff --git a/lib/mongo/cursor.rb b/lib/mongo/cursor.rb index 51952a0a72..1b659555e8 100644 --- a/lib/mongo/cursor.rb +++ b/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 diff --git a/test/functional/collection_test.rb b/test/functional/collection_test.rb index 7fdf65a435..f6cc07eb72 100644 --- a/test/functional/collection_test.rb +++ b/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 diff --git a/test/functional/cursor_test.rb b/test/functional/cursor_test.rb index 4e8d1c5745..17dc15a5d8 100644 --- a/test/functional/cursor_test.rb +++ b/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|