SERVER-59075 Create temporary RecordStore in HashAgg stage

28 files changed, 214 insertions, 49 deletions

diff --git a/src/mongo/db/catalog/multi_index_block.cpp b/src/mongo/db/catalog/multi_index_block.cpp
index 46084c3ffb2..eb987f7ed7d 100644
--- a/src/mongo/db/catalog/multi_index_block.cpp
+++ b/src/mongo/db/catalog/multi_index_block.cpp
@@ -956,7 +956,7 @@ void MultiIndexBlock::_writeStateToDisk(OperationContext* opCtx,
     auto obj = _constructStateObject(opCtx, collection);
     auto rs = opCtx->getServiceContext()
                   ->getStorageEngine()
-                  ->makeTemporaryRecordStoreForResumableIndexBuild(opCtx);
+                  ->makeTemporaryRecordStoreForResumableIndexBuild(opCtx, KeyFormat::Long);
 
     WriteUnitOfWork wuow(opCtx);
 
diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
index 74c74e2b3a5..5c2a13476ec 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
@@ -133,6 +133,44 @@ value::SlotAccessor* HashAggStage::getAccessor(CompileCtx& ctx, value::SlotId sl
     return ctx.getAccessor(slot);
 }
 
+namespace {
+// This check makes sure we are safe to spill to disk without the need to abandon the current
+// snapshot.
+void assertIgnoreConflictsWriteBehavior(OperationContext* opCtx) {
+    tassert(5907502,
+            "The operation must be ignoring conflicts and allowing writes or enforcing prepare "
+            "conflicts entirely",
+            opCtx->recoveryUnit()->getPrepareConflictBehavior() !=
+                PrepareConflictBehavior::kIgnoreConflicts);
+}
+}  // namespace
+void HashAggStage::doDetachFromOperationContext() {
+    if (_recordStore) {
+        _recordStore->doDetachFromOperationContext();
+    }
+}
+
+void HashAggStage::doAttachToOperationContext(OperationContext* opCtx) {
+    if (_recordStore) {
+        _recordStore->doAttachToOperationContext(opCtx);
+    }
+}
+
+void HashAggStage::makeTemporaryRecordStore() {
+    tassert(
+        5907500,
+        "HashAggStage attempted to write to disk in an environment which is not prepared to do so",
+        _opCtx->getServiceContext());
+    tassert(5907501,
+            "No storage engine so HashAggStage cannot spill to disk",
+            _opCtx->getServiceContext()->getStorageEngine());
+    assertIgnoreConflictsWriteBehavior(_opCtx);
+    _recordStore = std::make_unique<RecordStore>(
+        _opCtx,
+        _opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
+            _opCtx, KeyFormat::String));
+}
+
 void HashAggStage::open(bool reOpen) {
     auto optTimer(getOptTimer(_opCtx));
 
@@ -191,9 +229,12 @@ void HashAggStage::open(bool reOpen) {
                 long estimatedSizeForOneRow =
                     it->first.memUsageForSorter() + it->second.memUsageForSorter();
                 long long estimatedTotalSize = _ht->size() * estimatedSizeForOneRow;
-                uassert(5859000,
-                        "Need to spill to disk",
-                        estimatedTotalSize < _approxMemoryUseInBytesBeforeSpill);
+
+                if (estimatedTotalSize >= _approxMemoryUseInBytesBeforeSpill) {
+                    // TODO SERVER-58436: Remove this uassert when spilling is implemented.
+                    uasserted(5859000, "Need to spill to disk");
+                    makeTemporaryRecordStore();
+                }
             }
         }
 
@@ -234,6 +275,12 @@ PlanState HashAggStage::getNext() {
     }
 
     if (_htIt == _ht->end()) {
+        if (_recordStore && _seekKeysAccessors.empty()) {
+            // A record store was created to spill to disk. Clean it up.
+            assertIgnoreConflictsWriteBehavior(_opCtx);
+            _recordStore->deleteTemporaryRecordStoreIfLock();
+            _recordStore.reset();
+        }
         return trackPlanState(PlanState::IS_EOF);
     }
 
