1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
|
/**
* Copyright (C) 2019-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 <unordered_map>
#include "mongo/db/exec/sbe/expressions/expression.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"
#include "mongo/db/storage/temporary_record_store.h"
#include "mongo/stdx/unordered_map.h"
namespace mongo {
namespace sbe {
/**
* Performs a hash-based aggregation. Appears as the "group" stage in debug output. Groups the input
* based on the provided vector of group-by slots, 'gbs'. The 'aggs' parameter is a map from
* 'SlotId' to expression. This defines a set of output slots whose values will be computed based on
* the corresponding aggregate expressions. Each distinct grouping will produce a single output,
* consisting of the values of the group-by keys and the results of the aggregate functions.
*
* Since the data must be buffered in a hash table, this is a "binding reflector". This means slots
* from the 'input' tree are not visible higher in tree. Stages higher in the tree can only see the
* slots holding the group-by keys as well as those holding the corresponding aggregate values.
*
* The optional 'seekKeys', if provided, limit the results returned from the hash table only to
* those equal to seekKeys.
*
* The 'optimizedClose' flag controls whether we can close the child subtree right after building
* the hash table. If true it means that we do not expect the subtree to be reopened.
*
* The optional 'collatorSlot', if provided, changes the definition of string equality used when
* determining whether two group-by keys are equal. For instance, the plan may require us to do a
* case-insensitive group on a string field.
*
* Debug string representation:
*
* group [<group by slots>] [slot_1 = expr_1, ..., slot_n = expr_n] [<seek slots>]? reopen?
* collatorSlot? childStage
*/
class HashAggStage final : public PlanStage {
public:
HashAggStage(std::unique_ptr<PlanStage> input,
value::SlotVector gbs,
value::SlotMap<std::unique_ptr<EExpression>> aggs,
value::SlotVector seekKeysSlots,
bool optimizedClose,
boost::optional<value::SlotId> collatorSlot,
bool allowDiskUse,
PlanNodeId planNodeId,
bool participateInTrialRunTracking = true);
std::unique_ptr<PlanStage> clone() const final;
void prepare(CompileCtx& ctx) final;
value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final;
void open(bool reOpen) final;
PlanState getNext() final;
void close() final;
std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final;
const SpecificStats* getSpecificStats() const final;
std::vector<DebugPrinter::Block> debugPrint() const final;
size_t estimateCompileTimeSize() const final;
protected:
void doSaveState(bool relinquishCursor) override;
void doRestoreState(bool relinquishCursor) override;
void doDetachFromOperationContext() override;
void doAttachToOperationContext(OperationContext* opCtx) override;
void doDetachFromTrialRunTracker() override;
TrialRunTrackerAttachResultMask doAttachToTrialRunTracker(
TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) override;
private:
using TableType = stdx::unordered_map<value::MaterializedRow,
value::MaterializedRow,
value::MaterializedRowHasher,
value::MaterializedRowEq>;
using HashKeyAccessor = value::MaterializedRowKeyAccessor<TableType::iterator>;
using HashAggAccessor = value::MaterializedRowValueAccessor<TableType::iterator>;
void makeTemporaryRecordStore();
/**
* Spills a key and value pair to the '_recordStore' where the semantics are insert or update
* depending on the 'update' flag. When the 'update' flag is true this method already expects
* the 'key' to be inserted into the '_recordStore', otherwise the 'key' and 'val' pair are
* fresh.
*
* This method expects the key to be seralized into a KeyString::Value so that the key is
* memcmp-able and lookups can be done to update the 'val' in the '_recordStore'. Note that the
* 'typeBits' are needed to reconstruct the spilled 'key' when calling 'getNext' to deserialize
* the 'key' to a MaterializedRow. Since the '_recordStore' only stores the memcmp-able part of
* the KeyString we need to carry the 'typeBits' separately, and we do this by appending the
* 'typeBits' to the end of the serialized 'val' buffer and store them at the leaves of the
* backing B-tree of the '_recordStore'. used as the RecordId.
*/
void spillValueToDisk(const RecordId& key,
const value::MaterializedRow& val,
const KeyString::TypeBits& typeBits,
bool update);
void spillRowToDisk(const value::MaterializedRow& key,
const value::MaterializedRow& defaultVal);
/**
* We check amount of used memory every T processed incoming records, where T is calculated
* based on the estimated used memory and its recent growth. When the memory limit is exceeded,
* 'checkMemoryUsageAndSpillIfNecessary()' will create '_recordStore' and might spill some of
* the already accumulated data into it.
*/
struct MemoryCheckData {
const double checkpointMargin = internalQuerySBEAggMemoryUseCheckMargin.load();
const long atMostCheckFrequency = internalQuerySBEAggMemoryCheckPerAdvanceAtMost.load();
const long atLeastMemoryCheckFrequency =
internalQuerySBEAggMemoryCheckPerAdvanceAtLeast.load();
// The check frequency upper bound, which start at 'atMost' and exponentially backs off
// to 'atLeast' as more data is accumulated. If 'atLeast' is less than 'atMost', the memory
// checks will be done every 'atLeast' incoming records.
long memoryCheckFrequency = 1;
// The number of incoming records to process before the next memory checkpoint.
long nextMemoryCheckpoint = 0;
// The counter of the incoming records between memory checkpoints.
long memoryCheckpointCounter = 0;
long long lastEstimatedMemoryUsage = 0;
MemoryCheckData() {
memoryCheckFrequency = std::min(atMostCheckFrequency, atLeastMemoryCheckFrequency);
}
};
void checkMemoryUsageAndSpillIfNecessary(MemoryCheckData& mcd);
const value::SlotVector _gbs;
const value::SlotMap<std::unique_ptr<EExpression>> _aggs;
const boost::optional<value::SlotId> _collatorSlot;
const bool _allowDiskUse;
const value::SlotVector _seekKeysSlots;
// When this operator does not expect to be reopened (almost always) then it can close the child
// early.
const bool _optimizedClose{true};
value::SlotAccessorMap _outAccessors;
std::vector<value::SlotAccessor*> _inKeyAccessors;
// Accessors for the key stored in '_ht', a SwitchAccessor is used so we can produce the key
// from either the '_ht' or the '_recordStore'.
std::vector<std::unique_ptr<HashKeyAccessor>> _outHashKeyAccessors;
std::vector<std::unique_ptr<value::SwitchAccessor>> _outKeyAccessors;
// Accessor for the agg state value stored in the '_recordStore' when data is spilled to disk.
value::MaterializedRow _aggKeyRecordStore{0};
value::MaterializedRow _aggValueRecordStore{0};
std::vector<std::unique_ptr<value::MaterializedSingleRowAccessor>> _outRecordStoreKeyAccessors;
std::vector<std::unique_ptr<value::MaterializedSingleRowAccessor>> _outRecordStoreAggAccessors;
// This buffer stores values for the spilled '_aggKeyRecordStore' that's loaded into memory from
// the '_recordStore'. Values in the '_aggKeyRecordStore' row are pointers that point to data in
// this buffer.
BufBuilder _aggKeyRSBuffer;
std::vector<value::SlotAccessor*> _seekKeysAccessors;
value::MaterializedRow _seekKeys;
// Accesors for the agg state in '_ht', a SwitchAccessor is used so we can produce the agg state
// from either the '_ht' or the '_recordStore' when draining the HashAgg stage.
std::vector<std::unique_ptr<value::SwitchAccessor>> _outAggAccessors;
std::vector<std::unique_ptr<HashAggAccessor>> _outHashAggAccessors;
std::vector<std::unique_ptr<vm::CodeFragment>> _aggCodes;
// Only set if collator slot provided on construction.
value::SlotAccessor* _collatorAccessor = nullptr;
boost::optional<TableType> _ht;
TableType::iterator _htIt;
vm::ByteCode _bytecode;
bool _compiled{false};
bool _childOpened{false};
// Memory tracking and spilling to disk.
const long long _approxMemoryUseInBytesBeforeSpill =
internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.load();
std::unique_ptr<TemporaryRecordStore> _recordStore;
bool _drainingRecordStore{false};
std::unique_ptr<SeekableRecordCursor> _rsCursor;
HashAggStats _specificStats;
// If provided, used during a trial run to accumulate certain execution stats. Once the trial
// run is complete, this pointer is reset to nullptr.
TrialRunTracker* _tracker{nullptr};
};
} // namespace sbe
} // namespace mongo
|
