SERVER-73602 Add the stub of the TIMESERIES_WRITE query stage

9 files changed, 170 insertions, 0 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 7a5931d98db..d198bd8e5e3 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -1354,6 +1354,7 @@ env.Library(
         'exec/subplan.cpp',
         'exec/text_match.cpp',
         'exec/text_or.cpp',
+        'exec/timeseries_write.cpp',
         'exec/trial_period_utils.cpp',
         'exec/trial_stage.cpp',
         'exec/unpack_timeseries_bucket.cpp',
diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index 4801769c441..de30eef1a55 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -1161,6 +1161,26 @@ struct UnpackTimeseriesBucketStats final : public SpecificStats {
     size_t nBucketsUnpacked = 0u;
 };
 
+struct TimeseriesWriteStats final : public SpecificStats {
+    std::unique_ptr<SpecificStats> clone() const final {
+        return std::make_unique<TimeseriesWriteStats>(*this);
+    }
+
+    uint64_t estimateObjectSizeInBytes() const {
+        return sizeof(*this);
+    }
+
+    void acceptVisitor(PlanStatsConstVisitor* visitor) const final {
+        visitor->visit(this);
+    }
+
+    void acceptVisitor(PlanStatsMutableVisitor* visitor) final {
+        visitor->visit(this);
+    }
+
+    size_t measurementsDeleted = 0u;
+};
+
 struct SampleFromTimeseriesBucketStats final : public SpecificStats {
     std::unique_ptr<SpecificStats> clone() const final {
         return std::make_unique<SampleFromTimeseriesBucketStats>(*this);
diff --git a/src/mongo/db/exec/plan_stats_visitor.h b/src/mongo/db/exec/plan_stats_visitor.h
index 175bbf1e91b..953eaa85921 100644
--- a/src/mongo/db/exec/plan_stats_visitor.h
+++ b/src/mongo/db/exec/plan_stats_visitor.h
@@ -80,6 +80,7 @@ struct DocumentSourceLookupStats;
 struct UnionWithStats;
 struct DocumentSourceFacetStats;
 struct UnpackTimeseriesBucketStats;
+struct TimeseriesWriteStats;
 struct SampleFromTimeseriesBucketStats;
 
 /**
@@ -141,6 +142,7 @@ public:
     virtual void visit(tree_walker::MaybeConstPtr<IsConst, UnionWithStats> stats) = 0;
     virtual void visit(tree_walker::MaybeConstPtr<IsConst, DocumentSourceFacetStats> stats) = 0;
     virtual void visit(tree_walker::MaybeConstPtr<IsConst, UnpackTimeseriesBucketStats> stats) = 0;
+    virtual void visit(tree_walker::MaybeConstPtr<IsConst, TimeseriesWriteStats> stats) = 0;
     virtual void visit(
         tree_walker::MaybeConstPtr<IsConst, SampleFromTimeseriesBucketStats> stats) = 0;
 };
@@ -196,6 +198,7 @@ struct PlanStatsVisitorBase : public PlanStatsVisitor<IsConst> {
     void visit(tree_walker::MaybeConstPtr<IsConst, UnionWithStats> stats) override {}
     void visit(tree_walker::MaybeConstPtr<IsConst, DocumentSourceFacetStats> stats) override {}
     void visit(tree_walker::MaybeConstPtr<IsConst, UnpackTimeseriesBucketStats> stats) override {}
+    void visit(tree_walker::MaybeConstPtr<IsConst, TimeseriesWriteStats> stats) override {}
     void visit(
         tree_walker::MaybeConstPtr<IsConst, SampleFromTimeseriesBucketStats> stats) override {}
 };
diff --git a/src/mongo/db/exec/timeseries_write.cpp b/src/mongo/db/exec/timeseries_write.cpp
new file mode 100644
index 00000000000..e3d109866ef
--- /dev/null
+++ b/src/mongo/db/exec/timeseries_write.cpp
@@ -0,0 +1,53 @@
+/**
+ *    Copyright (C) 2023-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.
+ */
+
+#include "mongo/db/exec/timeseries_write.h"
+
+namespace mongo {
+
+const char* TimeseriesWriteStage::kStageType = "TS_WRITE";
+
+bool TimeseriesWriteStage::isEOF() {
+    return true;
+}
+
+std::unique_ptr<PlanStageStats> TimeseriesWriteStage::getStats() {
+    _commonStats.isEOF = isEOF();
+    auto ret = std::make_unique<PlanStageStats>(_commonStats, stageType());
+    ret->specific = std::make_unique<TimeseriesWriteStats>(_specificStats);
+    for (const auto& child : _children) {
+        ret->children.emplace_back(child->getStats());
+    }
+    return ret;
+}
+
+PlanStage::StageState TimeseriesWriteStage::doWork(WorkingSetID* out) {
+    return PlanStage::IS_EOF;
+}
+}  // namespace mongo
diff --git a/src/mongo/db/exec/timeseries_write.h b/src/mongo/db/exec/timeseries_write.h
new file mode 100644
index 00000000000..f55770c2c1c
--- /dev/null
+++ b/src/mongo/db/exec/timeseries_write.h
@@ -0,0 +1,83 @@
+/**
+ *    Copyright (C) 2023-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.
+ */
+
+
+#pragma once
+
+#include "mongo/db/exec/requires_collection_stage.h"
+
+namespace mongo {
+
+class TimeseriesWriteStage final : public RequiresMutableCollectionStage {
+public:
+    static const char* kStageType;
+
+    TimeseriesWriteStage(ExpressionContext* expCtx,
+                         WorkingSet* ws,
+                         std::unique_ptr<PlanStage> child,
+                         const CollectionPtr& coll,
+                         std::unique_ptr<MatchExpression> residualPredicate,
+                         const MatchExpression* originalBucketPredicate);
+
+    StageType stageType() const {
+        return STAGE_TIMESERIES_WRITE;
+    }
+
+    bool isEOF() final;
+
+    std::unique_ptr<PlanStageStats> getStats();
+
+    const SpecificStats* getSpecificStats() const {
+        return &_specificStats;
+    }
+
+    PlanStage::StageState doWork(WorkingSetID* id);
+
+protected:
+    void doSaveStateRequiresCollection() final {}
+
+    void doRestoreStateRequiresCollection() final {}
+
+private:
+    WorkingSet* _ws;
+
+    //
+    // Main execution machinery data structures.
+    //
+
+    // Determines the measurements to delete from this bucket, and by inverse, those to keep
+    // unmodified.
+    std::unique_ptr<MatchExpression> _residualPredicate;
+
+    std::vector<WorkingSetID> _unchangedMeasurements;
+    std::vector<WorkingSetID> _deletedMeasurements;
+
+    TimeseriesWriteStats _specificStats{};
+};
+}  //  namespace mongo
diff --git a/src/mongo/db/query/classic_stage_builder.cpp b/src/mongo/db/query/classic_stage_builder.cpp
index a11db550e04..2410384ebcf 100644
--- a/src/mongo/db/query/classic_stage_builder.cpp
+++ b/src/mongo/db/query/classic_stage_builder.cpp
@@ -421,6 +421,7 @@ std::unique_ptr<PlanStage> ClassicStageBuilder::build(const QuerySolutionNode* r
         case STAGE_TRIAL:
         case STAGE_UNKNOWN:
         case STAGE_UNPACK_TIMESERIES_BUCKET:
+        case STAGE_TIMESERIES_WRITE:
         case STAGE_SENTINEL:
         case STAGE_COLUMN_SCAN:
         case STAGE_UPDATE: {
diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
index 677267fc8b7..cc97c25e7f3 100644
--- a/src/mongo/db/query/plan_executor_impl.cpp
+++ b/src/mongo/db/query/plan_executor_impl.cpp
@@ -611,6 +611,11 @@ long long PlanExecutorImpl::executeDelete() {
             const SpecificStats* stats = _root->child()->getSpecificStats();
             return static_cast<const DeleteStats*>(stats)->docsDeleted;
         }
+        case StageType::STAGE_TIMESERIES_WRITE: {
+            const auto* tsWriteStats =
+                static_cast<const TimeseriesWriteStats*>(_root->getSpecificStats());
+            return tsWriteStats->measurementsDeleted;
+        }
         default: {
             invariant(StageType::STAGE_DELETE == _root->stageType() ||
                       StageType::STAGE_BATCHED_DELETE == _root->stageType());
diff --git a/src/mongo/db/query/stage_types.cpp b/src/mongo/db/query/stage_types.cpp
index 3de3fa2086e..934fe3107ec 100644
--- a/src/mongo/db/query/stage_types.cpp
+++ b/src/mongo/db/query/stage_types.cpp
@@ -74,6 +74,7 @@ StringData stageTypeToString(StageType stageType) {
         {STAGE_SUBPLAN, "SUBPLAN"_sd},
         {STAGE_TEXT_OR, "TEXT_OR"_sd},
         {STAGE_TEXT_MATCH, "TEXT_MATCH"_sd},
+        {STAGE_TIMESERIES_WRITE, "TIMESERIES_WRITE"_sd},
         {STAGE_TRIAL, "TRIAL"_sd},
         {STAGE_UNKNOWN, "UNKNOWN"_sd},
         {STAGE_UNPACK_TIMESERIES_BUCKET, "UNPACK_TIMESERIES_BUCKET"_sd},
diff --git a/src/mongo/db/query/stage_types.h b/src/mongo/db/query/stage_types.h
index f5ee3a70d4c..6398573f543 100644
--- a/src/mongo/db/query/stage_types.h
+++ b/src/mongo/db/query/stage_types.h
@@ -117,6 +117,9 @@ enum StageType {
     STAGE_TEXT_OR,
     STAGE_TEXT_MATCH,
 
+    // Stage for modifying bucket documents in a time-series bucket collection.
+    STAGE_TIMESERIES_WRITE,
+
     // Stage for choosing between two alternate plans based on an initial trial period.
     STAGE_TRIAL,