118 lines
3.8 KiB
JavaScript
118 lines
3.8 KiB
JavaScript
// Tests that stage rewrite optimizations for $graphLookup work correctly.
|
|
//
|
|
// This test makes assumptions about how the explain output will be formatted.
|
|
// @tags: [
|
|
// assumes_unsharded_collection,
|
|
// do_not_wrap_aggregations_in_facets,
|
|
// requires_pipeline_optimization,
|
|
// ]
|
|
(function() {
|
|
"use strict";
|
|
|
|
load('jstests/aggregation/extras/utils.js');
|
|
load("jstests/libs/analyze_plan.js");
|
|
load("jstests/libs/fixture_helpers.js");
|
|
|
|
const coll = db.graphlookup_rewrite;
|
|
coll.drop();
|
|
|
|
assert.commandWorked(coll.insertMany([
|
|
{"from": "a", "foo": 1},
|
|
{"from": "b", "to": "a", "foo": 2},
|
|
{"from": "c", "to": "b", "foo": 3},
|
|
{"from": "d", "to": "b", "foo": 4},
|
|
{"from": "e", "to": "c", "foo": 5},
|
|
{"from": "f", "to": "d", "foo": 6}
|
|
]));
|
|
|
|
const admin = db.getSiblingDB("admin");
|
|
|
|
const setPipelineOptimizationMode = (mode) => {
|
|
FixtureHelpers.runCommandOnEachPrimary(
|
|
{db: admin, cmdObj: {configureFailPoint: "disablePipelineOptimization", mode}});
|
|
};
|
|
|
|
// Get initial optimization mode.
|
|
const pipelineOptParameter = assert.commandWorked(
|
|
db.adminCommand({getParameter: 1, "failpoint.disablePipelineOptimization": 1}));
|
|
const oldMode =
|
|
pipelineOptParameter["failpoint.disablePipelineOptimization"].mode ? 'alwaysOn' : 'off';
|
|
|
|
function assertStagesAndOutput(
|
|
{pipeline = [], expectedStages = [], optimizedAwayStages = [], fieldsToSkip = [], msg = ""}) {
|
|
setPipelineOptimizationMode("off");
|
|
|
|
const explain = coll.explain().aggregate(pipeline);
|
|
const output = coll.aggregate(pipeline).toArray();
|
|
|
|
for (const stage of expectedStages) {
|
|
assert(aggPlanHasStage(explain, stage),
|
|
`${msg}: missing stage ${stage}: ${tojson(explain)}`);
|
|
}
|
|
for (const stage of optimizedAwayStages) {
|
|
assert(!aggPlanHasStage(explain, stage),
|
|
`${msg}: stage ${stage} not optimized away: ${tojson(explain)}`);
|
|
}
|
|
|
|
setPipelineOptimizationMode("alwaysOn");
|
|
|
|
const expectedOutput = coll.aggregate(pipeline).toArray();
|
|
assert(orderedArrayEq(output, expectedOutput, true, fieldsToSkip), msg);
|
|
}
|
|
|
|
const graphLookup = {
|
|
$graphLookup: {
|
|
from: "graphlookup_rewrite",
|
|
startWith: "$from",
|
|
connectFromField: "from",
|
|
connectToField: "to",
|
|
as: "out"
|
|
}
|
|
};
|
|
|
|
assertStagesAndOutput({
|
|
pipeline: [graphLookup, {$sort: {foo: 1}}],
|
|
expectedStages: ["SORT", "COLLSCAN", "$graphLookup"],
|
|
optimizedAwayStages: ["$sort"],
|
|
msg: "$graphLookup should swap with $sort if there is no internal $unwind"
|
|
});
|
|
|
|
assertStagesAndOutput({
|
|
pipeline: [graphLookup, {$limit: 100}],
|
|
expectedStages: ["LIMIT", "COLLSCAN", "$graphLookup"],
|
|
optimizedAwayStages: ["$limit"],
|
|
msg: "$graphLookup should swap with $limit if there is no internal $unwind"
|
|
});
|
|
|
|
assertStagesAndOutput({
|
|
pipeline: [graphLookup, {$skip: 100}],
|
|
expectedStages: ["SKIP", "COLLSCAN", "$graphLookup"],
|
|
optimizedAwayStages: ["$skip"],
|
|
msg: "$graphLookup should swap with $skip if there is no internal $unwind"
|
|
});
|
|
|
|
assertStagesAndOutput({
|
|
pipeline: [graphLookup, {$sort: {foo: 1}}, {$limit: 100}],
|
|
expectedStages: ["SORT", "COLLSCAN", "$graphLookup"],
|
|
optimizedAwayStages: ["LIMIT", "$limit"],
|
|
msg: "$graphLookup should swap with $limit and $sort, and $sort should absorb $limit if " +
|
|
"there is no internal $unwind"
|
|
});
|
|
|
|
assertStagesAndOutput({
|
|
pipeline: [graphLookup, {$sort: {out: 1}}],
|
|
expectedStages: ["COLLSCAN", "$graphLookup", "$sort"],
|
|
msg: "$graphLookup should not swap with $sort if sort uses fields created by $graphLookup"
|
|
});
|
|
|
|
assertStagesAndOutput({
|
|
pipeline: [graphLookup, {$unwind: "$out"}, {$sort: {foo: 1}}],
|
|
expectedStages: ["COLLSCAN", "$graphLookup", "$sort"],
|
|
msg: "$graphLookup with an internal $unwind should not swap with $sort",
|
|
fieldsToSkip: ["out"]
|
|
});
|
|
|
|
// Reset optimization mode.
|
|
setPipelineOptimizationMode(oldMode);
|
|
})();
|