SERVER-41964 find() execution code

3 files changed, 248 insertions, 8 deletions

diff --git a/src/mongo/db/exec/find_projection_executor.cpp b/src/mongo/db/exec/find_projection_executor.cpp
index 8c44b4d8793..af40bd10114 100644
--- a/src/mongo/db/exec/find_projection_executor.cpp
+++ b/src/mongo/db/exec/find_projection_executor.cpp
@@ -35,6 +35,20 @@ namespace mongo {
 namespace projection_executor {
 namespace {
 /**
+ * Holds various parameters required to apply a $slice projection. Populated from the arguments
+ * to 'applySliceProjection()'.
+ */
+struct SliceParams {
+    const FieldPath& path;
+    const boost::optional<int> skip;
+    const int limit;
+};
+
+Value applySliceProjectionHelper(const Document& input,
+                                 const SliceParams& params,
+                                 size_t fieldPathIndex);
+
+/**
  * Extracts an element from the array 'arr' at position 'elemIndex'. The 'elemIndex' string
  * parameter must hold a value which can be converted to an unsigned integer. If 'elemIndex' is not
  * within array boundaries, an empty Value is returned.
@@ -44,6 +58,107 @@ Value extractArrayElement(const Value& arr, const std::string& elemIndex) {
     invariant(index);
     return arr[*index];
 }
+
+/**
+ * Returns a portion of the 'array', skipping a number of elements as indicated by the 'skip'
+ * parameter, either from the beginning of the array (if 'skip' is positive), or from the end
+ * of the array (if 'skip' is negative). The 'limit' indicates the number of items to return.
+ * If 'limit' is negative, the last 'limit' number of items in the array are returned.
+ *
+ * If the 'skip' is specified, the 'limit' cannot be negative.
+ */
+Value sliceArray(const std::vector<Value>& array, boost::optional<int> skip, int limit) {
+    auto start = 0ll;
+    auto forward = 0ll;
+    const long long len = array.size();
+
+    if (!skip) {
+        if (limit < 0) {
+            start = std::max(0ll, len + limit);
+            forward = len - start;
+        } else {
+            forward = std::min(len, static_cast<long long>(limit));
+        }
+    } else {
+        // We explicitly disallow a negative limit when skip is specified.
+        invariant(limit >= 0);
+
+        if (*skip < 0) {
+            start = std::max(0ll, len + *skip);
+            forward = std::min(len - start, static_cast<long long>(limit));
+        } else {
+            start = std::min(len, static_cast<long long>(*skip));
+            forward = std::min(len - start, static_cast<long long>(limit));
+        }
+    }
+
+    invariant(start + forward >= 0);
+    invariant(start + forward <= len);
+    return Value{std::vector<Value>(array.cbegin() + start, array.cbegin() + start + forward)};
+}
+
+/**
+ * Applies a $slice projection to the array at the given 'params.path'. For each array element,
+ * recursively calls 'applySliceProjectionHelper' if the element is a Document, storing the result
+ * in the output array, otherwise just stores the element in the output unmodified.
+ *
+ * Note we do not expand arrays within arrays this way. For example, {a: [[{b: 1}]]} has no values
+ * on the path "a.b", but {a: [{b: 1}]} does, so nested arrays are stored within the output array
+ * as regular values.
+ */
+Value applySliceProjectionToArray(const std::vector<Value>& array,
+                                  const SliceParams& params,
+                                  size_t fieldPathIndex) {
+    std::vector<Value> output;
+    output.reserve(array.size());
+
+    for (const auto& elem : array) {
+        output.push_back(
+            elem.getType() == BSONType::Object
+                ? applySliceProjectionHelper(elem.getDocument(), params, fieldPathIndex)
+                : elem);
+    }
+
+    return Value{output};
+}
+
+/**
+ * This is a helper function which implements the $slice projection. The strategy for applying a
+ * $slice projection is as follows:
+ *     * Pick the current path component from the current 'params.path' and store the value from the
+ *       'input' doc at this sub-path in 'val'.
+ *     * If 'val' is an array and we're at the last component in the 'params.path' - slice the array
+ *       and exit recursion, otherwise recursively apply the $slice projection to each element
+ *       in the array, and store the result in 'val'.
+ *     * If the field value is a document, apply the $slice projection to this document, and store
+ *       the result in 'val'.
+ *     * Store the computed 'val' in the 'output' document under the current field name.
+ */
+Value applySliceProjectionHelper(const Document& input,
+                                 const SliceParams& params,
+                                 size_t fieldPathIndex) {
+    invariant(fieldPathIndex < params.path.getPathLength());
+
+    auto fieldName = params.path.getFieldName(fieldPathIndex);
+    Value val{input[fieldName]};
+
+    switch (val.getType()) {
+        case BSONType::Array:
+            val = (fieldPathIndex + 1 == params.path.getPathLength())
+                ? sliceArray(val.getArray(), params.skip, params.limit)
+                : applySliceProjectionToArray(val.getArray(), params, fieldPathIndex + 1);
+            break;
+        case BSONType::Object:
+            val = applySliceProjectionHelper(val.getDocument(), params, fieldPathIndex + 1);
+            break;
+        default:
+            break;
+    }
+
+    MutableDocument output(input);
+    output.setField(fieldName, val);
+    return Value{output.freeze()};
+}
 }  // namespace
 
 void applyPositionalProjection(const Document& input,
@@ -131,5 +246,15 @@ void applyElemMatchProjection(const Document& input,
 
     output->setField(path.fullPath(), Value{std::vector<Value>{matchingElem}});
 }
+
+Document applySliceProjection(const Document& input,
+                              const FieldPath& path,
+                              boost::optional<int> skip,
+                              int limit) {
+    auto params = SliceParams{path, skip, limit};
+    auto val = applySliceProjectionHelper(input, params, 0);
+    invariant(val.getType() == BSONType::Object);
+    return val.getDocument();
+}
 }  // namespace projection_executor
 }  // namespace mongo
diff --git a/src/mongo/db/exec/find_projection_executor.h b/src/mongo/db/exec/find_projection_executor.h
index 70075de15eb..5e9b22bca26 100644
--- a/src/mongo/db/exec/find_projection_executor.h
+++ b/src/mongo/db/exec/find_projection_executor.h
@@ -79,5 +79,29 @@ void applyElemMatchProjection(const Document& input,
                               const MatchExpression& matchExpr,
                               const FieldPath& path,
                               MutableDocument* output);
+/**
+ * Applies a $slice projection on the array at the given 'path' on the 'input' document. The applied
+ * projection is returned as a Document. The 'skip' parameter indicates the number of items in the
+ * array to skip and the 'limit' indicates the number of items to return.
+ *
+ * If any of the 'skip' or 'limit' is negative, it is applied to the last items in the array (e.g.,
+ * {$slice: -5} returns the last five items in array, and {$slice: [-20, 10]} returns 10 items,
+ * beginning with the item that is 20th from the last item of the array.
+ *
+ * If the 'skip' is specified, the 'limit' cannot be negative.
+ *
+ * For example, given:
+ *
+ *   - the 'input' document {foo: [{bar: 1}, {bar: 2}, {bar: 3}, {bar: 4}]}
+ *   - the path of 'foo'
+ *   - the 'skip' of -3 and 'limit' of 2
+ *
+ * The resulting document will contain the following element: {foo: [{bar: 2}, {bar: 3}]}.
+ */
+Document applySliceProjection(const Document& input,
+                              const FieldPath& path,
+                              boost::optional<int> skip,
+                              int limit);
+
 }  // namespace projection_executor
 }  // namespace mongo
