SERVER-67338 Fix ownership of HashAggStage key in SBE

1 files changed, 10 insertions, 17 deletions

diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
index e05f4d7cf63..1dcbc500ec8 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
@@ -354,25 +354,18 @@ void HashAggStage::open(bool reOpen) {
                 key.reset(idx++, false, tag, val);
             }
 
-            if (!_recordStore) {
-                // The memory limit hasn't been reached yet, accumulate state in '_ht'.
-                auto [it, inserted] = _ht->try_emplace(std::move(key), value::MaterializedRow{0});
-                if (inserted) {
-                    // Copy keys.
-                    const_cast<value::MaterializedRow&>(it->first).makeOwned();
-                    // Initialize accumulators.
-                    it->second.resize(_outAggAccessors.size());
-                }
-                // Always update the state in the '_ht' for the branch when data hasn't been
-                // spilled to disk.
+
+            if (_htIt = _ht->find(key); !_recordStore && _htIt == _ht->end()) {
+                // The memory limit hasn't been reached yet, insert a new key in '_ht' by copying
+                // the key. Note as a future optimization, we should avoid the lookup in the find()
+                // call and the emplace.
+                key.makeOwned();
+                auto [it, _] = _ht->emplace(std::move(key), value::MaterializedRow{0});
+                // Initialize accumulators.
+                it->second.resize(_outAggAccessors.size());
                 _htIt = it;
-                updateAggStateHt = true;
-            } else {
-                // The memory limit has been reached, accumulate state in '_ht' only if we
-                // find the key in '_ht'.
-                _htIt = _ht->find(key);
-                updateAggStateHt = _htIt != _ht->end();
             }
+            updateAggStateHt = _htIt != _ht->end();
 
             if (updateAggStateHt) {
                 // Accumulate state in '_ht' by pointing the '_outAggAccessors' the