diff options
| -rw-r--r-- | src/mongo/db/exec/sbe/SConscript | 1 | ||||
| -rw-r--r-- | src/mongo/db/exec/sbe/sbe_loop_join_test.cpp | 166 | ||||
| -rw-r--r-- | src/mongo/db/exec/sbe/stages/loop_join.h | 20 |
3 files changed, 179 insertions, 8 deletions
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript index c9395ca88e6..bba4be32267 100644 --- a/src/mongo/db/exec/sbe/SConscript +++ b/src/mongo/db/exec/sbe/SConscript @@ -177,6 +177,7 @@ env.CppUnitTest( 'sbe_hash_join_test.cpp', 'sbe_key_string_test.cpp', 'sbe_limit_skip_test.cpp', + 'sbe_loop_join_test.cpp', 'sbe_math_builtins_test.cpp', 'sbe_merge_join_test.cpp', 'sbe_mkobj_test.cpp', diff --git a/src/mongo/db/exec/sbe/sbe_loop_join_test.cpp b/src/mongo/db/exec/sbe/sbe_loop_join_test.cpp new file mode 100644 index 00000000000..95fc1156a7a --- /dev/null +++ b/src/mongo/db/exec/sbe/sbe_loop_join_test.cpp @@ -0,0 +1,166 @@ +/** + * Copyright (C) 2021-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +/** + * This file contains tests for sbe::LoopJoinStage. + */ + +#include "mongo/db/exec/sbe/sbe_plan_stage_test.h" +#include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/util/assert_util.h" + +namespace mongo::sbe { + +using LoopJoinStageTest = PlanStageTestFixture; + +TEST_F(LoopJoinStageTest, LoopJoinNoPredicate) { + auto ctx = makeCompileCtx(); + + // Build a scan for the outer loop. + auto [outerScanSlot, outerScanStage] = generateVirtualScan(BSON_ARRAY(1 << 2)); + + // Build a scan for the inner loop. + auto [innerScanSlot, innerScanStage] = generateVirtualScan(BSON_ARRAY(3 << 4 << 5)); + + // Build and prepare for execution loop join of the two scan stages. + auto loopJoin = makeS<LoopJoinStage>(std::move(outerScanStage), + std::move(innerScanStage), + makeSV(outerScanSlot) /*outerProjects*/, + makeSV() /*outerCorrelated*/, + nullptr /*predicate*/, + kEmptyPlanNodeId); + + prepareTree(ctx.get(), loopJoin.get()); + auto outer = loopJoin->getAccessor(*ctx, outerScanSlot); + auto inner = loopJoin->getAccessor(*ctx, innerScanSlot); + + // Expected output: cartesian product of the two scans. + std::vector<std::pair<int, int>> expected{{1, 3}, {1, 4}, {1, 5}, {2, 3}, {2, 4}, {2, 5}}; + int i = 0; + for (auto st = loopJoin->getNext(); st == PlanState::ADVANCED; st = loopJoin->getNext(), i++) { + ASSERT_LT(i, expected.size()); + + auto [outerTag, outerVal] = outer->copyOrMoveValue(); + assertValuesEqual(outerTag, outerVal, value::TypeTags::NumberInt32, expected[i].first); + + auto [innerTag, innerVal] = inner->copyOrMoveValue(); + assertValuesEqual(innerTag, innerVal, value::TypeTags::NumberInt32, expected[i].second); + } +} + +TEST_F(LoopJoinStageTest, LoopJoinConstTruePredicate) { + auto ctx = makeCompileCtx(); + + // Build a scan for the outer loop. + auto [outerScanSlot, outerScanStage] = generateVirtualScan(BSON_ARRAY(1 << 2)); + + // Build a scan for the inner loop. + auto [innerScanSlot, innerScanStage] = generateVirtualScan(BSON_ARRAY(3 << 4 << 5)); + + // Build and prepare for execution loop join of the two scan stages. + auto loopJoin = makeS<LoopJoinStage>( + std::move(outerScanStage), + std::move(innerScanStage), + makeSV(outerScanSlot) /*outerProjects*/, + makeSV() /*outerCorrelated*/, + stage_builder::makeConstant(sbe::value::TypeTags::Boolean, true) /*predicate*/, + kEmptyPlanNodeId); + prepareTree(ctx.get(), loopJoin.get()); + auto outer = loopJoin->getAccessor(*ctx, outerScanSlot); + auto inner = loopJoin->getAccessor(*ctx, innerScanSlot); + + // Expected output: cartesian product of the two scans. + std::vector<std::pair<int, int>> expected{{1, 3}, {1, 4}, {1, 5}, {2, 3}, {2, 4}, {2, 5}}; + int i = 0; + for (auto st = loopJoin->getNext(); st == PlanState::ADVANCED; st = loopJoin->getNext(), i++) { + ASSERT_LT(i, expected.size()); + + auto [outerTag, outerVal] = outer->copyOrMoveValue(); + assertValuesEqual(outerTag, outerVal, value::TypeTags::NumberInt32, expected[i].first); + + auto [innerTag, innerVal] = inner->copyOrMoveValue(); + assertValuesEqual(innerTag, innerVal, value::TypeTags::NumberInt32, expected[i].second); + } +} + +TEST_F(LoopJoinStageTest, LoopJoinConstFalsePredicate) { + auto ctx = makeCompileCtx(); + + // Build a scan for the outer loop. + auto [outerScanSlot, outerScanStage] = generateVirtualScan(BSON_ARRAY(1 << 2)); + + // Build a scan for the inner loop. + auto [innerScanSlot, innerScanStage] = generateVirtualScan(BSON_ARRAY(3 << 4 << 5)); + + // Build and prepare for execution loop join of the two scan stages. + auto loopJoin = makeS<LoopJoinStage>( + std::move(outerScanStage), + std::move(innerScanStage), + makeSV(outerScanSlot) /*outerProjects*/, + makeSV() /*outerCorrelated*/, + stage_builder::makeConstant(sbe::value::TypeTags::Boolean, false) /*predicate*/, + kEmptyPlanNodeId); + prepareTree(ctx.get(), loopJoin.get()); + + // Executing the stage should produce no results because of the predicate filter. + ASSERT(PlanState::IS_EOF == loopJoin->getNext()); +} + +TEST_F(LoopJoinStageTest, LoopJoinEqualityPredicate) { + auto ctx = makeCompileCtx(); + + // Build a scan for the outer loop. + auto [outerScanSlot, outerScanStage] = generateVirtualScan(BSON_ARRAY(1 << 2 << 3 << 4 << 1)); + + // Build a scan for the inner loop. + auto [innerScanSlot, innerScanStage] = generateVirtualScan(BSON_ARRAY(3 << 1 << 5 << 3 << 3)); + + // Build and prepare for execution loop join of the two scan stages. + auto predicate = makeE<EPrimBinary>( + EPrimBinary::eq, makeE<EVariable>(outerScanSlot), makeE<EVariable>(innerScanSlot)); + auto loopJoin = makeS<LoopJoinStage>(std::move(outerScanStage), + std::move(innerScanStage), + makeSV(outerScanSlot) /*outerProjects*/, + makeSV() /*outerCorrelated*/, + std::move(predicate), + kEmptyPlanNodeId); + prepareTree(ctx.get(), loopJoin.get()); + auto inner = loopJoin->getAccessor(*ctx, innerScanSlot); + + // Expected output should filter from the inner array elements that exist in the outer. + std::vector<int> expected{1, 3, 3, 3, 1}; + int i = 0; + for (auto st = loopJoin->getNext(); st == PlanState::ADVANCED; st = loopJoin->getNext(), i++) { + ASSERT_LT(i, expected.size()); + + auto [innerTag, innerVal] = inner->copyOrMoveValue(); + assertValuesEqual(innerTag, innerVal, value::TypeTags::NumberInt32, expected[i]); + } +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/loop_join.h b/src/mongo/db/exec/sbe/stages/loop_join.h index 22e3a93ffa2..2a7b0920751 100644 --- a/src/mongo/db/exec/sbe/stages/loop_join.h +++ b/src/mongo/db/exec/sbe/stages/loop_join.h @@ -84,19 +84,25 @@ private: return ret; } - // Set of variables coming from the outer side. + void openInner(); + + // Set of variables coming from the outer side. These are _not_ visible to the inner side, + // unless also added to '_outerCorrelated'. const value::SlotVector _outerProjects; + // Set of correlated variables from the outer side that are visible on the inner side. const value::SlotVector _outerCorrelated; - // If not set then this is a cross product. + + // Predicate to filter the joint set. If not set then the result is a cross product. + // Note: the predicate resolves the slots it's using through this stage's public accessors, + // meaning that if they are coming from the 'outer', they must be projected by the 'outer'. const std::unique_ptr<EExpression> _predicate; + vm::ByteCode _bytecode; + std::unique_ptr<vm::CodeFragment> _predicateCode; + // '_outerProjects' as a set (for faster checking of accessors, provided by the 'outer' child). value::SlotSet _outerRefs; - std::vector<value::SlotAccessor*> _correlatedAccessors; - std::unique_ptr<vm::CodeFragment> _predicateCode; - - vm::ByteCode _bytecode; bool _reOpenInner{false}; bool _outerGetNext{false}; LoopJoinStats _specificStats; @@ -104,7 +110,5 @@ private: // Tracks whether or not we're reading from the left child or the right child. // This is necessary for yielding. bool _isReadingLeftSide = false; - - void openInner(); }; } // namespace mongo::sbe |