diff --git a/src/mongo/db/exec/find_projection_executor_test.cpp b/src/mongo/db/exec/find_projection_executor_test.cpp
index 113b59daf45..ffbf556840f 100644
--- a/src/mongo/db/exec/find_projection_executor_test.cpp
+++ b/src/mongo/db/exec/find_projection_executor_test.cpp
@@ -33,15 +33,16 @@
 #include "mongo/db/matcher/expression_parser.h"
 #include "mongo/db/pipeline/document_value_test_util.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
+#include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
 namespace projection_executor {
 namespace positional_projection_tests {
-Document applyPositional(const BSONObj& match,
-                         const std::string& path,
-                         const Document& input,
-                         const Document& output = {}) {
+auto applyPositional(const BSONObj& match,
+                     const std::string& path,
+                     const Document& input,
+                     const Document& output = {}) {
     MutableDocument doc(output);
     boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
     auto matchExpr = uassertStatusOK(MatchExpressionParser::parse(match, expCtx));
@@ -123,10 +124,10 @@ TEST(PositionalProjection, CanMergeWithExistingFieldsInOutputDocument) {
 }  // namespace positional_projection_tests
 
 namespace elem_match_projection_tests {
-Document applyElemMatch(const BSONObj& match,
-                        const std::string& path,
-                        const Document& input,
-                        const Document& output = {}) {
+auto applyElemMatch(const BSONObj& match,
+                    const std::string& path,
+                    const Document& input,
+                    const Document& output = {}) {
     MutableDocument doc(output);
     boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
     auto matchObj = BSON(path << BSON("$elemMatch" << match));
@@ -207,5 +208,95 @@ TEST(ElemMatchProjection, RemovesFieldFromOutputDocumentIfItContainsNumericSubfi
                        applyElemMatch(fromjson("{$gt: 2}"), "foo", doc, doc));
 }
 }  // namespace elem_match_projection_tests
+
+namespace slice_projection_tests {
+DEATH_TEST(SliceProjection,
+           ShouldFailIfNegativeLimitSpecifiedWithPositiveSkip,
+           "Invariant failure limit >= 0") {
+    auto doc = Document{fromjson("{a: [1,2,3,4]}")};
+    applySliceProjection(doc, "a", 1, -1);
+}
+
+DEATH_TEST(SliceProjection,
+           ShouldFailIfNegativeLimitSpecifiedWithNegativeSkip,
+           "Invariant failure limit >= 0") {
+    auto doc = Document{fromjson("{a: [1,2,3,4]}")};
+    applySliceProjection(doc, "a", -1, -1);
+}
+
+TEST(SliceProjection, CorrectlyProjectsSimplePath) {
+    auto doc = Document{fromjson("{a: [1,2,3,4]}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [1,2,3]}")},
+                       applySliceProjection(doc, "a", boost::none, 3));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [2,3,4]}")},
+                       applySliceProjection(doc, "a", boost::none, -3));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [2]}")}, applySliceProjection(doc, "a", -3, 1));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [2,3,4]}")}, applySliceProjection(doc, "a", -3, 4));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [4]}")}, applySliceProjection(doc, "a", 3, 1));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [1,2,3,4]}")}, applySliceProjection(doc, "a", -5, 5));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: []}")}, applySliceProjection(doc, "a", 5, 2));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [1,2]}")}, applySliceProjection(doc, "a", -5, 2));
+    ASSERT_DOCUMENT_EQ(
+        Document{fromjson("{a: [1,2,3,4]}")},
+        applySliceProjection(doc, "a", boost::none, std::numeric_limits<int>::max()));
+    ASSERT_DOCUMENT_EQ(
+        Document{fromjson("{a: [1,2,3,4]}")},
+        applySliceProjection(doc, "a", boost::none, std::numeric_limits<int>::min()));
+    ASSERT_DOCUMENT_EQ(
+        Document{fromjson("{a: [1,2,3,4]}")},
+        applySliceProjection(
+            doc, "a", std::numeric_limits<int>::min(), std::numeric_limits<int>::max()));
+    ASSERT_DOCUMENT_EQ(
+        Document{fromjson("{a: []}")},
+        applySliceProjection(
+            doc, "a", std::numeric_limits<int>::max(), std::numeric_limits<int>::max()));
+
+    doc = Document{fromjson("{a: [{b: 1, c: 1}, {b: 2, c: 2}, {b: 3, c: 3}], d: 2}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: 1, c: 1}], d: 2}")},
+                       applySliceProjection(doc, "a", boost::none, 1));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: 3, c: 3}], d: 2}")},
+                       applySliceProjection(doc, "a", boost::none, -1));
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: 2, c: 2}], d: 2}")},
+                       applySliceProjection(doc, "a", 1, 1));
+
+    doc = Document{fromjson("{a: 1}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: 1}")},
+                       applySliceProjection(doc, "a", boost::none, 2));
+}
+
+TEST(SliceProjection, CorrectlyProjectsDottedPath) {
+    auto doc = Document{fromjson("{a: {b: [1,2,3], c: 1}}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: {b: [1,2], c: 1}}")},
+                       applySliceProjection(doc, "a.b", boost::none, 2));
+
+    doc = Document{fromjson("{a: {b: [1,2,3], c: 1}, d: 1}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: {b: [1,2], c: 1}, d: 1}")},
+                       applySliceProjection(doc, "a.b", boost::none, 2));
+
+    doc = Document{fromjson("{a: {b: [[1,2], [3,4], [5,6]]}}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: {b: [[1,2], [3,4]]}}")},
+                       applySliceProjection(doc, "a.b", boost::none, 2));
+
+    doc = Document{fromjson("{a: [{b: {c: [1,2,3,4]}}, {b: {c: [5,6,7,8]}}], d: 1}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: {c: [4]}}, {b: {c: [8]}}], d: 1}")},
+                       applySliceProjection(doc, "a.b.c", -1, 2));
+
+    doc = Document{fromjson("{a: {b: 1}}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: {b: 1}}")},
+                       applySliceProjection(doc, "a.b", boost::none, 2));
+
+    doc = Document{fromjson("{a: [{b: [1,2,3], c: 1}]}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: [3], c: 1}]}")},
+                       applySliceProjection(doc, "a.b", boost::none, -1));
+
+    doc = Document{fromjson("{a: [{b: [1,2,3], c: 4}, {b: [5,6,7], c: 8}]}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: [3], c: 4}, {b: [7], c: 8}]}")},
+                       applySliceProjection(doc, "a.b", boost::none, -1));
+
+    doc = Document{fromjson("{a: [{b: [{x:1, c: [1, 2]}, {y: 1, c: [3, 4]}]}], z: 1}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: [{x:1, c: [1]}, {y: 1, c: [3]}]}], z: 1}")},
+                       applySliceProjection(doc, "a.b.c", boost::none, 1));
+}
+}  // namespace slice_projection_tests
 }  // namespace projection_executor
 }  // namespace mongo