SERVER-115463 Use comparator with kSortArrays in sbe_non_leading_match_pbt.js (#46710)

GitOrigin-RevId: 1000d78f2f664f6bfa61a4ae0e615ad2f551b1c1

eslint.config.mjs

@@ -125,6 +125,7 @@ export default [
                 _rand: true,
                 _replMonitorStats: true,
                 _resultSetsEqualNormalized: true,
+                _resultSetsEqualUnorderedWithUnorderedArrays: true,
                 _resultSetsEqualUnordered: true,
                 _setShellFailPoint: true,
                 _srand: true,

jstests/aggregation/sources/match/sbe_non_leading_match_pbt.js

@@ -15,15 +15,22 @@ import {getCollectionModel} from "jstests/libs/property_test_helpers/models/coll
 import {makeWorkloadModel} from "jstests/libs/property_test_helpers/models/workload_models.js";
 import {testProperty} from "jstests/libs/property_test_helpers/property_testing_utils.js";
 import {isSlowBuild} from "jstests/libs/query/aggregation_pipeline_utils.js";
-import {sbePushdownEligibleAggModel} from "jstests/libs/property_test_helpers/common_models.js";
+import {checkSbeStatus, kSbeRestricted, kSbeDisabled} from "jstests/libs/query/sbe_util.js";
+import {trySbeRestrictedPushdownEligibleAggModel} from "jstests/libs/property_test_helpers/common_models.js";
 
 if (isSlowBuild(db)) {
     jsTest.log.info("Returning early because debug is on, opt is off, or a sanitizer is enabled.");
     quit();
 }
 
-const numRuns = 10;
-const numQueriesPerRun = 10;
+const sbeStatus = checkSbeStatus(db);
+if (sbeStatus === kSbeDisabled) {
+    jsTest.log.info("SBE is disabled, skipping test.");
+    quit();
+}
+
+const numRuns = 6;
+const numQueriesPerRun = 6;
 const jsTestLogExplain = false;
 
 const controlName = "match_pbt_control";
@@ -42,17 +49,18 @@ const correctnessProperty = createCorrectnessProperty(
 // The inner side of the lookup may be out-of-order between control and experiment. There are
 // $unwind's sprinkled in as a workaround. This blocks SBE pushdown for some suffix of the
 // pipeline.
-// TODO SERVER-115463 use a new comparator with kSortArrays.
 testProperty(
     correctnessProperty,
     {controlColl, experimentColl},
     // Control collection is foreign side.
     makeWorkloadModel({
         collModel: getCollectionModel(),
-        aggModel: sbePushdownEligibleAggModel(controlName),
+        aggModel: trySbeRestrictedPushdownEligibleAggModel(controlName),
         numQueriesPerRun,
     }),
     numRuns,
+    undefined /*examples*/,
+    true /*sortArrays*/,
 );
 
 testProperty(
@@ -61,8 +69,10 @@ testProperty(
     // Experiment collection is foreign side.
     makeWorkloadModel({
         collModel: getCollectionModel(),
-        aggModel: sbePushdownEligibleAggModel(experimentName),
+        aggModel: trySbeRestrictedPushdownEligibleAggModel(experimentName),
         numQueriesPerRun,
     }),
     numRuns,
+    undefined /*examples*/,
+    true /*sortArrays*/,
 );

jstests/libs/property_test_helpers/common_models.js

@@ -7,10 +7,11 @@ import {
     simpleProjectArb,
     getAggPipelineArb,
     getQueryAndOptionsModel,
-    getSbePushdownEligibleAggPipelineArb,
+    getTrySbeRestrictedPushdownEligibleAggPipelineArb,
+    getTrySbeEnginePushdownEligibleAggPipelineArb,
+    getSbeFullPushdownEligibleAggPipelineArb,
 } from "jstests/libs/property_test_helpers/models/query_models.js";
 import {getMatchArb} from "jstests/libs/property_test_helpers/models/match_models.js";
-import {getEqLookupUnwindArb} from "jstests/libs/property_test_helpers/models/lookup_models.js";
 import {groupArb} from "jstests/libs/property_test_helpers/models/group_models.js";
 import {makeWorkloadModel} from "jstests/libs/property_test_helpers/models/workload_models.js";
 import {fc} from "jstests/third_party/fast_check/fc-3.1.0.js";
@@ -99,17 +100,43 @@ export function groupThenMatchAggModel({isTS = false, is83orAbove = true} = {})
     });
 }
 
-export function sbePushdownEligibleAggModel(foreignName, {isTS = false, is83orAbove = true} = {}) {
+export function trySbeRestrictedPushdownEligibleAggModel(foreignName, {isTS = false, is83orAbove = true} = {}) {
     let aggArb = fc.record({
-        pipeline: getSbePushdownEligibleAggPipelineArb(foreignName, {isTS: isTS}),
+        pipeline: getTrySbeRestrictedPushdownEligibleAggPipelineArb(foreignName, {isTS: isTS}),
     });
-
     // Older versions suffer from SERVER-101007
     // TODO SERVER-114269 remove this check.
     if (!is83orAbove) {
         aggArb = aggArb.filter(({pipeline}) => getNestedProperties(pipeline, "$elemMatch").length == 0);
     }
+    return aggArb.map(({pipeline}) => {
+        return {pipeline, "options": {}};
+    });
+}
 
+export function trySbeEnginePushdownEligibleAggModel(foreignName, {isTS = false, is83orAbove = true} = {}) {
+    let aggArb = fc.record({
+        pipeline: getTrySbeEnginePushdownEligibleAggPipelineArb(foreignName, {isTS: isTS}),
+    });
+    // Older versions suffer from SERVER-101007
+    // TODO SERVER-114269 remove this check.
+    if (!is83orAbove) {
+        aggArb = aggArb.filter(({pipeline}) => getNestedProperties(pipeline, "$elemMatch").length == 0);
+    }
+    return aggArb.map(({pipeline}) => {
+        return {pipeline, "options": {}};
+    });
+}
+
+export function sbeFullPushdownEligibleAggModel(foreignName, {isTS = false, is83orAbove = true} = {}) {
+    let aggArb = fc.record({
+        pipeline: getSbeFullPushdownEligibleAggPipelineArb(foreignName, {isTS: isTS}),
+    });
+    // Older versions suffer from SERVER-101007
+    // TODO SERVER-114269 remove this check.
+    if (!is83orAbove) {
+        aggArb = aggArb.filter(({pipeline}) => getNestedProperties(pipeline, "$elemMatch").length == 0);
+    }
     return aggArb.map(({pipeline}) => {
         return {pipeline, "options": {}};
     });

jstests/libs/property_test_helpers/models/lookup_models.js

@@ -5,6 +5,18 @@
 import {assignableFieldArb, fieldArb} from "jstests/libs/property_test_helpers/models/basic_models.js";
 import {fc} from "jstests/third_party/fast_check/fc-3.1.0.js";
 
+export function getEqLookupArb(from) {
+    return fc
+        .record({
+            localField: fieldArb,
+            foreignField: fieldArb,
+            as: assignableFieldArb,
+        })
+        .map(({localField, foreignField, as}) => {
+            return {$lookup: {from, localField, foreignField, as}};
+        });
+}
+
 export function getEqLookupUnwindArb(from) {
     return fc
         .record({

jstests/libs/property_test_helpers/models/query_models.js

@@ -15,7 +15,7 @@ import {
 } from "jstests/libs/property_test_helpers/models/basic_models.js";
 import {collationArb} from "jstests/libs/property_test_helpers/models/collation_models.js";
 import {groupArb} from "jstests/libs/property_test_helpers/models/group_models.js";
-import {getEqLookupUnwindArb} from "jstests/libs/property_test_helpers/models/lookup_models.js";
+import {getEqLookupArb, getEqLookupUnwindArb} from "jstests/libs/property_test_helpers/models/lookup_models.js";
 import {getMatchArb} from "jstests/libs/property_test_helpers/models/match_models.js";
 import {oneof} from "jstests/libs/property_test_helpers/models/model_utils.js";
 import {fc} from "jstests/third_party/fast_check/fc-3.1.0.js";
@@ -129,15 +129,53 @@ export function getAggPipelineArb({allowOrs = true, deterministicBag = true, all
     return fc.array(oneof(...stages), {minLength: 1, maxLength: 6});
 }
 
-export function getSbePushdownEligibleAggPipelineArb(
+export function getTrySbeRestrictedPushdownEligibleAggPipelineArb(
     foreignCollName,
     {allowOrs = true, deterministicBag = true, allowedStages = [], isTS = false} = {},
 ) {
-    const stages = [groupArb, getEqLookupUnwindArb(foreignCollName), getMatchArb()];
+    // This list is ordered from simplest to most complex. This works best for fast check minimization.
+    const stages = [getMatchArb(), groupArb, getEqLookupArb(foreignCollName)];
+    return fc.array(oneof(...stages), {minLength: 1, maxLength: 6});
+}
+
+export function getTrySbeEnginePushdownEligibleAggPipelineArb(
+    foreignCollName,
+    {allowOrs = true, deterministicBag = true, allowedStages = [], isTS = false} = {},
+) {
+    // This list is ordered from simplest to most complex. This works best for fast check minimization.
+    // Not yet included, $window and $unwind.
+    const stages = [];
+    if (!deterministicBag) {
+        stages.push(limitArb, skipArb);
+    }
+    stages.push(
+        simpleProjectArb,
+        getMatchArb(allowOrs),
+        addFieldsConstArb,
+        computedProjectArb,
+        addFieldsVarArb,
+        getSortArb(),
+        groupArb,
+        getEqLookupArb(foreignCollName),
+        getEqLookupUnwindArb(foreignCollName),
+    );
     // eqLookupUnwind returns a javascript array; flatten that here.
     return fc.array(oneof(...stages), {minLength: 1, maxLength: 6}).map((item) => item.flat());
 }
 
+export function getSbeFullPushdownEligibleAggPipelineArb(
+    foreignCollName,
+    {allowOrs = true, deterministicBag = true, allowedStages = [], isTS = false} = {},
+) {
+    // TODO: Add $unwind/$search/$searchMeta arb when available
+    return getTrySbeEnginePushdownEligibleAggPipelineArb(foreignCollName, {
+        allowOrs,
+        deterministicBag,
+        allowedStages,
+        isTS,
+    });
+}
+
 /*
  * Our full model for aggregation pipelines. The full model is an object that contains a pipeline
  * and options. See `getAllowedStages` for description of `allowOrs`

jstests/libs/property_test_helpers/property_testing_utils.js

@@ -83,7 +83,7 @@ const okIndexCreationErrorCodes = [
  * TODO SERVER-98132 redesign getQuery to be more opaque about how many query shapes and constants
  * there are.
  */
-function runProperty(propertyFn, namespaces, workload) {
+function runProperty(propertyFn, namespaces, workload, sortArrays) {
     let {collSpec, queries, extraParams} = workload;
     const {controlColl, experimentColl} = namespaces;
 
@@ -106,7 +106,7 @@ function runProperty(propertyFn, namespaces, workload) {
     });
 
     const testHelpers = {
-        comp: _resultSetsEqualUnordered,
+        comp: sortArrays === true ? _resultSetsEqualUnorderedWithUnorderedArrays : _resultSetsEqualUnordered,
         numQueryShapes: queries.length,
         leafParametersPerFamily,
     };
@@ -149,7 +149,7 @@ function reporter(propertyFn, namespaces) {
  * failure, `runProperty` is called again in the reporter, and prints out more details about the
  * failed property.
  */
-export function testProperty(propertyFn, namespaces, workloadModel, numRuns, examples) {
+export function testProperty(propertyFn, namespaces, workloadModel, numRuns, examples, sortArrays) {
     assert.eq(typeof propertyFn, "function");
     assert(Object.keys(namespaces).every((collName) => collName === "controlColl" || collName === "experimentColl"));
     assert.eq(typeof numRuns, "number");

src/mongo/shell/shell_utils.cpp

@@ -1217,6 +1217,18 @@ BSONObj _resultSetsEqualNormalized(const BSONObj& input, void*) {
     return BSON("" << compareNormalizedResultSets(input, opts));
 }
 
+/*
+ * Takes two arrays of documents, and returns whether they contain the same set of BSON Objects. The
+ * BSON do not need to be in the same order for this to return true. Has no special logic for
+ * handling double/NumberDecimal closeness. Used in property based jstests.
+ * Sorts arrays in order to compare the inner side of lookup results.
+ */
+BSONObj _resultSetsEqualUnorderedWithUnorderedArrays(const BSONObj& input, void*) {
+    auto opts = NormalizationOpts::kSortResults | NormalizationOpts::kSortBSON |
+        NormalizationOpts::kSortArrays;
+    return BSON("" << compareNormalizedResultSets(input, opts));
+}
+
 /*
  * Takes two arrays of documents, and returns whether they contain the same set of BSON Objects. The
  * BSON do not need to be in the same order for this to return true. Has no special logic for
@@ -1288,6 +1300,8 @@ void installShellUtils(Scope& scope) {
     scope.injectNative("_decimal128Limit", _decimal128Limit);
     scope.injectNative("_fnvHashToHexString", _fnvHashToHexString);
     scope.injectNative("_resultSetsEqualUnordered", _resultSetsEqualUnordered);
+    scope.injectNative("_resultSetsEqualUnorderedWithUnorderedArrays",
+                       _resultSetsEqualUnorderedWithUnorderedArrays);
     scope.injectNative("_resultSetsEqualNormalized", _resultSetsEqualNormalized);
     scope.injectNative("_compareStringsWithCollation", _compareStringsWithCollation);
 

src/mongo/shell/shell_utils.d.ts

@@ -29,6 +29,18 @@ declare function _replMonitorStats()
  * @returns True if the result sets compare equal and false otherwise.
  */
 declare function _resultSetsEqualNormalized(a: object[], b: object[]): boolean
+/**
+ * Compares two result sets after sorting arrays.
+ *
+ * @param a First result set.
+ * @param b Second result set.
+ *
+ * @throws {Error} If the size of the BSON representation of 'a' and 'b' exceeds the BSON size limit
+ *                 (~16mb).
+ *
+ * @returns True if the result sets compare equal after sorting arrays and false otherwise.
+ */
+declare function _resultSetsEqualUnorderedWithUnorderedArrays(a: object[], b: object[]): boolean
 declare function _resultSetsEqualUnordered()
 declare function _setShellFailPoint()
 declare function _srand()