SERVER-55239 Support 2dsphere index on time-series collection metadata fields

6 files changed, 82 insertions, 16 deletions

diff --git a/src/mongo/db/commands/run_aggregate.cpp b/src/mongo/db/commands/run_aggregate.cpp
index 7626cfc06cc..789581c7337 100644
--- a/src/mongo/db/commands/run_aggregate.cpp
+++ b/src/mongo/db/commands/run_aggregate.cpp
@@ -778,6 +778,9 @@ Status runAggregate(OperationContext* opCtx,
 
         pipeline->optimizePipeline();
 
+        constexpr bool alreadyOptimized = true;
+        pipeline->validateCommon(alreadyOptimized);
+
         // Check if the pipeline has a $geoNear stage, as it will be ripped away during the build
         // query executor phase below (to be replaced with a $geoNearCursorStage later during the
         // executor attach phase).
diff --git a/src/mongo/db/pipeline/document_source_geo_near.h b/src/mongo/db/pipeline/document_source_geo_near.h
index b4bc3c460e6..d9fbc6e669a 100644
--- a/src/mongo/db/pipeline/document_source_geo_near.h
+++ b/src/mongo/db/pipeline/document_source_geo_near.h
@@ -51,7 +51,7 @@ public:
 
     StageConstraints constraints(Pipeline::SplitState pipeState) const final {
         return {StreamType::kStreaming,
-                PositionRequirement::kFirst,
+                PositionRequirement::kFirstAfterOptimization,
                 HostTypeRequirement::kAnyShard,
                 DiskUseRequirement::kNoDiskUse,
                 FacetRequirement::kNotAllowed,
@@ -112,6 +112,13 @@ public:
     }
 
     /**
+     * Set the field over which to apply the "near" predicate.
+     */
+    void setKeyField(const FieldPath& newPath) {
+        keyFieldPath = newPath;
+    }
+
+    /**
      * A scaling factor to apply to the distance, if specified by the user.
      */
     boost::optional<double> getDistanceMultiplier() const {
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
index 228fe86b4ce..6d8b0432259 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
@@ -43,6 +43,7 @@
 #include "mongo/db/matcher/expression_algo.h"
 #include "mongo/db/matcher/expression_internal_expr_comparison.h"
 #include "mongo/db/pipeline/document_source_add_fields.h"
+#include "mongo/db/pipeline/document_source_geo_near.h"
 #include "mongo/db/pipeline/document_source_group.h"
 #include "mongo/db/pipeline/document_source_match.h"
 #include "mongo/db/pipeline/document_source_project.h"
@@ -279,6 +280,7 @@ bool fieldIsComputed(BucketSpec spec, std::string field) {
                 expression::isPathPrefixOf(s, field);
         });
 }
+
 }  // namespace
 
 DocumentSourceInternalUnpackBucket::DocumentSourceInternalUnpackBucket(
@@ -960,6 +962,41 @@ Pipeline::SourceContainer::iterator DocumentSourceInternalUnpackBucket::doOptimi
         }
     }
 
