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

GitOrigin-RevId: 96cfa4c555d156e505c866a66b82ac8482bce7c0

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,13 +15,23 @@ 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,
+    trySbeEnginePushdownEligibleAggModel,
+} 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 sbeStatus = checkSbeStatus(db);
+if (sbeStatus === kSbeDisabled) {
+    jsTest.log.info("SBE is disabled, skipping test.");
+    quit();
+}
+
 const numRuns = 10;
 const numQueriesPerRun = 10;
 const jsTestLogExplain = false;
@@ -39,20 +49,30 @@ const correctnessProperty = createCorrectnessProperty(
     jsTestLogExplain,
 );
 
+const getAggModel = (foreignName) => {
+    if (sbeStatus === kSbeRestricted) {
+        return trySbeRestrictedPushdownEligibleAggModel(foreignName);
+    } else {
+        // Not distinguishing between trySbeEngine and featureFlagSbeFull here.
+        return trySbeEnginePushdownEligibleAggModel(foreignName);
+    }
+};
+
 // 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: getAggModel(controlName),
         numQueriesPerRun,
     }),
     numRuns,
+    undefined /*examples*/,
+    true /*sortArrays*/,
 );
 
 testProperty(
@@ -61,8 +81,10 @@ testProperty(
     // Experiment collection is foreign side.
     makeWorkloadModel({
         collModel: getCollectionModel(),
-        aggModel: sbePushdownEligibleAggModel(experimentName),
+        aggModel: getAggModel(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";
@@ -80,6 +80,28 @@ export function getSortArb(maxNumSortComponents = 1) {
 export const limitArb = fc.record({$limit: fc.integer({min: 1, max: 5})});
 export const skipArb = fc.record({$skip: fc.integer({min: 1, max: 5})});
 
+// Generates a standalone $unwind stage. The path field should reference an array field.
+// According to the schema, "array" is the array field, but we allow other fields as well
+// for flexibility (they may be arrays in some documents).
+export const unwindArb = fc
+    .record({
+        path: fieldArb,
+        preserveNullAndEmptyArrays: fc.boolean(),
+        includeArrayIndex: fc.boolean(),
+        indexFieldName: assignableFieldArb,
+    })
+    .map(({path, preserveNullAndEmptyArrays, includeArrayIndex, indexFieldName}) => {
+        const unwindSpec = {path: "$" + path};
+        if (preserveNullAndEmptyArrays) {
+            unwindSpec.preserveNullAndEmptyArrays = true;
+        }
+        if (includeArrayIndex) {
+            // includeArrayIndex specifies the field name to store the array index.
+            unwindSpec.includeArrayIndex = indexFieldName;
+        }
+        return {$unwind: unwindSpec};
+    });
+
 /*
  * Return the arbitraries for agg stages that are allowed given:
  *    - `allowOrs` for whether we allow $or in $match
@@ -129,11 +151,66 @@ 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()];
+    const stages = [groupArb, getEqLookupArb(foreignCollName), getMatchArb()];
+    return fc.array(oneof(...stages), {minLength: 1, maxLength: 6});
+}
+
+export function getTrySbeEnginePushdownEligibleAggPipelineArb(
+    foreignCollName,
+    {allowOrs = true, deterministicBag = true, allowedStages = [], isTS = false} = {},
+) {
+    // Not yet included, $window and $unwind.
+    const stages = [
+        groupArb,
+        getEqLookupArb(foreignCollName),
+        getEqLookupUnwindArb(foreignCollName),
+        getMatchArb(allowOrs),
+        simpleProjectArb,
+        computedProjectArb,
+        addFieldsConstArb,
+        addFieldsVarArb,
+        getSortArb(),
+    ];
+    if (!deterministicBag) {
+        stages.push(limitArb, skipArb);
+    }
+    // eqLookupUnwind returns a javascript array; flatten that here.
+    return fc.array(oneof(...stages), {minLength: 1, maxLength: 6}).map((item) => item.flat());
+}
+
+// According to findSbeCompatibleStagesForPushdown, SbeCompatibility::requiresSbeFull (featureFlagSbeFull)
+// is a superset of requiresTrySbe (trySbeEngine) and additionally allows:
+// - $unwind (SbeCompatibility::requiresSbeFull)
+// - $search/$searchMeta (SbeCompatibility::requiresSbeFull with featureFlagSearchInSbe)
+// Note: Currently, there are no aggregation model arbitraries defined for standalone $unwind or $search
+// stages, so this function includes all stages from trySbeEngine. If models for these stages are added
+// in the future, they should be included here.
+export function getSbeFullPushdownEligibleAggPipelineArb(
+    foreignCollName,
+    {allowOrs = true, deterministicBag = true, allowedStages = [], isTS = false} = {},
+) {
+    // Start with all stages from trySbeEngine (requiresTrySbe is a subset of requiresSbeFull)
+    const stages = [
+        groupArb,
+        getEqLookupArb(foreignCollName),
+        getEqLookupUnwindArb(foreignCollName),
+        getMatchArb(allowOrs),
+        simpleProjectArb,
+        computedProjectArb,
+        addFieldsConstArb,
+        addFieldsVarArb,
+        getSortArb(),
+    ];
+    if (!deterministicBag) {
+        stages.push(limitArb, skipArb);
+    }
+    // Add standalone $unwind (requires SbeCompatibility::requiresSbeFull)
+    stages.push(unwindArb);
+    // TODO: Add $search/$searchMeta arb when available
     // eqLookupUnwind returns a javascript array; flatten that here.
     return fc.array(oneof(...stages), {minLength: 1, maxLength: 6}).map((item) => item.flat());
 }

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,
     };
@@ -145,7 +145,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");
@@ -165,7 +165,7 @@ export function testProperty(propertyFn, namespaces, workloadModel, numRuns, exa
         fc.property(workloadModel, (workload) => {
             // Only return if the property passed or not. On failure,
             // `runProperty` is called again and more details are exposed.
-            const result = runProperty(propertyFn, namespaces, workload);
+            const result = runProperty(propertyFn, namespaces, workload, sortArrays);
             // If it failed for the first time, print that out so we have the first failure available
             // in case shrinking fails.
             if (!result.passed && alwaysPassed) {

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,7 @@ declare function _replMonitorStats()
  * @returns True if the result sets compare equal and false otherwise.
  */
 declare function _resultSetsEqualNormalized(a: object[], b: object[]): boolean
+declare function _resultSetsEqualUnorderedWithUnorderedArrays()
 declare function _resultSetsEqualUnordered()
 declare function _setShellFailPoint()
 declare function _srand()