Add some additional query tests.

4 files changed, 209 insertions, 0 deletions

diff --git a/jstests/explaina.js b/jstests/explaina.js
new file mode 100644
index 00000000000..8185bf41d59
--- /dev/null
+++ b/jstests/explaina.js
@@ -0,0 +1,47 @@
+// Check explain results when an in order plan is selected among mixed in order and out of order
+// plans.
+
+t = db.jstests_explaina;
+t.drop();
+
+t.ensureIndex( { a:1 } );
+t.ensureIndex( { b:1 } );
+
+for( i = 0; i < 1000; ++i ) {
+    t.save( { a:i, b:i%3 } );
+}
+
+// Query with an initial set of documents.
+explain1 = t.find( { a:{ $gte:0 }, b:2 } ).sort( { a:1 } ).explain( true );
+
+for( i = 1000; i < 2000; ++i ) {
+    t.save( { a:i, b:i%3 } );
+}
+
+// Query with some additional documents.
+explain2 = t.find( { a:{ $gte:0 }, b:2 } ).sort( { a:1 } ).explain( true );
+
+function plan( explain, cursor ) {
+    for( i in explain.allPlans ) {
+        e = explain.allPlans[ i ];
+        if ( e.cursor == cursor ) {
+            return e;
+        }
+    }
+    assert( false );
+}
+
+// Check query totals.
+assert.eq( 333, explain1.n );
+assert.eq( 666, explain2.n );
+
+// Check totals for the selected in order a:1 plan.
+assert.eq( 333, plan( explain1, "BtreeCursor a_1" ).n );
+assert.eq( 1000, plan( explain1, "BtreeCursor a_1" ).nscanned );
+assert.eq( 666, plan( explain2, "BtreeCursor a_1" ).n );
+assert.eq( 2000, plan( explain2, "BtreeCursor a_1" ).nscanned );
+
+// Check that results only examined after the a:1 plan is selected will not affect plan explain
+// output for other plans.
+assert.eq( plan( explain1, "BtreeCursor b_1" ), plan( explain2, "BtreeCursor b_1" ) );
+assert.eq( plan( explain1, "BasicCursor" ), plan( explain2, "BasicCursor" ) );
diff --git a/jstests/finda.js b/jstests/finda.js
new file mode 100644
index 00000000000..c17bcc3c696
--- /dev/null
+++ b/jstests/finda.js
@@ -0,0 +1,98 @@
+// Tests where the QueryOptimizerCursor enters takeover mode during a query rather than a get more.
+
+t = db.jstests_finda;
+t.drop();
+
+numDocs = 200;
+
+function clearQueryPlanCache() {
+    t.ensureIndex( { c:1 } );
+    t.dropIndex( { c:1 } );
+}
+
+function assertAllFound( matches ) {
+//    printjson( matches );
+    found = new Array( numDocs );
+    for( i = 0; i < numDocs; ++i ) {
+        found[ i ] = false;
+    }
+    for( i in matches ) {
+        m = matches[ i ];
+        found[ m._id ] = true;
+    }
+    for( i = 0; i < numDocs; ++i ) {
+        assert( found[ i ], i );
+    }
+}
+
+function makeCursor( query, projection, sort, batchSize, returnKey ) {
+    cursor = t.find( query, projection );
+    if ( sort ) {
+        cursor.sort( sort );
+    }
+    if ( batchSize ) {
+        cursor.batchSize( batchSize );
+    }
+    if ( returnKey ) {
+        cursor._addSpecial( "$returnKey", true );
+    }
+    return cursor;
+}
+
+function checkCursorWithBatchSizeProjection( query, projection, sort, batchSize,
+                                            expectedLeftInBatch ) {
+    clearQueryPlanCache();
+    cursor = makeCursor( query, projection, sort, batchSize );
+    assert.eq( expectedLeftInBatch, cursor.objsLeftInBatch() );
+    assertAllFound( cursor.toArray() );
+}
+
+function checkCursorWithBatchSize( query, sort, batchSize, expectedLeftInBatch ) {
+    checkCursorWithBatchSizeProjection( query, {}, sort, batchSize, expectedLeftInBatch );
+    checkCursorWithBatchSizeProjection( query, { a:1, _id:1 }, sort, batchSize,
+                                       expectedLeftInBatch );
+    // In the cases tested, when expectedLeftInBatch is high enough takeover will occur during
+    // the query operation rather than getMore and the last few matches should properly return keys
+    // from the a,_id index.
+    clearQueryPlanCache();
+    if ( expectedLeftInBatch > 110 ) {
+        cursor = makeCursor( query, {}, sort, batchSize, true );
+        lastNonAIndexResult = -1;
+        for( i = 0; i < expectedLeftInBatch; ++i ) {
+            next = cursor.next();
+            // Identify the query plan used by checking the fields of a returnKey query.
+            if ( !friendlyEqual( [ 'a', '_id' ], Object.keySet( next ) ) ) {
+                lastNonAIndexResult = i;
+            }
+        }
+        // The last results should come from the a,_id index.
+        assert.lt( lastNonAIndexResult, expectedLeftInBatch - 5 );
+    }
+}
+
+function queryWithPlanTypes( withDups ) {
+    t.drop();
+    for( i = 1; i < numDocs; ++i ) {
+        t.save( { _id:i, a:i, b:0 } );
+    }
+    if ( withDups ) {
+        t.save( { _id:0, a:[ 0, numDocs ], b:0 } ); // Add a dup on a:1 index.
+    }
+    else {
+        t.save( { _id:0, a:0, b:0 } );
+    }
+    t.ensureIndex( { a:1, _id:1 } ); // Include _id for a covered index projection.
+
+    // All plans in order.
+    checkCursorWithBatchSize( { a:{ $gte:0 } }, null, 150, 150 );
+
+    // All plans out of order.
+    checkCursorWithBatchSize( { a:{ $gte:0 } }, { c:1 }, null, 200 );
+
+    // Some plans in order, some out of order.
+    checkCursorWithBatchSize( { a:{ $gte:0 }, b:0 }, { a:1 }, 150, 150 );
+    checkCursorWithBatchSize( { a:{ $gte:0 }, b:0 }, { a:1 }, null, 101 );
+}
+
+queryWithPlanTypes( false );
+queryWithPlanTypes( true );
diff --git a/jstests/queryoptimizer8.js b/jstests/queryoptimizer8.js
new file mode 100644
index 00000000000..3218960be41
--- /dev/null
+++ b/jstests/queryoptimizer8.js
@@ -0,0 +1,48 @@
+// Test failover from an in order plan to an out of order plan and vice versa.
+
+t = db.jstests_queryoptimizer8;
+t.drop();
+
+t.ensureIndex( { a:1 } );
+t.ensureIndex( { b:1 } );
+
+t.save( { a:1, b:0 } );
+t.save( { a:0, b:1 } );
+for( i = 0; i < 1000; ++i ) {
+    t.save( { a:1, b:1 } );
+}
+
+aPreferableQuery = { a:0, b:1 };
+bPreferableQuery = { a:1, b:0 };
+
+// Run an explain query to record a specified query plan for the query pattern.
+function recordPlan( plan, query, sort, scanAndOrder ) {
+    explain = t.find( query ).sort( sort ).explain();
+    assert.eq( 'BtreeCursor ' + plan, explain.cursor );
+    assert.eq( scanAndOrder, explain.scanAndOrder );
+}
+
+// Check the plan used for a query.
+function checkPlanUsed( plan, query, sort ) {
+    // Run the query and check its result.
+    results = t.find( query ).sort( sort ).toArray();
+    assert.eq( 1, results.length );
+    result = results[ 0 ];
+    assert.eq( query.a, result.a );
+    assert.eq( query.b, result.b );
+    // Check the plan used for the above query by examining explain's oldPlan field.
+    assert.eq( 'BtreeCursor ' + plan, t.find( query ).sort( sort ).explain( true ).oldPlan.cursor );
+}
+
+function checkOrdered( ordered ) {
+    sort = ordered ? { a:1 } : {};
+    
+    recordPlan( 'a_1', aPreferableQuery, sort, false );
+    checkPlanUsed( 'b_1', bPreferableQuery, sort );
+
+    recordPlan( 'b_1', bPreferableQuery, sort, ordered );
+    checkPlanUsed( 'a_1', aPreferableQuery, sort );
+}
+
+checkOrdered( false );
+checkOrdered( true );
diff --git a/jstests/sortj.js b/jstests/sortj.js
new file mode 100644
index 00000000000..79f85cb8e3d
--- /dev/null
+++ b/jstests/sortj.js
@@ -0,0 +1,16 @@
+// Test an in memory sort memory assertion after a plan has "taken over" in the query optimizer
+// cursor.
+
+t = db.jstests_sortj;
+t.drop();
+
+t.ensureIndex( { a:1 } );
+
+big = new Array( 100000 ).toString();
+for( i = 0; i < 1000; ++i ) {
+    t.save( { a:1, b:big } );
+}
+
+assert.throws( function() {
+              t.find( { a:{ $gte:0 }, c:null } ).sort( { d:1 } ).itcount();
+              } );