@@ -269,6 +316,12 @@ void HashAggStage::close() {
 
     trackClose();
     _ht = boost::none;
+    if (_recordStore) {
+        // A record store was created to spill to disk. Clean it up.
+        assertIgnoreConflictsWriteBehavior(_opCtx);
+        _recordStore->deleteTemporaryRecordStoreIfLock();
+        _recordStore.reset();
+    }
 
     if (_childOpened) {
         _children[0]->close();
diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.h b/src/mongo/db/exec/sbe/stages/hash_agg.h
index 67302aeb02a..9363c3db142 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.h
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.h
@@ -32,6 +32,7 @@
 #include <unordered_map>
 
 #include "mongo/db/exec/sbe/expressions/expression.h"
+#include "mongo/db/exec/sbe/stages/record_store.h"
 #include "mongo/db/exec/sbe/stages/stages.h"
 #include "mongo/db/exec/sbe/vm/vm.h"
 #include "mongo/db/query/query_knobs_gen.h"
@@ -88,7 +89,13 @@ public:
     std::vector<DebugPrinter::Block> debugPrint() const final;
     size_t estimateCompileTimeSize() const final;
 
+protected:
+    void doDetachFromOperationContext() override;
+    void doAttachToOperationContext(OperationContext* opCtx) override;
+
 private:
+    void makeTemporaryRecordStore();
+
     using TableType = stdx::unordered_map<value::MaterializedRow,
                                           value::MaterializedRow,
                                           value::MaterializedRowHasher,
@@ -132,6 +139,10 @@ private:
 
     bool _compiled{false};
     bool _childOpened{false};
+
+    // Used when spilling to disk.
+    std::unique_ptr<RecordStore> _recordStore;
 };
+
 }  // namespace sbe
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/stages/record_store.h b/src/mongo/db/exec/sbe/stages/record_store.h
new file mode 100644
index 00000000000..da33ba80d0c
--- /dev/null
+++ b/src/mongo/db/exec/sbe/stages/record_store.h
@@ -0,0 +1,86 @@
+/**
+ *    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.
+ */
+
+#pragma once
+
+#include "mongo/db/operation_context.h"
+#include "mongo/db/storage/temporary_record_store.h"
+
+namespace mongo {
+namespace sbe {
+/**
+ * A way of interacting with the TemporaryRecordStore in PlanStages.
+ * TODO SERVER-60639: Remove this when SERVER-60639 is done.
+ */
+class RecordStore {
+public:
+    RecordStore(OperationContext* opCtx, std::unique_ptr<TemporaryRecordStore> rs)
+        : _opCtx(opCtx), _rs(std::move(rs)){};
+
+    RecordStore(RecordStore&& other) : _opCtx(other._opCtx), _rs(std::move(other._rs)) {
+        other._opCtx = nullptr;
+    }
+    ~RecordStore() {
+        deleteTemporaryRecordStoreIfLock();
+    }
+
+    RecordStore& operator=(RecordStore&& other) {
+        if (this == &other) {
+            return *this;
+        }
+        _opCtx = other._opCtx;
+        _rs = std::move(other._rs);
+        other._opCtx = nullptr;
+        return *this;
+    }
+
+    void doDetachFromOperationContext() {
+        _opCtx = nullptr;
+    }
+
+    void doAttachToOperationContext(OperationContext* opCtx) {
+        _opCtx = opCtx;
+    }
+
+    void deleteTemporaryRecordStoreIfLock() {
+        // A best-effort to destroy the TRS, we can only do this if the query is holding a
+        // global lock.
+        if (_opCtx->lockState()->wasGlobalLockTaken()) {
+            _rs->finalizeTemporaryTable(_opCtx, TemporaryRecordStore::FinalizationAction::kDelete);
+        }
+    }
+    // add attachToOperationContext, detachFromOperationContext, call this from
+    // HashAgg::attachToOperationContext, detachFromOperationContext.
+
+private:
+    OperationContext* _opCtx;
+    std::unique_ptr<TemporaryRecordStore> _rs;
+};
+}  // namespace sbe
+}  // namespace mongo
diff --git a/src/mongo/db/index/duplicate_key_tracker.cpp b/src/mongo/db/index/duplicate_key_tracker.cpp
index da9f931c079..8c02a28fe99 100644
--- a/src/mongo/db/index/duplicate_key_tracker.cpp
+++ b/src/mongo/db/index/duplicate_key_tracker.cpp
@@ -49,8 +49,8 @@ static constexpr StringData kKeyField = "key"_sd;
 
 DuplicateKeyTracker::DuplicateKeyTracker(OperationContext* opCtx, const IndexCatalogEntry* entry)
     : _indexCatalogEntry(entry),
