SERVER-54236 Implement $count accumulator

3 files changed, 19 insertions, 1 deletions

diff --git a/src/mongo/db/pipeline/document_source_set_window_fields.cpp b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
index d8d691eff34..9a4f1e0314e 100644
--- a/src/mongo/db/pipeline/document_source_set_window_fields.cpp
+++ b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
@@ -316,6 +316,11 @@ DocumentSource::GetNextResult DocumentSourceInternalSetWindowFields::doGetNext()
         return curStat;
     }
     auto curDoc = curStat.getDocument();
+    if (_partitionBy) {
+        uassert(ErrorCodes::TypeMismatch,
+                "Cannot 'partitionBy' an expression of type array",
+                !_partitionBy->get()->evaluate(curDoc, &pExpCtx->variables).isArray());
+    }
     MutableDocument outDoc(curDoc);
     for (auto& output : _executableOutputs) {
         // Currently only support unbounded windows and run on the merging shard -- we don't need
diff --git a/src/mongo/db/pipeline/window_function/partition_iterator.cpp b/src/mongo/db/pipeline/window_function/partition_iterator.cpp
index 0e57c72152d..435cb7a765b 100644
--- a/src/mongo/db/pipeline/window_function/partition_iterator.cpp
+++ b/src/mongo/db/pipeline/window_function/partition_iterator.cpp
@@ -125,7 +125,10 @@ void PartitionIterator::getNextDocument() {
 
     auto doc = getNextRes.releaseDocument();
     if (_partitionExpr) {
-        auto curKey = _partitionExpr->evaluate(doc, nullptr);
+        auto curKey = _partitionExpr->evaluate(doc, &_expCtx->variables);
+        uassert(ErrorCodes::TypeMismatch,
+                "Cannot 'partitionBy' an expression of type array",
+                !curKey.isArray());
         if (_state == IteratorState::kNotInitialized) {
             _nextPartition = NextPartitionState{std::move(doc), std::move(curKey)};
             advanceToNextPartition();
diff --git a/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp b/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp
index 04b8ae787fe..97787282b10 100644
--- a/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp
+++ b/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp
@@ -167,5 +167,15 @@ TEST_F(PartitionIteratorTest, EmptyCollectionReturnsEOF) {
     ASSERT_ADVANCE_RESULT(PartitionIterator::AdvanceResult::kEOF, partIter.advance());
 }
 
+TEST_F(PartitionIteratorTest, PartitionByArrayErrs) {
+    const auto docs = std::deque<DocumentSource::GetNextResult>{Document{{"key", 1}},
+                                                                Document{fromjson("{key: [1]}")}};
+    const auto mock = DocumentSourceMock::createForTest(docs, getExpCtx());
+    auto key = ExpressionFieldPath::createPathFromString(
+        getExpCtx().get(), "key", getExpCtx()->variablesParseState);
+    auto partIter = PartitionIterator(getExpCtx().get(), mock.get(), *key);
+    ASSERT_DOCUMENT_EQ(docs[0].getDocument(), *partIter[0]);
+    ASSERT_THROWS_CODE(*partIter[1], AssertionException, ErrorCodes::TypeMismatch);
+}
 }  // namespace
 }  // namespace mongo