mongo/jstests/aggregation/sources/setWindowFields/optimize.js

/**
 * Test that redundant sorts are removed/swapped.
 *
 * @tags: [
 *   assumes_unsharded_collection,
 *   do_not_wrap_aggregations_in_facets,
 *   requires_pipeline_optimization,
 *   # This feature flag adjusts the desugaring a bit - requesting 'outputSortKeyMetadata' from the
 *   # $sort stage.
 *   featureFlagRankFusionBasic,
 *   requires_fcv_81,
 * ]
 */
import {aggPlanHasStage, getAggPlanStages} from "jstests/libs/query/analyze_plan.js";

// Find how many stages of the plan are 'stageName'.
function numberOfStages(explain, stageName) {
    return getAggPlanStages(explain, stageName).length;
}

const coll = db[jsTestName()];
coll.drop();

assert.commandWorked(
    coll.insert([
        {a: {b: 3, c: 3}, b: 2, c: 1},
        {a: {b: 3, c: 3}, b: 4, c: 1},
        {a: {b: 1, c: 1}, b: 4, c: 1},
        {a: {b: 2, c: 2}, b: 1, c: 1},
        {a: {b: 1, c: 1}, b: 4, c: 1},
        {a: {b: 3, c: 3}, b: 2, c: 1},
        {a: {b: 2, c: 2}, b: 1, c: 1},
        {a: {b: 2, c: 2}, b: 1, c: 1},
        {a: {b: 2, c: 2}, b: 3, c: 1},
        {a: {b: 1, c: 1}, b: 4, c: 1},
    ]),
);

const explain1 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1, b: 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: 1}},
]);

// Redundant $sort should be removed.
assert.eq(1, numberOfStages(explain1, "$sort"), explain1);
// We keep the more specific sort.
assert.docEq(
    [{$sort: {sortKey: {a: 1, b: 1}, outputSortKeyMetadata: true}}],
    getAggPlanStages(explain1, "$sort"),
    explain1,
);

const explain2 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1, b: 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: -1}},
]);

// $sort is not redundant, should not be removed.
assert.eq(2, numberOfStages(explain2, "$sort"), explain2);

const explain3 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: 1, b: -1}},
]);

// $sort should be swapped with $_internalSetWindowFields, and the extra one removed.
assert.eq(1, numberOfStages(explain3, "$sort"), explain3);
// The sort we keep should be the more specific one.
assert.docEq(
    [{$sort: {sortKey: {a: 1, b: -1}, outputSortKeyMetadata: true}}],
    getAggPlanStages(explain3, "$sort"),
    explain3,
);

const explain4 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1},
            output: {a: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: 1}},
]);

// Sort field is modified, can't be swapped or removed.
assert.eq(2, numberOfStages(explain4, "$sort"), explain4);

const explain5 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1},
            output: {"a.b": {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: 1}},
]);

// Sort field is modified, can't be swapped or removed.
assert.eq(2, numberOfStages(explain5, "$sort"), explain5);

const explain6 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: 1}},
    {$limit: 5},
]);

// $sort + $limit should not be merged.
assert(aggPlanHasStage(explain6, "$limit"), explain6);
// The $limit should not prevent us from removing the redundant $sort.
assert.eq(1, numberOfStages(explain6, "$sort"), explain6);

const explain7 = coll
    .explain()
    .aggregate([
        {$_internalInhibitOptimization: {}},
        {$setWindowFields: {partitionBy: "$a.b", output: {sum: {$sum: "$c"}}}},
        {$sort: {"a.b": 1}},
    ]);

// Sort should be removed if sorting and partitioning on same field.
assert.eq(1, numberOfStages(explain7, "$sort"), explain7);

const explain8 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            partitionBy: "$a.c",
            sortBy: {"a.b": 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {"a.b": 1}},
]);

// Sort should not be removed since primary sort field is "a.c".
assert.eq(2, numberOfStages(explain8, "$sort"), explain8);

const explain9 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            partitionBy: "$a.c",
            sortBy: {"a.b": 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {"a.c": 1}},
]);

// Sort should be removed since primary sort field is "a.c".
assert.eq(1, numberOfStages(explain9, "$sort"), explain9);

// Multiple redundant sorts are dropped.
const explain10 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1, b: 1, c: 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: 1, b: 1}},
    {$sort: {a: 1}},
    {$sort: {a: 1, b: 1}},
]);
assert.eq(1, numberOfStages(explain10, "$sort"), explain10);
assert.docEq(
    [{$sort: {sortKey: {a: 1, b: 1, c: 1}, outputSortKeyMetadata: true}}],
    getAggPlanStages(explain10, "$sort"),
    explain10,
);

// Multiple compatible sorts are pushed down.
const explain11 = coll.explain().aggregate([
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: 1, b: 1, c: 1}},
    {$sort: {a: 1, b: 1}},
    {$sort: {a: 1, b: 1, c: 1}},
]);
assert.eq(1, numberOfStages(explain11, "$sort"), explain11);
assert.docEq(
    [{$sort: {sortKey: {a: 1, b: 1, c: 1}, outputSortKeyMetadata: true}}],
    getAggPlanStages(explain11, "$sort"),
    explain11,
);

// An incompatible $meta sort should not be dropped or pushed down.
coll.createIndex({"$**": "text"});
const explain12 = coll.explain().aggregate([
    {$match: {$text: {$search: "hi"}}},
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: 1},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: {$meta: "textScore"}}},
]);
assert.eq(2, numberOfStages(explain12, "$sort"), explain12);

// For now, we don't handle $meta at all: it won't be optimized even if it's compatible.
const explain13 = coll.explain().aggregate([
    {$match: {$text: {$search: "hi"}}},
    {$_internalInhibitOptimization: {}},
    {
        $setWindowFields: {
            sortBy: {a: {$meta: "textScore"}},
            output: {sum: {$sum: "$c", window: {documents: ["unbounded", "current"]}}},
        },
    },
    {$sort: {a: {$meta: "textScore"}}},
]);
assert.eq(2, numberOfStages(explain13, "$sort"), explain13);