-      _keyConstraintsTable(
-          opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(opCtx)) {
+      _keyConstraintsTable(opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
+          opCtx, KeyFormat::Long)) {
 
     invariant(_indexCatalogEntry->descriptor()->unique());
 }
diff --git a/src/mongo/db/index/index_build_interceptor.cpp b/src/mongo/db/index/index_build_interceptor.cpp
index 88a0690a8c6..eb6f8893f09 100644
--- a/src/mongo/db/index/index_build_interceptor.cpp
+++ b/src/mongo/db/index/index_build_interceptor.cpp
@@ -59,8 +59,8 @@ MONGO_FAIL_POINT_DEFINE(hangIndexBuildDuringDrainWritesPhaseSecond);
 
 IndexBuildInterceptor::IndexBuildInterceptor(OperationContext* opCtx, IndexCatalogEntry* entry)
     : _indexCatalogEntry(entry),
-      _sideWritesTable(
-          opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(opCtx)),
+      _sideWritesTable(opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
+          opCtx, KeyFormat::Long)),
       _skippedRecordTracker(opCtx, entry, boost::none) {
 
     if (entry->descriptor()->unique()) {
diff --git a/src/mongo/db/index/skipped_record_tracker.cpp b/src/mongo/db/index/skipped_record_tracker.cpp
index 7891cda5478..e40795fe18a 100644
--- a/src/mongo/db/index/skipped_record_tracker.cpp
+++ b/src/mongo/db/index/skipped_record_tracker.cpp
@@ -75,7 +75,8 @@ void SkippedRecordTracker::record(OperationContext* opCtx, const RecordId& recor
     // Lazily initialize table when we record the first document.
     if (!_skippedRecordsTable) {
         _skippedRecordsTable =
-            opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(opCtx);
+            opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
+                opCtx, KeyFormat::Long);
     }
 
     writeConflictRetry(
diff --git a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
index 9a3ccf3b9c5..36256027436 100644
--- a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
+++ b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
@@ -152,7 +152,6 @@ void assertIgnorePrepareConflictsBehavior(const boost::intrusive_ptr<ExpressionC
             expCtx->opCtx->recoveryUnit()->getPrepareConflictBehavior() !=
                 PrepareConflictBehavior::kIgnoreConflicts);
 }
-
 }  // namespace
 
 std::unique_ptr<TransactionHistoryIteratorBase>
@@ -671,11 +670,12 @@ void CommonMongodProcessInterface::writeRecordsToRecordStore(
         wuow.commit();
     });
 }