+    // Attempt to push geoNear on the metaField past $_internalUnpackBucket.
+    if (auto nextNear = dynamic_cast<DocumentSourceGeoNear*>(std::next(itr)->get())) {
+        // Currently we only support geo indexes on the meta field, and we enforce this by
+        // requiring the key field to be set so we can check before we try to look up indexes.
+        auto keyField = nextNear->getKeyField();
+        uassert(5892921,
+                "Must specify 'key' option for $geoNear on a time-series collection",
+                keyField);
+
+        auto metaField = _bucketUnpacker.bucketSpec().metaField;
+        uassert(
+            4581294,
+            "Must specify part of metadata field as 'key' for $geoNear on a time-series collection",
+            metaField && *metaField == keyField->front());
+
+        // Currently we do not support query for $geoNear on a bucket
+        uassert(
+            1938439,
+            "Must not specify 'query' for $geoNear on a time-series collection; use $match instead",
+            nextNear->getQuery().binaryEqual(BSONObj()));
+
+        // Make sure we actually re-write the key field for the buckets collection so we can
+        // locate the index.
+        static const FieldPath baseMetaFieldPath{timeseries::kBucketMetaFieldName};
+        nextNear->setKeyField(keyField->getPathLength() > 1
+                                  ? baseMetaFieldPath.concat(keyField->tail())
+                                  : baseMetaFieldPath);
+
+        // Save the source, remove it, and then push it down.
+        auto source = *std::next(itr);
+        container->erase(std::next(itr));
+        container->insert(itr, source);
+        return std::prev(itr) == container->begin() ? std::prev(itr) : std::prev(std::prev(itr));
+    }
+
     // Attempt to map predicates on bucketed fields to predicates on the control field.
     if (auto nextMatch = dynamic_cast<DocumentSourceMatch*>(std::next(itr)->get());
         nextMatch && !_triedBucketLevelFieldsPredicatesPushdown) {
diff --git a/src/mongo/db/pipeline/pipeline.cpp b/src/mongo/db/pipeline/pipeline.cpp
index cf298b0a492..5219a3ab9b1 100644
--- a/src/mongo/db/pipeline/pipeline.cpp
+++ b/src/mongo/db/pipeline/pipeline.cpp
@@ -179,7 +179,8 @@ std::unique_ptr<Pipeline, PipelineDeleter> Pipeline::parse(
     }
 
     // Next run through the common validation rules that apply to every pipeline.
-    pipeline->validateCommon();
+    constexpr bool alreadyOptimized = false;
+    pipeline->validateCommon(alreadyOptimized);
 
     pipeline->stitch();
     return pipeline;
@@ -190,12 +191,13 @@ std::unique_ptr<Pipeline, PipelineDeleter> Pipeline::create(
     std::unique_ptr<Pipeline, PipelineDeleter> pipeline(new Pipeline(std::move(stages), expCtx),
                                                         PipelineDeleter(expCtx->opCtx));
 
-    pipeline->validateCommon();
+    constexpr bool alreadyOptimized = false;
+    pipeline->validateCommon(alreadyOptimized);
     pipeline->stitch();
     return pipeline;
 }
 
-void Pipeline::validateCommon() const {
+void Pipeline::validateCommon(bool alreadyOptimized) const {
     size_t i = 0;
 
     uassert(ErrorCodes::FailedToParse,
@@ -209,8 +211,14 @@ void Pipeline::validateCommon() const {
         // Verify that all stages adhere to their PositionRequirement constraints.
         uassert(40602,
                 str::stream() << stage->getSourceName()
-                              << " is only valid as the first stage in a pipeline.",
+                              << " is only valid as the first stage in a pipeline",
                 !(constraints.requiredPosition == PositionRequirement::kFirst && i != 0));
+        uassert(40603,
+                str::stream() << stage->getSourceName()
+                              << " is only valid as the first stage in an optimized pipeline",
+                !(alreadyOptimized &&
+                  constraints.requiredPosition == PositionRequirement::kFirstAfterOptimization &&
+                  i != 0));
 
         auto matchStage = dynamic_cast<DocumentSourceMatch*>(stage.get());
         uassert(17313,
@@ -651,6 +659,9 @@ std::unique_ptr<Pipeline, PipelineDeleter> Pipeline::makePipeline(
         pipeline->optimizePipeline();
     }
 
+    constexpr bool alreadyOptimized = true;
+    pipeline->validateCommon(alreadyOptimized);
+
     if (opts.attachCursorSource) {
         pipeline = expCtx->mongoProcessInterface->attachCursorSourceToPipeline(
             pipeline.release(), opts.shardTargetingPolicy, std::move(opts.readConcern));
diff --git a/src/mongo/db/pipeline/pipeline.h b/src/mongo/db/pipeline/pipeline.h
index 1bff0d2e64c..eded6e32740 100644
--- a/src/mongo/db/pipeline/pipeline.h
+++ b/src/mongo/db/pipeline/pipeline.h
@@ -342,6 +342,16 @@ public:
         StringData targetStageName, std::function<bool(const DocumentSource* const)> predicate);
 
     /**
+     * Performs common validation for top-level or facet pipelines. Throws if the pipeline is
+     * invalid.
+     *
+     * Includes checking for illegal stage positioning. For example, $out must be at the end, while
+     * a $match stage with a text query must be at the start. Note that this method accepts an
+     * initial source as the first stage, which is illegal for $facet pipelines.
+     */
+    void validateCommon(bool alreadyOptimized) const;
+
+    /**
      * PipelineD is a "sister" class that has additional functionality for the Pipeline. It exists
      * because of linkage requirements. Pipeline needs to function in mongod and mongos. PipelineD
      * contains extra functionality required in mongod, and which can't appear in mongos because the
@@ -373,16 +383,6 @@ private:
     static void stitch(SourceContainer* container);
 
     /**
-     * Performs common validation for top-level or facet pipelines. Throws if the pipeline is
-     * invalid.
-     *
-     * Includes checking for illegal stage positioning. For example, $out must be at the end, while
-     * a $match stage with a text query must be at the start. Note that this method accepts an
-     * initial source as the first stage, which is illegal for $facet pipelines.
-     */
-    void validateCommon() const;
-
-    /**
      * Returns Status::OK if the pipeline can run on mongoS, or an error with a message explaining
      * why it cannot.
      */
diff --git a/src/mongo/db/pipeline/stage_constraints.h b/src/mongo/db/pipeline/stage_constraints.h
index b520dfdbbdb..e955039feb2 100644
--- a/src/mongo/db/pipeline/stage_constraints.h
+++ b/src/mongo/db/pipeline/stage_constraints.h
@@ -52,7 +52,15 @@ struct StageConstraints {
      * A PositionRequirement stipulates what specific position the stage must occupy within the
      * pipeline, if any.
      */
-    enum class PositionRequirement { kNone, kFirst, kLast };
+    enum class PositionRequirement {
+        kNone,
+        kFirst,
+        // User can specify this stage anywhere, as long as the system can move the stage to be
+        // first. If pipeline optimization is disabled, then the stage must be first prior to
+        // optimization.
+        kFirstAfterOptimization,
+        kLast
+    };
 
     /**
      * A HostTypeRequirement defines where this stage is permitted to be executed when the