SERVER-99348 Collect a deterministic sample for testing purpose of Sampling CE (#31831)

GitOrigin-RevId: 405e3fe

Author: RuoxinXu

jstests/noPassthroughWithMongod/query/cbr/cbr_sampling.js

@@ -38,11 +38,15 @@ function assertCollscanUsesSampling(query) {
     assert.eq(plan.estimatesMetadata.ceSource, "Sampling", plan);
 }
 
-function assertAllPlansUseSampling(query) {
+function assertAllPlansUseSampling(query, ce) {
     const explain = coll.find(query).explain();
     [getWinningPlanFromExplain(explain), ...getRejectedPlans(explain)].forEach(plan => {
         assert.eq(plan.estimatesMetadata.ceSource, "Sampling", plan);
-        assert.gt(plan.cardinalityEstimate, 0, plan);
+        if (ce === undefined) {
+            assert.gt(plan.cardinalityEstimate, 0, plan);
+        } else {
+            assert.close(plan.cardinalityEstimate, ce, plan);
+        }
         assert.gt(plan.inputStage.cardinalityEstimate, 0, plan);
         assert.gt(plan.inputStage.numKeysEstimate, 0, plan);
     });
@@ -73,7 +77,28 @@ try {
     assertAllPlansUseSampling({b: {$lt: 100}});
     assertAllPlansUseSampling({a: {$lt: 100}, b: {$lt: 500}});
     assertAllPlansUseSampling({a: {$lt: 100}, b: {$lt: 500}, c: {$exists: true}});
+
+    // Test the sequential scan sampling method that generates the sample by scanning the first N
+    // documents of the collection. The sample generated by this method is repeatable as long as
+    // the collection was populated in the same way.
+    assert.commandWorked(
+        db.adminCommand({setParameter: 1, internalQuerySamplingBySequentialScan: true}));
+    // Run the same query multiple times to ensure the sample is repeatable.
+    assertAllPlansUseSampling({a: {$lt: 100}}, 260.41666);
+    assertAllPlansUseSampling({a: {$lt: 100}}, 260.41666);
+    assertAllPlansUseSampling({a: {$lt: 100}}, 260.41666);
+
+    // Require a sample larger than the collection and test that a full scan of the collection was
+    // done to collect the sample.
+    assert.commandWorked(db.adminCommand({setParameter: 1, samplingMarginOfError: 1.0}));
+    // Since the sample was actually generated from all the documents in the collection. The 'ce' of
+    // this predicate should include all documents, which is 1000.
+    assertAllPlansUseSampling({a: {$lt: 1000}}, 1000);
+    assertAllPlansUseSampling({a: {$lt: 1000}}, 1000);
 } finally {
     // Ensure that query knob doesn't leak into other testcases in the suite.
     assert.commandWorked(db.adminCommand({setParameter: 1, planRankerMode: "multiPlanning"}));
+    assert.commandWorked(
+        db.adminCommand({setParameter: 1, internalQuerySamplingBySequentialScan: false}));
+    assert.commandWorked(db.adminCommand({setParameter: 1, samplingMarginOfError: 5.0}));
 }

src/mongo/db/query/ce/sampling_estimator_impl.cpp

@@ -160,19 +160,9 @@ std::vector<BSONObj> SamplingEstimatorImpl::getIndexKeys(const IndexBounds& boun
                                nullptr /*collator*/,
                                boost::none);
 
-    // This function converts an index key to a BSONObj in order to compare with the IndexBounds.
-    auto keyStringToBson = [](const key_string::Value& keyString) {
-        BSONObjBuilder bob;
-        auto keyStringObj = key_string::toBson(keyString, Ordering::make(BSONObj()));
-        for (auto&& keyStringElem : keyStringObj) {
-            bob.append(keyStringElem);
-        }
-        return bob.obj();
-    };
-
     std::vector<BSONObj> indexKeys;
     for (auto&& keyString : keyStrings) {
-        indexKeys.push_back(keyStringToBson(keyString));
+        indexKeys.push_back(key_string::toBson(keyString, Ordering::make(BSONObj())));
     }
     return indexKeys;
 }
@@ -303,10 +293,12 @@ SamplingEstimatorImpl::generateChunkSamplingPlan(PlanYieldPolicy* sbeYieldPolicy
         0 /* nodeId */);
 
 
+    sbe::value::SlotVector outerProjectsSlots;
+    sbe::value::SlotVector outerCorrelatedSlots{*outerRid};
     auto loopJoinStage = sbe::makeS<sbe::LoopJoinStage>(std::move(outerStage),
                                                         std::move(innerStage),
-                                                        sbe::value::SlotVector{},
-                                                        sbe::value::SlotVector{*outerRid},
+                                                        outerProjectsSlots,
+                                                        outerCorrelatedSlots,
                                                         nullptr /* predicate */,
                                                         0 /* _nodeId */);
 
@@ -414,6 +406,38 @@ void SamplingEstimatorImpl::generateChunkSample() {
     return;
 }
 
+void SamplingEstimatorImpl::generateSampleBySeqScanningForTesting() {
+    // Create a CanonicalQuery for the sampling plan.
+    auto cq = makeCanonicalQuery(_collections.getMainCollection()->ns(), _opCtx, _sampleSize);
+    auto sbeYieldPolicy = PlanYieldPolicySBE::make(
+        _opCtx, PlanYieldPolicy::YieldPolicy::YIELD_AUTO, _collections, cq->nss());
+
+    auto staticData = std::make_unique<stage_builder::PlanStageStaticData>();
+    sbe::value::SlotIdGenerator ids;
+    staticData->resultSlot = ids.generate();
+    const CollectionPtr& collection = _collections.getMainCollection();
+    // Scan the first '_sampleSize' documents sequentially from the start of the target collection
+    // in order to generate a repeatable sample.
+    auto stage = makeScanStage(
+        collection, staticData->resultSlot, boost::none, boost::none, false, sbeYieldPolicy.get());
+    stage = sbe::makeS<sbe::LimitSkipStage>(
+        std::move(stage),
+        sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt64,
+                                   sbe::value::bitcastFrom<int64_t>(_sampleSize)),
+        nullptr /* skip */,
+        0 /* nodeId */);
+
+    stage_builder::PlanStageData data{
+        stage_builder::Environment{std::make_unique<sbe::RuntimeEnvironment>()},
+        std::move(staticData)};
+    auto plan =
+        std::make_pair<std::unique_ptr<sbe::PlanStage>, mongo::stage_builder::PlanStageData>(
+            std::move(stage), std::move(data));
+    executeSamplingQueryAndSample(plan, std::move(cq), std::move(sbeYieldPolicy));
+
+    return;
+}
+
 CardinalityEstimate SamplingEstimatorImpl::estimateCardinality(const MatchExpression* expr) const {
     size_t cnt = 0;
     for (const auto& doc : _sample) {
@@ -527,6 +551,11 @@ SamplingEstimatorImpl::SamplingEstimatorImpl(OperationContext* opCtx,
       _sampleSize(sampleSize),
       _numChunks(numChunks),
       _collectionCard(collectionCard) {
+    if (internalQuerySamplingBySequentialScan.load()) {
+        // This is only used for testing purposes when a repeatable sample is needed.
+        generateSampleBySeqScanningForTesting();
+        return;
+    }
 
     if (sampleSize >= collectionCard.cardinality().v()) {
         // If the required sample is larger than the collection, the sample is generated from all

src/mongo/db/query/ce/sampling_estimator_impl.h

@@ -225,6 +225,13 @@ class SamplingEstimatorImpl : public SamplingEstimator {
         std::unique_ptr<CanonicalQuery> cq,
         std::unique_ptr<PlanYieldPolicySBE> sbeYieldPolicy);
 
+    /**
+     * Generates a sample by sequentially scanning documents from the start of the target
+     * collection. The sample is generated from the first '_sampleSize' documents of the collection.
+     * This sampling method is only used for testing purposes where a repeatable sample is needed.
+     */
+    void generateSampleBySeqScanningForTesting();
+
     /*
      * The SamplingEstimator calculates the size of a sample based on the confidence level and
      * margin of error required.

src/mongo/db/query/query_knobs.idl

@@ -1511,6 +1511,16 @@ server_parameters:
       gt: 0
     redact: false
 
+  internalQuerySamplingBySequentialScan:
+    description: "Indicate sampling CE to generate a repeatable sample by sequentially scanning
+    documents from the start of the target collection. This sampling method is useful for testing
+    purposes and should not be used in production." 
+    set_at: [ startup, runtime ]
+    cpp_varname: "internalQuerySamplingBySequentialScan"
+    cpp_vartype: AtomicWord<bool>
+    default: false
+    redact: false
+
 # Note for adding additional query knobs:
 #
 # When adding a new query knob, you should consider whether or not you need to add an 'on_update'