+
 std::unique_ptr<TemporaryRecordStore> CommonMongodProcessInterface::createTemporaryRecordStore(
-    const boost::intrusive_ptr<ExpressionContext>& expCtx) const {
+    const boost::intrusive_ptr<ExpressionContext>& expCtx, KeyFormat keyFormat) const {
     assertIgnorePrepareConflictsBehavior(expCtx);
     return expCtx->opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
-        expCtx->opCtx);
+        expCtx->opCtx, keyFormat);
 }
 
 Document CommonMongodProcessInterface::readRecordFromRecordStore(
diff --git a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h
index 16421e54a31..c13e6fb45a7 100644
--- a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h
@@ -102,7 +102,7 @@ public:
                                            const NamespaceString& outputNs) const final;
 
     std::unique_ptr<TemporaryRecordStore> createTemporaryRecordStore(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx) const final;
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, KeyFormat keyFormat) const final;
 
     void writeRecordsToRecordStore(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                                    RecordStore* rs,
diff --git a/src/mongo/db/pipeline/process_interface/mongo_process_interface.h b/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
index 9cf6043dfc6..e578947c154 100644
--- a/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
@@ -483,7 +483,7 @@ public:
      * Create a temporary record store.
      */
     virtual std::unique_ptr<TemporaryRecordStore> createTemporaryRecordStore(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx) const = 0;
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, KeyFormat keyFormat) const = 0;
 
     /**
      * Write the records in 'records' to the record store. Record store must already exist. Asserts
diff --git a/src/mongo/db/pipeline/process_interface/mongos_process_interface.h b/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
index f1dd0eab9b8..17603103419 100644
--- a/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
@@ -252,7 +252,7 @@ public:
         boost::optional<BSONObj> readConcern = boost::none) final;
 
     std::unique_ptr<TemporaryRecordStore> createTemporaryRecordStore(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx) const final {
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, KeyFormat keyFormat) const final {
         MONGO_UNREACHABLE;
     }
     void writeRecordsToRecordStore(const boost::intrusive_ptr<ExpressionContext>& expCtx,
diff --git a/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h b/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
index 51adcf96f6e..580f032ba96 100644
--- a/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
@@ -282,7 +282,7 @@ public:
     }
 
     std::unique_ptr<TemporaryRecordStore> createTemporaryRecordStore(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx) const {
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, KeyFormat keyFormat) const {
         MONGO_UNREACHABLE;
     }
 
diff --git a/src/mongo/db/pipeline/window_function/spillable_cache.cpp b/src/mongo/db/pipeline/window_function/spillable_cache.cpp
index d3ebb07ae0b..6c4d0804e8c 100644
--- a/src/mongo/db/pipeline/window_function/spillable_cache.cpp
+++ b/src/mongo/db/pipeline/window_function/spillable_cache.cpp
@@ -120,7 +120,8 @@ void SpillableCache::spillToDisk() {
                 "engine configured",
                 _expCtx->opCtx->getServiceContext()->getStorageEngine());
         _usedDisk = true;
-        _diskCache = _expCtx->mongoProcessInterface->createTemporaryRecordStore(_expCtx);
+        _diskCache =
+            _expCtx->mongoProcessInterface->createTemporaryRecordStore(_expCtx, KeyFormat::Long);
     }
     // If we've freed things from cache before writing to disk, we need to update
     // '_diskWrittenIndex' to be the actual index of the document we're going to write.
diff --git a/src/mongo/db/pipeline/window_function/spillable_cache_test.cpp b/src/mongo/db/pipeline/window_function/spillable_cache_test.cpp
index bacd4eb59b7..83cb374ebc0 100644
--- a/src/mongo/db/pipeline/window_function/spillable_cache_test.cpp
+++ b/src/mongo/db/pipeline/window_function/spillable_cache_test.cpp
@@ -43,12 +43,12 @@ namespace {
 class MongoProcessInterfaceForTest : public StubMongoProcessInterface {
 public:
     std::unique_ptr<TemporaryRecordStore> createTemporaryRecordStore(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx) const override {
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, KeyFormat keyFormat) const override {
         expCtx->opCtx->recoveryUnit()->abandonSnapshot();
         expCtx->opCtx->recoveryUnit()->setPrepareConflictBehavior(
             PrepareConflictBehavior::kIgnoreConflictsAllowWrites);
         return expCtx->opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
-            expCtx->opCtx);
+            expCtx->opCtx, keyFormat);
     }
 
     void writeRecordsToRecordStore(const boost::intrusive_ptr<ExpressionContext>& expCtx,
diff --git a/src/mongo/db/storage/devnull/devnull_kv_engine.cpp b/src/mongo/db/storage/devnull/devnull_kv_engine.cpp
index 144b289e27e..4b8ef17e47d 100644
--- a/src/mongo/db/storage/devnull/devnull_kv_engine.cpp
+++ b/src/mongo/db/storage/devnull/devnull_kv_engine.cpp
@@ -60,8 +60,11 @@ public:
 
 class DevNullRecordStore : public RecordStore {
 public:
-    DevNullRecordStore(StringData ns, StringData identName, const CollectionOptions& options)
-        : RecordStore(ns, identName), _options(options) {
+    DevNullRecordStore(StringData ns,
+                       StringData identName,
+                       const CollectionOptions& options,
+                       KeyFormat keyFormat)
+        : RecordStore(ns, identName), _options(options), _keyFormat(keyFormat) {
         _numInserts = 0;
         _dummy = BSON("_id" << 1);
     }
@@ -85,7 +88,7 @@ public:
     }
 
     virtual KeyFormat keyFormat() const {
-        return KeyFormat::Long;
+        return _keyFormat;
     }
 
     virtual int64_t storageSize(OperationContext* opCtx,
@@ -155,6 +158,7 @@ public:
 
 private:
     CollectionOptions _options;
+    KeyFormat _keyFormat;
     long long _numInserts;
     BSONObj _dummy;
 };
@@ -240,12 +244,13 @@ std::unique_ptr<RecordStore> DevNullKVEngine::getRecordStore(OperationContext* o
     if (ident == "_mdb_catalog") {
         return std::make_unique<EphemeralForTestRecordStore>(ns, ident, &_catalogInfo);
     }
-    return std::make_unique<DevNullRecordStore>(ns, ident, options);
+    return std::make_unique<DevNullRecordStore>(ns, ident, options, KeyFormat::Long);
 }
 
 std::unique_ptr<RecordStore> DevNullKVEngine::makeTemporaryRecordStore(OperationContext* opCtx,
-                                                                       StringData ident) {
-    return std::make_unique<DevNullRecordStore>("" /* ns */, ident, CollectionOptions());
+                                                                       StringData ident,
+                                                                       KeyFormat keyFormat) {
+    return std::make_unique<DevNullRecordStore>("" /* ns */, ident, CollectionOptions(), keyFormat);
 }
 
 std::unique_ptr<SortedDataInterface> DevNullKVEngine::getSortedDataInterface(
diff --git a/src/mongo/db/storage/devnull/devnull_kv_engine.h b/src/mongo/db/storage/devnull/devnull_kv_engine.h
index e3f3f2ae3f3..7327d8541fc 100644
--- a/src/mongo/db/storage/devnull/devnull_kv_engine.h
+++ b/src/mongo/db/storage/devnull/devnull_kv_engine.h
@@ -62,7 +62,8 @@ public:
                                                         const CollectionOptions& options);
 
     virtual std::unique_ptr<RecordStore> makeTemporaryRecordStore(OperationContext* opCtx,
-                                                                  StringData ident) override;
+                                                                  StringData ident,
+                                                                  KeyFormat keyFormat) override;
 
     virtual Status createSortedDataInterface(OperationContext* opCtx,
                                              const CollectionOptions& collOptions,
diff --git a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.cpp b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.cpp
index f51e580fdc6..49b63faac74 100644
--- a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.cpp
+++ b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.cpp
@@ -84,9 +84,10 @@ Status KVEngine::importRecordStore(OperationContext* opCtx,
 }
 
 std::unique_ptr<mongo::RecordStore> KVEngine::makeTemporaryRecordStore(OperationContext* opCtx,
-                                                                       StringData ident) {
+                                                                       StringData ident,
+                                                                       KeyFormat keyFormat) {
     std::unique_ptr<mongo::RecordStore> recordStore =
-        std::make_unique<RecordStore>("", ident, KeyFormat::Long, false);
+        std::make_unique<RecordStore>("", ident, keyFormat, false);
     stdx::lock_guard lock(_identsLock);
     _idents[ident.toString()] = true;
     return recordStore;
diff --git a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.h b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.h
index 11b88951580..5d8db6a1c56 100644
--- a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.h
+++ b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_kv_engine.h
@@ -68,8 +68,8 @@ public:
                                                                StringData ident,
                                                                const CollectionOptions& options);
 
-    virtual std::unique_ptr<mongo::RecordStore> makeTemporaryRecordStore(OperationContext* opCtx,
-                                                                         StringData ident) override;
+    virtual std::unique_ptr<mongo::RecordStore> makeTemporaryRecordStore(
+        OperationContext* opCtx, StringData ident, KeyFormat keyFormat) override;
 
     virtual Status createSortedDataInterface(OperationContext* opCtx,
                                              const CollectionOptions& collOptions,
diff --git a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp
index 3bbcb413c6f..ab5a81442e1 100644
--- a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp
+++ b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp
@@ -76,7 +76,8 @@ public:
         return Status::OK();
     }
     std::unique_ptr<RecordStore> makeTemporaryRecordStore(OperationContext* opCtx,
-                                                          StringData ident) override {
+                                                          StringData ident,
+                                                          KeyFormat keyFormat) override {
         return {};
     }
     Status createSortedDataInterface(OperationContext* opCtx,
diff --git a/src/mongo/db/storage/kv/kv_engine.h b/src/mongo/db/storage/kv/kv_engine.h
index 1398e2b6d1c..c8b641d5e51 100644
--- a/src/mongo/db/storage/kv/kv_engine.h
+++ b/src/mongo/db/storage/kv/kv_engine.h
@@ -100,7 +100,8 @@ public:
                                      const CollectionOptions& options) = 0;
 
     virtual std::unique_ptr<RecordStore> makeTemporaryRecordStore(OperationContext* opCtx,
-                                                                  StringData ident) = 0;
+                                                                  StringData ident,
+                                                                  KeyFormat keyFormat) = 0;
 
     /**
      * Similar to createRecordStore but this imports from an existing table with the provided ident
diff --git a/src/mongo/db/storage/kv/kv_engine_test_harness.cpp b/src/mongo/db/storage/kv/kv_engine_test_harness.cpp
index 8d721f8cc26..e37fcd9281e 100644
--- a/src/mongo/db/storage/kv/kv_engine_test_harness.cpp
+++ b/src/mongo/db/storage/kv/kv_engine_test_harness.cpp
@@ -303,7 +303,7 @@ TEST_F(KVEngineTestHarness, TemporaryRecordStoreSimple) {
     std::unique_ptr<RecordStore> rs;
     {
         auto opCtx = _makeOperationContext(engine);
-        rs = engine->makeTemporaryRecordStore(opCtx.get(), ident);
+        rs = engine->makeTemporaryRecordStore(opCtx.get(), ident, KeyFormat::Long);
         ASSERT(rs);
     }
 
diff --git a/src/mongo/db/storage/storage_engine.h b/src/mongo/db/storage/storage_engine.h
index 11b4ab609b8..80bc03ebff7 100644
--- a/src/mongo/db/storage/storage_engine.h
+++ b/src/mongo/db/storage/storage_engine.h
@@ -377,8 +377,8 @@ public:
      * Creates a temporary RecordStore on the storage engine. On startup after an unclean shutdown,
      * the storage engine will drop any un-dropped temporary record stores.
      */
-    virtual std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStore(
-        OperationContext* opCtx) = 0;
+    virtual std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStore(OperationContext* opCtx,
+                                                                           KeyFormat keyFormat) = 0;
 
     /**
      * Creates a temporary RecordStore on the storage engine for a resumable index build. On
@@ -386,7 +386,7 @@ public:
      * record stores.
      */
     virtual std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStoreForResumableIndexBuild(
-        OperationContext* opCtx) = 0;
+        OperationContext* opCtx, KeyFormat keyFormat) = 0;
 
     /**
      * Creates a temporary RecordStore on the storage engine from an existing ident on disk. On
diff --git a/src/mongo/db/storage/storage_engine_impl.cpp b/src/mongo/db/storage/storage_engine_impl.cpp
index 29785aabd35..fbd4f0f4214 100644
--- a/src/mongo/db/storage/storage_engine_impl.cpp
+++ b/src/mongo/db/storage/storage_engine_impl.cpp
@@ -930,17 +930,18 @@ Status StorageEngineImpl::repairRecordStore(OperationContext* opCtx,
 }
 
 std::unique_ptr<TemporaryRecordStore> StorageEngineImpl::makeTemporaryRecordStore(
-    OperationContext* opCtx) {
+    OperationContext* opCtx, KeyFormat keyFormat) {
     std::unique_ptr<RecordStore> rs =
-        _engine->makeTemporaryRecordStore(opCtx, _catalog->newInternalIdent());
+        _engine->makeTemporaryRecordStore(opCtx, _catalog->newInternalIdent(), keyFormat);
     LOGV2_DEBUG(22258, 1, "Created temporary record store", "ident"_attr = rs->getIdent());
     return std::make_unique<TemporaryKVRecordStore>(getEngine(), std::move(rs));
 }
 
 std::unique_ptr<TemporaryRecordStore>
-StorageEngineImpl::makeTemporaryRecordStoreForResumableIndexBuild(OperationContext* opCtx) {
-    std::unique_ptr<RecordStore> rs =
-        _engine->makeTemporaryRecordStore(opCtx, _catalog->newInternalResumableIndexBuildIdent());
+StorageEngineImpl::makeTemporaryRecordStoreForResumableIndexBuild(OperationContext* opCtx,
+                                                                  KeyFormat keyFormat) {
+    std::unique_ptr<RecordStore> rs = _engine->makeTemporaryRecordStore(
+        opCtx, _catalog->newInternalResumableIndexBuildIdent(), keyFormat);
     LOGV2_DEBUG(4921500,
                 1,
                 "Created temporary record store for resumable index build",
diff --git a/src/mongo/db/storage/storage_engine_impl.h b/src/mongo/db/storage/storage_engine_impl.h
index c8d3e5cfee1..7ccf2261b1a 100644
--- a/src/mongo/db/storage/storage_engine_impl.h
+++ b/src/mongo/db/storage/storage_engine_impl.h
@@ -109,10 +109,10 @@ public:
                                      const NamespaceString& nss) override;
 
     virtual std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStore(
-        OperationContext* opCtx) override;
+        OperationContext* opCtx, KeyFormat keyFormat) override;
 
     virtual std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStoreForResumableIndexBuild(
-        OperationContext* opCtx) override;
+        OperationContext* opCtx, KeyFormat keyFormat) override;
 
     virtual std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStoreFromExistingIdent(
         OperationContext* opCtx, StringData ident) override;
diff --git a/src/mongo/db/storage/storage_engine_mock.h b/src/mongo/db/storage/storage_engine_mock.h
index 01b7f0409c0..3ad87268784 100644
--- a/src/mongo/db/storage/storage_engine_mock.h
+++ b/src/mongo/db/storage/storage_engine_mock.h
@@ -87,11 +87,12 @@ public:
                              const NamespaceString& ns) final {
         return Status::OK();
     }
-    std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStore(OperationContext* opCtx) final {
+    std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStore(OperationContext* opCtx,
+                                                                   KeyFormat keyFormat) final {
         return {};
     }
     std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStoreForResumableIndexBuild(
-        OperationContext* opCtx) final {
+        OperationContext* opCtx, KeyFormat keyFormat) final {
         return {};
     }
     std::unique_ptr<TemporaryRecordStore> makeTemporaryRecordStoreFromExistingIdent(
diff --git a/src/mongo/db/storage/storage_engine_test_fixture.h b/src/mongo/db/storage/storage_engine_test_fixture.h
index 3677e4026db..81ee4eb7a75 100644
--- a/src/mongo/db/storage/storage_engine_test_fixture.h
+++ b/src/mongo/db/storage/storage_engine_test_fixture.h
@@ -91,7 +91,7 @@ public:
     }
 
     std::unique_ptr<TemporaryRecordStore> makeTemporary(OperationContext* opCtx) {
-        return _storageEngine->makeTemporaryRecordStore(opCtx);
+        return _storageEngine->makeTemporaryRecordStore(opCtx, KeyFormat::Long);
     }
 
     /**
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.cpp
index 48e361a39d7..2c278ef7f18 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.cpp
@@ -1599,7 +1599,8 @@ std::unique_ptr<SortedDataInterface> WiredTigerKVEngine::getSortedDataInterface(
 }
 
 std::unique_ptr<RecordStore> WiredTigerKVEngine::makeTemporaryRecordStore(OperationContext* opCtx,
-                                                                          StringData ident) {
+                                                                          StringData ident,
+                                                                          KeyFormat keyFormat) {
     invariant(!_readOnly || !recoverToOplogTimestamp.empty());
 
     _ensureIdentPath(ident);
@@ -1626,7 +1627,7 @@ std::unique_ptr<RecordStore> WiredTigerKVEngine::makeTemporaryRecordStore(Operat
     params.ident = ident.toString();
     params.engineName = _canonicalName;
     params.isCapped = false;
-    params.keyFormat = KeyFormat::Long;
+    params.keyFormat = keyFormat;
     params.overwrite = true;
     params.isEphemeral = _ephemeral;
     params.cappedCallback = nullptr;
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.h b/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.h
index fe75e451a36..8d951088935 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.h
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_kv_engine.h
@@ -137,7 +137,8 @@ public:
                                                 const CollectionOptions& options) override;
 
     std::unique_ptr<RecordStore> makeTemporaryRecordStore(OperationContext* opCtx,
-                                                          StringData ident) override;
+                                                          StringData ident,
+                                                          KeyFormat keyFormat) override;
 
     Status createSortedDataInterface(OperationContext* opCtx,
                                      const CollectionOptions& collOptions,